Add functions that return PIDs of various processes

Dockerfile.tmpl

@@ -2,10 +2,12 @@ FROM postgres:${PG_VERSION}-alpine
 
 ENV PYTHON=python${PYTHON_VERSION}
 RUN if [ "${PYTHON_VERSION}" = "2" ] ; then \
-	apk add --no-cache curl python2 py-virtualenv py-pip; \
+	apk add --no-cache curl python2 python2-dev build-base musl-dev \
+    linux-headers py-virtualenv py-pip; \
 	fi
 RUN if [ "${PYTHON_VERSION}" = "3" ] ; then \
-	apk add --no-cache curl python3 py-virtualenv; \
+	apk add --no-cache curl python3 python3-dev build-base musl-dev \
+    linux-headers py-virtualenv; \
 	fi
 ENV LANG=C.UTF-8
 

run_tests.sh

@@ -22,7 +22,7 @@ export VIRTUAL_ENV_DISABLE_PROMPT=1
 source $VENV_PATH/bin/activate
 
 # install utilities
-$PIP install coverage flake8
+$PIP install coverage flake8 psutil
 
 # install testgres' dependencies
 export PYTHONPATH=$(pwd)
@@ -47,7 +47,8 @@ time coverage run -a tests/test_simple.py
 unset PG_BIN
 
 
-# run tests (PG_CONFIG)
+# run tests (PG_CONFIG), also without psutil
+$PIP uninstall -y psutil
 export PG_CONFIG=$(which pg_config)
 time coverage run -a tests/test_simple.py
 unset PG_CONFIG

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._pid = 0
 
         self._connection = pglib.connect(
             database=dbname,
@@ -52,6 +53,12 @@ def node(self):
     def connection(self):
         return self._connection
 
+    @property
+    def pid(self):
+        if not self._pid:
+            self._pid = self.execute("select pg_backend_pid();")[0][0]
+        return self._pid
+
     @property
     def cursor(self):
         return self._cursor

testgres/enums.py

@@ -35,3 +35,18 @@ def __bool__(self):
 
     # for Python 2.x
     __nonzero__ = __bool__
+
+
+class ProcessType(Enum):
+    """
+    Types of postgres processes
+    """
+    Checkpointer = 'postgres: checkpointer'
+    BackgroundWriter = 'postgres: background writer'
+    WalWriter = 'postgres: walwriter'
+    AutovacuumLauncher = 'postgres: autovacuum launcher'
+    StatsCollector = 'postgres: stats collector'
+    LogicalReplicationLauncher = 'postgres: logical replication launcher'
+    WalReceiver = 'postgres: walreceiver'
+    WalSender = 'postgres: walsender'
+    Startup = 'postgres: startup'

testgres/node.py

@@ -6,11 +6,16 @@
 import subprocess
 import time
 
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
 from shutil import rmtree
 from six import raise_from
 from tempfile import mkstemp, mkdtemp
 
-from .enums import NodeStatus
+from .enums import NodeStatus, ProcessType
 
 from .cache import cached_initdb
 
@@ -48,7 +53,8 @@
     ExecUtilException,  \
     QueryException,     \
     StartNodeException, \
-    TimeoutException
+    TimeoutException,   \
+    TestgresException
 
 from .logger import TestgresLogger
 
@@ -116,7 +122,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 +427,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.
         """
@@ -430,6 +440,85 @@ def get_pid(self):
         # for clarity
         return 0
 
+    def get_child_processes(self):
+        ''' Returns child processes for this node '''
+
+        if psutil is None:
+            raise TestgresException("psutil module is not installed")
+
+        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 '''
+
+        alternative_names = {
+            ProcessType.LogicalReplicationLauncher: [
+                'postgres: bgworker: logical replication launcher'
+            ],
+            ProcessType.BackgroundWriter: [
+                'postgres: writer',
+            ],
+            ProcessType.WalWriter: [
+                'postgres: wal writer',
+            ],
+            ProcessType.WalReceiver: [
+                'postgres: wal receiver',
+            ],
+        }
+
+        children = self.get_child_processes()
+        if children is None:
+            return None
+
+        result = {}
+        for child in children:
+            line = ' '.join(child.cmdline())
+            for ptype in ProcessType:
+                if ptype == ProcessType.WalSender \
+                        and (line.startswith(ptype.value) or
+                             line.startswith('postgres: wal sender')):
+                    result.setdefault(ptype, [])
+                    result[ptype].append(child.pid)
+                    break
+                elif line.startswith(ptype.value):
+                    result[ptype] = child.pid
+                    break
+                elif ptype in alternative_names:
+                    names = alternative_names[ptype]
+                    for name in names:
+                        if line.startswith(name):
+                            result[ptype] = child.pid
+                            break
+
+        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 = ' '.join(child.cmdline())
+                    if (line.startswith(ProcessType.WalSender.value) or
+                        line.startswith('postgres: wal sender')) and \
+                            str(client_port) in line:
+                        return child.pid
+
+        return None
+
     def get_control_data(self):
         """
         Return contents of pg_control file.

tests/test_simple.py

@@ -21,7 +21,8 @@
     BackupException, \
     QueryException, \
     CatchUpException, \
-    TimeoutException
+    TimeoutException, \
+    TestgresException
 
 from testgres import \
     TestgresConfig, \
@@ -40,6 +41,7 @@
 
 from testgres import bound_ports
 from testgres.utils import pg_version_ge
+from testgres.enums import ProcessType
 
 
 def util_exists(util):
@@ -702,6 +704,59 @@ 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 = [
+            ProcessType.Checkpointer,
+            ProcessType.BackgroundWriter,
+            ProcessType.WalWriter,
+            ProcessType.AutovacuumLauncher,
+            ProcessType.StatsCollector,
+            ProcessType.WalSender,
+        ]
+        if pg_version_ge('10'):
+            master_processes.append(ProcessType.LogicalReplicationLauncher)
+
+        repl_processes = (
+            ProcessType.Startup,
+            ProcessType.Checkpointer,
+            ProcessType.BackgroundWriter,
+            ProcessType.StatsCollector,
+            ProcessType.WalReceiver,
+        )
+
+        with get_new_node('master') as master:
+            master.init().start()
+
+            self.assertIsNotNone(master.pid)
+            with master.connect() as con:
+                self.assertTrue(con.pid > 0)
+
+            with master.backup() as backup:
+                with backup.spawn_replica('repl', True) as repl:
+                    repl.start()
+                    if psutil is None:
+                        with self.assertRaises(TestgresException):
+                            master.auxiliary_pids
+                        with self.assertRaises(TestgresException):
+                            self.assertIsNone(repl.auxiliary_pids)
+                    else:
+                        master_pids = master.auxiliary_pids
+                        for ptype in master_processes:
+                            self.assertIn(ptype, master_pids)
+                        self.assertTrue(len(master_pids[ProcessType.WalSender]) == 1)
+
+                        repl_pids = repl.auxiliary_pids
+                        for ptype in repl_processes:
+                            self.assertIn(ptype, repl_pids)
+
+                        sender_pid = master_pids[ProcessType.WalSender][0]
+                        self.assertTrue(repl.get_walsender_pid() == sender_pid)
+
 
 if __name__ == '__main__':
     unittest.main()