Add functions that return PIDs of various processes

Author: ildus

testgres/connection.py

@@ -34,6 +34,7 @@ def __init__(self, node, dbname=None, username=None, password=None):
         username = username or default_username()
 
         self._node = node
+        self._backend_pid = None
 
         self._connection = pglib.connect(
             database=dbname,
@@ -52,6 +53,12 @@ def node(self):
     def connection(self):
         return self._connection
 
+    @property
+    def backend_pid(self):
+        if self._backend_pid is None:
+            self._backend_pid = self.execute("select pg_backend_pid();")[0][0]
+        return self._backend_pid
+
     @property
     def cursor(self):
         return self._cursor

testgres/node.py

@@ -5,6 +5,12 @@
 import six
 import subprocess
 import time
+import warnings
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
 
 from shutil import rmtree
 from six import raise_from
@@ -48,7 +54,8 @@
     ExecUtilException,  \
     QueryException,     \
     StartNodeException, \
-    TimeoutException
+    TimeoutException,   \
+    TestgresException
 
 from .logger import TestgresLogger
 
@@ -116,7 +123,11 @@ def __exit__(self, type, value, traceback):
 
     @property
     def pid(self):
-        return self.get_pid()
+        return self.get_main_pid()
+
+    @property
+    def auxiliary_pids(self):
+        return self.get_auxiliary_pids()
 
     @property
     def master(self):
@@ -417,7 +428,7 @@ def status(self):
             elif e.exit_code == 4:
                 return NodeStatus.Uninitialized
 
-    def get_pid(self):
+    def get_main_pid(self):
         """
         Return postmaster's PID if node is running, else 0.
         """
@@ -428,7 +439,73 @@ def get_pid(self):
                 return int(f.readline())
 
         # for clarity
-        return 0
+        return None
+
+    def get_child_processes(self):
+        ''' Returns child processes for this node '''
+
+        if psutil is None:
+            warnings.warn("psutil module is not installed")
+            return None
+
+        try:
+            postmaster = psutil.Process(self.pid)
+        except psutil.NoSuchProcess:
+            return None
+
+        return postmaster.children(recursive=True)
+
+    def get_auxiliary_pids(self):
+        ''' Returns dict with pids of auxiliary processes '''
+
+        children = self.get_child_processes()
+        if children is None:
+            return None
+
+        result = {}
+        for child in children:
+            line = child.cmdline()[0]
+            if line.startswith('postgres: checkpointer'):
+                result['checkpointer'] = child.pid
+            elif line.startswith('postgres: background writer'):
+                result['bgwriter'] = child.pid
+            elif line.startswith('postgres: walwriter'):
+                result['walwriter'] = child.pid
+            elif line.startswith('postgres: autovacuum launcher'):
+                result['autovacuum_launcher'] = child.pid
+            elif line.startswith('postgres: stats collector'):
+                result['stats'] = child.pid
+            elif line.startswith('postgres: logical replication launcher'):
+                result['logical_replication_launcher'] = child.pid
+            elif line.startswith('postgres: walreceiver'):
+                result['walreceiver'] = child.pid
+            elif line.startswith('postgres: walsender'):
+                result.setdefault('walsenders', [])
+                result['walsenders'].append(child.pid)
+            elif line.startswith('postgres: startup'):
+                result['startup'] = child.pid
+
+        return result
+
+    def get_walsender_pid(self):
+        ''' Returns pid of according walsender for replica '''
+
+        if not self._master:
+            raise TestgresException("This node is not a replica")
+
+        children = self._master.get_child_processes()
+        if children is None:
+            return None
+
+        sql = 'select application_name, client_port from pg_stat_replication'
+        for name, client_port in self._master.execute(sql):
+            if name == self.name:
+                for child in children:
+                    line = child.cmdline()[0]
+                    if line.startswith('postgres: walsender') and str(client_port) in line:
+                        return child.pid
+
+        return None
 
     def get_control_data(self):
         """

tests/test_simple.py

@@ -205,27 +205,27 @@ def test_status(self):
 
         # check statuses after each operation
         with get_new_node() as node:
-            self.assertEqual(node.pid, 0)
+            self.assertIsNone(node.pid)
             self.assertEqual(node.status(), NodeStatus.Uninitialized)
 
             node.init()
 
-            self.assertEqual(node.pid, 0)
+            self.assertIsNone(node.pid)
             self.assertEqual(node.status(), NodeStatus.Stopped)
 
             node.start()
 
-            self.assertNotEqual(node.pid, 0)
+            self.assertIsNotNone(node.pid)
             self.assertEqual(node.status(), NodeStatus.Running)
 
             node.stop()
 
-            self.assertEqual(node.pid, 0)
+            self.assertIsNone(node.pid)
             self.assertEqual(node.status(), NodeStatus.Stopped)
 
             node.cleanup()
 
-            self.assertEqual(node.pid, 0)
+            self.assertIsNone(node.pid)
             self.assertEqual(node.status(), NodeStatus.Uninitialized)
 
     def test_psql(self):
@@ -702,6 +702,53 @@ def test_version_management(self):
         self.assertTrue(b > c)
         self.assertTrue(a > c)
 
+    def test_pids(self):
+        try:
+            import psutil
+        except ImportError:
+            psutil = None
+
+        master_processes = (
+            'checkpointer',
+            'bgwriter',
+            'walwriter',
+            'autovacuum_launcher',
+            'stats',
+            'logical_replication_launcher',
+            'walsenders',
+        )
+        repl_processes = (
+            'startup',
+            'checkpointer',
+            'bgwriter',
+            'stats',
+            'walreceiver',
+        )
+
+        with get_new_node('master') as master:
+            master.init().start()
+
+            self.assertIsNotNone(master.pid)
+            with master.connect() as con:
+                self.assertTrue(con.backend_pid > 0)
+
+            with master.backup() as backup:
+                with backup.spawn_replica('repl', True) as repl:
+                    repl.start()
+                    if psutil is None:
+                        self.assertIsNone(master.auxiliary_pids)
+                        self.assertIsNone(repl.auxiliary_pids)
+                    else:
+                        master_pids = master.auxiliary_pids
+                        for name in master_processes:
+                            self.assertTrue(name in master_pids)
+                        self.assertTrue(len(master_pids['walsenders']) == 1)
+
+                        repl_pids = repl.auxiliary_pids
+                        for name in repl_processes:
+                            self.assertTrue(name in repl_pids)
+                        self.assertTrue(repl.get_walsender_pid() == master_pids['walsenders'][0])
+
 
 if __name__ == '__main__':
     unittest.main()