Merge branch 'master' into logical

zilder · zilder · commit f8b95c603a17 · 2018-03-19T17:43:45.000+03:00
diff --git a/.dockerignore b/.dockerignore
@@ -0,0 +1,4 @@
+dist
+env
+venv
+*.egg-info
diff --git a/Dockerfile.tmpl b/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
 
diff --git a/run_tests.sh b/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)
@@ -42,15 +42,17 @@ time coverage run -a tests/test_simple.py
 
 
 # run tests (PG_BIN)
-export PG_BIN=$(dirname $(which pg_config))
-time coverage run -a tests/test_simple.py
-unset PG_BIN
+time \
+	PG_BIN=$(dirname $(which pg_config)) \
+	ALT_CONFIG=1 \
+	coverage run -a tests/test_simple.py
 
 
 # run tests (PG_CONFIG)
-export PG_CONFIG=$(which pg_config)
-time coverage run -a tests/test_simple.py
-unset PG_CONFIG
+time \
+	PG_CONFIG=$(which pg_config) \
+	ALT_CONFIG=1 \
+	coverage run -a tests/test_simple.py
 
 
 # show coverage
diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
     from distutils.core import setup
 
 # Basic dependencies
-install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0"]
+install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil"]
 
 # Add compatibility enum class
 if sys.version_info < (3, 4):
@@ -17,7 +17,7 @@
     install_requires.append("ipaddress")
 
 setup(
-    version='1.5.0',
+    version='1.6.0',
     name='testgres',
     packages=['testgres'],
     description='Testing utility for PostgreSQL and its extensions',
diff --git a/testgres/connection.py b/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
diff --git a/testgres/enums.py b/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
diff --git a/testgres/node.py b/testgres/node.py
@@ -2,15 +2,15 @@
 
 import io
 import os
-import six
+import psutil
 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
 
@@ -49,7 +49,8 @@
     QueryException,     \
     StartNodeException, \
     TimeoutException,   \
-    InitNodeException
+    InitNodeException,  \
+    TestgresException
 
 from .logger import TestgresLogger
 
@@ -67,6 +68,26 @@
 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 __repr__(self):
+        return '{} : {}'.format(str(self.ptype), repr(self.process))
+
+
 class PostgresNode(object):
     def __init__(self, name=None, port=None, base_dir=None):
         """
@@ -119,7 +140,84 @@ 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.
+        """
+
+        # 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
+
+        msg = "Master doesn't send WAL to {}".format(self.name)
+        raise TestgresException(msg)
 
     @property
     def master(self):
@@ -428,19 +526,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.
@@ -895,7 +980,7 @@ def catchup(self, dbname=None, username=None):
         """
 
         if not self.master:
-            raise CatchUpException("Node doesn't have a master")
+            raise TestgresException("Node doesn't have a master")
 
         if pg_version_ge('10'):
             poll_lsn = "select pg_current_wal_lsn()::text"
@@ -1042,7 +1127,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('_', '-')
 
diff --git a/tests/README.md b/tests/README.md
@@ -21,7 +21,7 @@ export PG_BIN=/path/to/pg/bin
 
 ```bash
 # Set path to PostgreSQL and python version
-export PATH=$PATH:/path/to/pg/bin
+export PATH=/path/to/pg/bin:$PATH
 export PYTHON_VERSION=3  # or 2
 
 # Run tests
diff --git a/tests/test_simple.py b/tests/test_simple.py

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +dist
 +env
 +venv
 +*.egg-info