Merge branch 'ildus-master'

Author: funbringer

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)

testgres/connection.py

@@ -52,6 +52,10 @@ def node(self):
     def connection(self):
         return self._connection
 
+    @property
+    def pid(self):
+        return self.execute("select pg_backend_pid()")[0][0]
+
     @property
     def cursor(self):
         return self._cursor

testgres/enums.py

@@ -1,4 +1,5 @@
 from enum import Enum, IntEnum
+from six import iteritems
 
 
 class XLogMethod(Enum):
@@ -35,3 +36,53 @@ def __bool__(self):
 
     # for Python 2.x
     __nonzero__ = __bool__
+
+
+class ProcessType(Enum):
+    """
+    Types of processes
+    """
+
+    AutovacuumLauncher = 'autovacuum launcher'
+    BackgroundWriter = 'background writer'
+    Checkpointer = 'checkpointer'
+    LogicalReplicationLauncher = 'logical replication launcher'
+    Startup = 'startup'
+    StatsCollector = 'stats collector'
+    WalReceiver = 'wal receiver'
+    WalSender = 'wal sender'
+    WalWriter = 'wal writer'
+
+    # special value
+    Unknown = 'unknown'
+
+    @staticmethod
+    def from_process(process):
+        # yapf: disable
+        # legacy names for older releases of PG
+        alternative_names = {
+            ProcessType.LogicalReplicationLauncher: [
+                'logical replication worker'
+            ],
+            ProcessType.BackgroundWriter: [
+                'writer'
+            ],
+        }
+
+        # we deliberately cut special words and spaces
+        cmdline = ''.join(process.cmdline()) \
+                    .replace('postgres:', '', 1) \
+                    .replace('bgworker:', '', 1) \
+                    .replace(' ', '')
+
+        for ptype in ProcessType:
+            if cmdline.startswith(ptype.value.replace(' ', '')):
+                return ptype
+
+        for ptype, names in iteritems(alternative_names):
+            for name in names:
+                if cmdline.startswith(name.replace(' ', '')):
+                    return ptype
+
+        # default
+        return ProcessType.Unknown

testgres/node.py

@@ -2,15 +2,14 @@
 
 import io
 import os
-import six
 import subprocess
 import time
 
 from shutil import rmtree
-from six import raise_from
+from six import raise_from, iteritems
 from tempfile import mkstemp, mkdtemp
 
-from .enums import NodeStatus
+from .enums import NodeStatus, ProcessType
 
 from .cache import cached_initdb
 
@@ -48,7 +47,8 @@
     ExecUtilException,  \
     QueryException,     \
     StartNodeException, \
-    TimeoutException
+    TimeoutException,   \
+    TestgresException
 
 from .logger import TestgresLogger
 
@@ -64,6 +64,28 @@
 from .backup import NodeBackup
 
 
+class ProcessProxy(object):
+    """
+    Wrapper for psutil.Process
+
+    Attributes:
+        process: wrapped psutill.Process object
+        ptype: instance of ProcessType
+    """
+
+    def __init__(self, process):
+        self.process = process
+        self.ptype = ProcessType.from_process(process)
+
+    def __getattr__(self, name):
+        return getattr(self.process, name)
+
+    def __str__(self):
+        pid = self.process.pid
+        cmdline = ' '.join(self.process.cmdline()).strip()
+        return '{} [{}]'.format(cmdline, pid)
+
+
 class PostgresNode(object):
     def __init__(self, name=None, port=None, base_dir=None):
         """
@@ -116,7 +138,88 @@ def __exit__(self, type, value, traceback):
 
     @property
     def pid(self):
-        return self.get_pid()
+        """
+        Return postmaster's PID if node is running, else 0.
+        """
+
+        if self.status():
+            pid_file = os.path.join(self.data_dir, PG_PID_FILE)
+            with io.open(pid_file) as f:
+                return int(f.readline())
+
+        # for clarity
+        return 0
+
+    @property
+    def auxiliary_pids(self):
+        """
+        Returns a dict of { ProcessType : PID }.
+        """
+
+        result = {}
+
+        for process in self.auxiliary_processes:
+            if process.ptype not in result:
+                result[process.ptype] = []
+
+            result[process.ptype].append(process.pid)
+
+        return result
+
+    @property
+    def auxiliary_processes(self):
+        """
+        Returns a list of auxiliary processes.
+        Each process is represented by ProcessProxy object.
+        """
+
+        def is_aux(process):
+            return process.ptype != ProcessType.Unknown
+
+        return list(filter(is_aux, self.child_processes))
+
+    @property
+    def child_processes(self):
+        """
+        Returns a list of all child processes.
+        Each process is represented by ProcessProxy object.
+        """
+
+        try:
+            import psutil
+        except ImportError:
+            raise TestgresException("psutil module is not installed")
+
+        # get a list of postmaster's children
+        children = psutil.Process(self.pid).children()
+
+        return [ProcessProxy(p) for p in children]
+
+    @property
+    def source_walsender(self):
+        """
+        Returns master's walsender feeding this replica.
+        """
+
+        sql = """
+            select pid
+            from pg_catalog.pg_stat_replication
+            where application_name = $1
+        """
+
+        if not self.master:
+            raise TestgresException("Node doesn't have a master")
+
+        # master should be on the same host
+        assert self.master.host == self.host
+
+        with self.master.connect() as con:
+            for row in con.execute(sql, self.name):
+                for child in self.master.auxiliary_processes:
+                    if child.pid == int(row[0]):
+                        return child
+
+        raise QueryException("Master doesn't send WAL to {}", self.name)
 
     @property
     def master(self):
@@ -417,19 +520,6 @@ def status(self):
             elif e.exit_code == 4:
                 return NodeStatus.Uninitialized
 
-    def get_pid(self):
-        """
-        Return postmaster's PID if node is running, else 0.
-        """
-
-        if self.status():
-            pid_file = os.path.join(self.data_dir, PG_PID_FILE)
-            with io.open(pid_file) as f:
-                return int(f.readline())
-
-        # for clarity
-        return 0
-
     def get_control_data(self):
         """
         Return contents of pg_control file.
@@ -990,7 +1080,7 @@ def pgbench_run(self,
             "-U", username,
         ] + options
 
-        for key, value in six.iteritems(kwargs):
+        for key, value in iteritems(kwargs):
             # rename keys for pgbench
             key = key.replace('_', '-')
 

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):
@@ -59,6 +61,14 @@ def good_properties(f):
             return True
 
 
+def module_exists(module):
+    try:
+        __import__(module)
+        return True
+    except ImportError:
+        return False
+
+
 @contextmanager
 def removing(f):
     try:
@@ -702,6 +712,56 @@ def test_version_management(self):
         self.assertTrue(b > c)
         self.assertTrue(a > c)
 
+    @unittest.skipUnless(module_exists('psutil'), 'might be missing')
+    def test_child_pids(self):
+        master_processes = [
+            ProcessType.AutovacuumLauncher,
+            ProcessType.BackgroundWriter,
+            ProcessType.Checkpointer,
+            ProcessType.StatsCollector,
+            ProcessType.WalSender,
+            ProcessType.WalWriter,
+        ]
+
+        if pg_version_ge('10'):
+            master_processes.append(ProcessType.LogicalReplicationLauncher)
+
+        repl_processes = [
+            ProcessType.Startup,
+            ProcessType.WalReceiver,
+        ]
+
+        with get_new_node().init().start() as master:
+
+            # master node doesn't have a source walsender!
+            with self.assertRaises(TestgresException):
+                master.source_walsender
+
+            with master.replicate().start() as replica:
+
+                # test __str__ method
+                str(master.child_processes[0])
+
+                master_pids = master.auxiliary_pids
+                for ptype in master_processes:
+                    self.assertIn(ptype, master_pids)
+
+                replica_pids = replica.auxiliary_pids
+                for ptype in repl_processes:
+                    self.assertIn(ptype, replica_pids)
+
+                # there should be exactly 1 source walsender for replica
+                self.assertEqual(len(master_pids[ProcessType.WalSender]), 1)
+                pid1 = master_pids[ProcessType.WalSender][0]
+                pid2 = replica.source_walsender.pid
+                self.assertEqual(pid1, pid2)
+
+                replica.stop()
+
+                # there should be no walsender after we've stopped replica
+                with self.assertRaises(QueryException):
+                    replica.source_walsender
+
 
 if __name__ == '__main__':
     unittest.main()