Add support of external logging for postgresql nodes

Author: ildus

testgres/testgres.py

@@ -31,6 +31,11 @@
 import time
 import six
 
+import threading
+import logging
+import select
+import tempfile
+
 # Try to use psycopg2 by default. If psycopg2 isn"t available then use
 # pg8000 which is slower but much more portable because uses only
 # pure-Python code
@@ -44,6 +49,7 @@
 
 
 registered_nodes = []
+util_threads = []
 last_assigned_port = int(random.random() * 16384) + 49152
 pg_config_data = {}
 
@@ -74,6 +80,54 @@ def __str__(self):
         return '\n ERROR: {0}\n CMD: {1}'.format(self.error_text, self.cmd)
 
 
+class AsyncFileReader(threading.Thread):
+    '''
+    Helper class to implement asynchronous reading of a file
+    in a separate thread.
+    '''
+
+    def __init__(self, node_name, fd):
+        assert callable(fd.readline)
+
+        threading.Thread.__init__(self)
+
+        self.node_name = node_name
+        self.fd = fd
+        self.stop_event = threading.Event()
+        self.logger = logging.getLogger(node_name)
+        self.logger.setLevel(logging.INFO)
+
+    def run(self):
+        while self.fd in select.select([self.fd], [], [], 0)[0]:
+            line = self.fd.readline()
+            if line:
+                extra = {'node': self.node_name}
+                self.logger.info(line.strip(), extra=extra)
+            elif self.stopped():
+                break
+            else:
+                time.sleep(0.1)
+
+    def stop(self):
+        self.stop_event.set()
+
+    def stopped(self):
+        return self.stop_event.isSet()
+
+
+def log_watch(node_name, pg_logname):
+    ''' Starts thread for node that redirects postgresql logs
+        to python logging system
+    '''
+
+    reader = AsyncFileReader(node_name, open(pg_logname, 'r'))
+    reader.start()
+
+    global util_threads
+    util_threads.append(reader)
+    return reader
+
+
 class NodeConnection(object):
 
     """
@@ -148,7 +202,7 @@ def close(self):
 
 class PostgresNode(object):
 
-    def __init__(self, name, port, base_dir=None):
+    def __init__(self, name, port, base_dir=None, use_logging=False):
         self.name = name
         self.host = '127.0.0.1'
         self.port = port
@@ -160,6 +214,9 @@ def __init__(self, name, port, base_dir=None):
             os.makedirs(self.logs_dir)
         self.working = False
 
+        self.use_logging = use_logging
+        self.logger = None
+
     @property
     def data_dir(self):
         return os.path.join(self.base_dir, "data")
@@ -316,7 +373,15 @@ def pg_ctl(self, command, params={}, command_options=[]):
 
     def start(self, params={}):
         """ Starts cluster """
-        logfile = os.path.join(self.logs_dir, "postgresql.log")
+
+        if self.use_logging:
+            tmpfile = tempfile.NamedTemporaryFile('w', dir=self.logs_dir, delete=False)
+            logfile = tmpfile.name
+
+            self.logger = log_watch(self.name, logfile)
+        else:
+            logfile = os.path.join(self.logs_dir, "postgresql.log")
+
         _params = {
             "-D": self.data_dir,
             "-w": None,
@@ -623,7 +688,7 @@ def version_to_num(version):
     return num
 
 
-def get_new_node(name, base_dir=None):
+def get_new_node(name, base_dir=None, use_logging=False):
     global registered_nodes
     global last_assigned_port
 
@@ -647,7 +712,7 @@ def get_new_node(name, base_dir=None):
     #              socket.SOCK_STREAM,
     #              socket.getprotobyname("tcp"))
 
-    node = PostgresNode(name, port, base_dir)
+    node = PostgresNode(name, port, base_dir, use_logging=use_logging)
     registered_nodes.append(node)
     last_assigned_port = port
 
@@ -663,7 +728,12 @@ def clean_all():
 
 def stop_all():
     global registered_nodes
+    global util_threads
+
     for node in registered_nodes:
         # stop server if it still working
         if node.working:
             node.stop()
+
+    for thread in util_threads:
+        thread.stop()