testgres 1.10.1__tar.gz → 1.10.2__tar.gz

{testgres-1.10.1/testgres.egg-info → testgres-1.10.2}/PKG-INFO

@@ -1,15 +1,19 @@
 Metadata-Version: 2.1
 Name: testgres
-Version: 1.10.1
+Version: 1.10.2
 Summary: Testing utility for PostgreSQL and its extensions
 Home-page: https://github.com/postgrespro/testgres
 Author: Postgres Professional
 Author-email: testgres@postgrespro.ru
 License: PostgreSQL
 Keywords: test,testing,postgresql
-Platform: UNKNOWN
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: pg8000
+Requires-Dist: port-for>=0.4
+Requires-Dist: six>=1.9.0
+Requires-Dist: psutil
+Requires-Dist: packaging
 
 [![Build Status](https://travis-ci.com/postgrespro/testgres.svg?branch=master)](https://app.travis-ci.com/github/postgrespro/testgres/branches)
 [![codecov](https://codecov.io/gh/postgrespro/testgres/branch/master/graph/badge.svg)](https://codecov.io/gh/postgrespro/testgres)
@@ -220,5 +224,3 @@ def test_basic_query(self):
 [Dmitry Ivanov](https://github.com/funbringer)  
 [Ildus Kurbangaliev](https://github.com/ildus)  
 [Yury Zhuravlev](https://github.com/stalkerg)  
-
-

{testgres-1.10.1 → testgres-1.10.2}/setup.py

@@ -27,7 +27,7 @@ with open('README.md', 'r') as f:
     readme = f.read()
 
 setup(
-    version='1.10.1',
+    version='1.10.2',
     name='testgres',
     packages=['testgres', 'testgres.operations', 'testgres.helpers'],
     description='Testing utility for PostgreSQL and its extensions',

{testgres-1.10.1 → testgres-1.10.2}/testgres/config.py

@@ -2,6 +2,8 @@
 
 import atexit
 import copy
+import logging
+import os
 import tempfile
 
 from contextlib import contextmanager
@@ -10,6 +12,10 @@ from .consts import TMP_CACHE
 from .operations.os_ops import OsOperations
 from .operations.local_ops import LocalOperations
 
+log_level = os.getenv('LOGGING_LEVEL', 'WARNING').upper()
+log_format = os.getenv('LOGGING_FORMAT', '%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=log_level, format=log_format)
+
 
 class GlobalConfig(object):
     """

{testgres-1.10.1 → testgres-1.10.2}/testgres/connection.py

@@ -1,4 +1,5 @@
 # coding: utf-8
+import logging
 
 # we support both pg8000 and psycopg2
 try:
@@ -110,7 +111,7 @@ class NodeConnection(object):
         except ProgrammingError:
             return None
         except Exception as e:
-            print("Error executing query: {}\n {}".format(repr(e), query))
+            logging.error("Error executing query: {}\n {}".format(repr(e), query))
             return None
 
     def close(self):

{testgres-1.10.1 → testgres-1.10.2}/testgres/node.py

@@ -1,5 +1,5 @@
 # coding: utf-8
-
+import logging
 import os
 import random
 import signal
@@ -63,7 +63,6 @@ from .decorators import \
 
 from .defaults import \
     default_dbname, \
-    default_username, \
     generate_app_name
 
 from .exceptions import \
@@ -127,7 +126,7 @@ class ProcessProxy(object):
 
 
 class PostgresNode(object):
-    def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionParams = ConnectionParams(), bin_dir=None, prefix=None):
+    def __init__(self, name=None, base_dir=None, port=None, conn_params: ConnectionParams = ConnectionParams(), bin_dir=None, prefix=None):
         """
         PostgresNode constructor.
 
@@ -156,9 +155,9 @@ class PostgresNode(object):
         else:
             self.os_ops = LocalOperations(conn_params)
 
+        self.host = self.os_ops.host
         self.port = port or reserve_port()
 
-        self.host = self.os_ops.host
         self.ssh_key = self.os_ops.ssh_key
 
         # defaults for __exit__()
@@ -331,22 +330,42 @@ class PostgresNode(object):
         """
         return self._pg_version
 
-    def _try_shutdown(self, max_attempts):
+    def _try_shutdown(self, max_attempts, with_force=False):
         attempts = 0
+        node_pid = self.pid
 
-        # try stopping server N times
-        while attempts < max_attempts:
-            try:
-                self.stop()
-                break    # OK
-            except ExecUtilException:
-                pass    # one more time
-            except Exception:
-                # TODO: probably should kill stray instance
-                eprint('cannot stop node {}'.format(self.name))
-                break
-
-            attempts += 1
+        if node_pid > 0:
+            # try stopping server N times
+            while attempts < max_attempts:
+                try:
+                    self.stop()
+                    break    # OK
+                except ExecUtilException:
+                    pass    # one more time
+                except Exception:
+                    eprint('cannot stop node {}'.format(self.name))
+                    break
+
+                attempts += 1
+
+            # If force stopping is enabled and PID is valid
+            if with_force and node_pid != 0:
+                # If we couldn't stop the node
+                p_status_output = self.os_ops.exec_command(cmd=f'ps -o pid= -p {node_pid}', shell=True, ignore_errors=True).decode('utf-8')
+                if self.status() != NodeStatus.Stopped and p_status_output and str(node_pid) in p_status_output:
+                    try:
+                        eprint(f'Force stopping node {self.name} with PID {node_pid}')
+                        self.os_ops.kill(node_pid, signal.SIGKILL, expect_error=False)
+                    except Exception:
+                        # The node has already stopped
+                        pass
+
+                # Check that node stopped - print only column pid without headers
+                p_status_output = self.os_ops.exec_command(f'ps -o pid= -p {node_pid}', shell=True, ignore_errors=True).decode('utf-8')
+                if p_status_output and str(node_pid) in p_status_output:
+                    eprint(f'Failed to stop node {self.name}.')
+                else:
+                    eprint(f'Node {self.name} has been stopped successfully.')
 
     def _assign_master(self, master):
         """NOTE: this is a private method!"""
@@ -529,7 +548,9 @@ class PostgresNode(object):
                 u"host\treplication\tall\t127.0.0.1/32\t{}\n".format(auth_host),
                 u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host),
                 u"host\treplication\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host),
-                u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host)
+                u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host),
+                u"host\tall\tall\tall\t{}\n".format(auth_host),
+                u"host\treplication\tall\tall\t{}\n".format(auth_host)
             ]  # yapf: disable
 
             # write missing lines
@@ -683,8 +704,6 @@ class PostgresNode(object):
                         If False, waits for the instance to be in primary mode. Default is False.
                max_attempts:
         """
-        if not username:
-            username = default_username()
         self.start()
 
         if replica:
@@ -694,7 +713,7 @@ class PostgresNode(object):
         # Call poll_query_until until the expected value is returned
         self.poll_query_until(query=query,
                               dbname=dbname,
-                              username=username,
+                              username=username or self.os_ops.username,
                               suppress={InternalError,
                                         QueryException,
                                         ProgrammingError,
@@ -737,11 +756,10 @@ class PostgresNode(object):
                 if any(len(file) > 1 and 'Is another postmaster already '
                                          'running on port' in file[1].decode() for
                        file in files):
-                    print("Detected an issue with connecting to port {0}. "
-                          "Trying another port after a 5-second sleep...".format(self.port))
+                    logging.warning("Detected an issue with connecting to port {0}. "
+                                    "Trying another port after a 5-second sleep...".format(self.port))
                     self.port = reserve_port()
-                    options = {}
-                    options['port'] = str(self.port)
+                    options = {'port': str(self.port)}
                     self.set_auto_conf(options)
                     startup_retries -= 1
                     time.sleep(5)
@@ -915,13 +933,14 @@ class PostgresNode(object):
             self._should_free_port = False
             release_port(self.port)
 
-    def cleanup(self, max_attempts=3):
+    def cleanup(self, max_attempts=3, full=False):
         """
         Stop node if needed and remove its data/logs directory.
         NOTE: take a look at TestgresConfig.node_cleanup_full.
 
         Args:
             max_attempts: how many times should we try to stop()?
+            full: clean full base dir
 
         Returns:
             This instance of :class:`.PostgresNode`.
@@ -930,12 +949,12 @@ class PostgresNode(object):
         self._try_shutdown(max_attempts)
 
         # choose directory to be removed
-        if testgres_config.node_cleanup_full:
+        if testgres_config.node_cleanup_full or full:
             rm_dir = self.base_dir    # everything
         else:
             rm_dir = self.data_dir    # just data, save logs
 
-        self.os_ops.rmdirs(rm_dir, ignore_errors=True)
+        self.os_ops.rmdirs(rm_dir, ignore_errors=False)
 
         return self
 
@@ -967,15 +986,13 @@ class PostgresNode(object):
             >>> psql(query='select 3', ON_ERROR_STOP=1)
         """
 
-        # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
 
         psql_params = [
             self._get_bin_path("psql"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username,
+            "-U", username or self.os_ops.username,
             "-X",  # no .psqlrc
             "-A",  # unaligned output
             "-t",  # print rows only
@@ -1087,9 +1104,6 @@ class PostgresNode(object):
                 fname = self.os_ops.mkstemp(prefix=TMP_DUMP)
             return fname
 
-        # Set default arguments
-        dbname = dbname or default_dbname()
-        username = username or default_username()
         filename = filename or tmpfile()
 
         _params = [
@@ -1097,8 +1111,8 @@ class PostgresNode(object):
             "-p", str(self.port),
             "-h", self.host,
             "-f", filename,
-            "-U", username,
-            "-d", dbname,
+            "-U", username or self.os_ops.username,
+            "-d", dbname or default_dbname(),
             "-F", format.value
         ]  # yapf: disable
 
@@ -1118,7 +1132,7 @@ class PostgresNode(object):
 
         # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
+        username = username or self.os_ops.username
 
         _params = [
             self._get_bin_path("pg_restore"),
@@ -1171,7 +1185,6 @@ class PostgresNode(object):
         assert sleep_time > 0
         attempts = 0
         while max_attempts == 0 or attempts < max_attempts:
-            print(f"Pooling {attempts}")
             try:
                 res = self.execute(dbname=dbname,
                                    query=query,
@@ -1195,6 +1208,7 @@ class PostgresNode(object):
                     return    # done
 
             except tuple(suppress or []):
+                logging.info(f"Trying execute, attempt {attempts + 1}.\nQuery: {query}")
                 pass    # we're suppressing them
 
             time.sleep(sleep_time)
@@ -1388,15 +1402,13 @@ class PostgresNode(object):
         if options is None:
             options = []
 
-        # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
 
         _params = [
             self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username,
+            "-U", username or self.os_ops.username
         ] + options  # yapf: disable
 
         # should be the last one
@@ -1463,15 +1475,13 @@ class PostgresNode(object):
             >>> pgbench_run(time=10)
         """
 
-        # Set default arguments
         dbname = dbname or default_dbname()
-        username = username or default_username()
 
         _params = [
             self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
-            "-U", username,
+            "-U", username or self.os_ops.username
         ] + options  # yapf: disable
 
         for key, value in iteritems(kwargs):
@@ -1639,7 +1649,7 @@ class PostgresNode(object):
 
         self.os_ops.write(path, auto_conf, truncate=True)
 
-    def upgrade_from(self, old_node, options=None):
+    def upgrade_from(self, old_node, options=None, expect_error=False):
         """
         Upgrade this node from an old node using pg_upgrade.
 
@@ -1667,11 +1677,11 @@ class PostgresNode(object):
             "--old-datadir", old_node.data_dir,
             "--new-datadir", self.data_dir,
             "--old-port", str(old_node.port),
-            "--new-port", str(self.port),
+            "--new-port", str(self.port)
         ]
         upgrade_command += options
 
-        return self.os_ops.exec_command(upgrade_command)
+        return self.os_ops.exec_command(upgrade_command, expect_error=expect_error)
 
     def _get_bin_path(self, filename):
         if self.bin_dir:
@@ -1683,19 +1693,26 @@ class PostgresNode(object):
 
 class NodeApp:
 
-    def __init__(self, test_path, nodes_to_cleanup, os_ops=LocalOperations()):
-        self.test_path = test_path
-        self.nodes_to_cleanup = nodes_to_cleanup
+    def __init__(self, test_path=None, nodes_to_cleanup=None, os_ops=LocalOperations()):
+        if test_path:
+            if os.path.isabs(test_path):
+                self.test_path = test_path
+            else:
+                self.test_path = os.path.join(os_ops.cwd(), test_path)
+        else:
+            self.test_path = os_ops.cwd()
+        self.nodes_to_cleanup = nodes_to_cleanup if nodes_to_cleanup else []
         self.os_ops = os_ops
 
     def make_empty(
             self,
-            base_dir=None):
+            base_dir=None,
+            port=None):
         real_base_dir = os.path.join(self.test_path, base_dir)
         self.os_ops.rmdirs(real_base_dir, ignore_errors=True)
         self.os_ops.makedirs(real_base_dir)
 
-        node = PostgresNode(base_dir=real_base_dir)
+        node = PostgresNode(base_dir=real_base_dir, port=port)
         node.should_rm_dirs = True
         self.nodes_to_cleanup.append(node)
 
@@ -1704,6 +1721,7 @@ class NodeApp:
     def make_simple(
             self,
             base_dir=None,
+            port=None,
             set_replication=False,
             ptrack_enable=False,
             initdb_params=[],
@@ -1711,7 +1729,7 @@ class NodeApp:
             checksum=True):
         if checksum and '--data-checksums' not in initdb_params:
             initdb_params.append('--data-checksums')
-        node = self.make_empty(base_dir)
+        node = self.make_empty(base_dir, port)
         node.init(
             initdb_params=initdb_params, allow_streaming=set_replication)
 

{testgres-1.10.1 → testgres-1.10.2}/testgres/operations/local_ops.py

@@ -1,9 +1,11 @@
 import getpass
+import logging
 import os
 import shutil
 import stat
 import subprocess
 import tempfile
+import time
 
 import psutil
 
@@ -19,13 +21,18 @@ except ImportError:
 
 CMD_TIMEOUT_SEC = 60
 error_markers = [b'error', b'Permission denied', b'fatal']
+err_out_markers = [b'Failure']
 
 
-def has_errors(output):
+def has_errors(output=None, error=None):
     if output:
         if isinstance(output, str):
             output = output.encode(get_default_encoding())
-        return any(marker in output for marker in error_markers)
+        return any(marker in output for marker in err_out_markers)
+    if error:
+        if isinstance(error, str):
+            error = error.encode(get_default_encoding())
+        return any(marker in error for marker in error_markers)
     return False
 
 
@@ -38,7 +45,7 @@ class LocalOperations(OsOperations):
         self.host = conn_params.host
         self.ssh_key = None
         self.remote = False
-        self.username = conn_params.username or self.get_user()
+        self.username = conn_params.username or getpass.getuser()
 
     @staticmethod
     def _raise_exec_exception(message, command, exit_code, output):
@@ -100,15 +107,16 @@ class LocalOperations(OsOperations):
                 raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
 
     def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=False,
-                     text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None):
+                     text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None,
+                     ignore_errors=False):
         """
         Execute a command in a subprocess and handle the output based on the provided parameters.
         """
         process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
         if get_process:
             return process
-        if process.returncode != 0 or (has_errors(error) and not expect_error):
-            self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode, error)
+        if not ignore_errors and ((process.returncode != 0 or has_errors(output=output, error=error)) and not expect_error):
+            self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode, error or output)
 
         if verbose:
             return process.returncode, output, error
@@ -130,10 +138,6 @@ class LocalOperations(OsOperations):
         # Check if the directory is already in PATH
         os.environ[var_name] = var_val
 
-    # Get environment variables
-    def get_user(self):
-        return self.username or getpass.getuser()
-
     def get_name(self):
         return os.name
 
@@ -146,8 +150,27 @@ class LocalOperations(OsOperations):
         except FileExistsError:
             pass
 
-    def rmdirs(self, path, ignore_errors=True):
-        return rmtree(path, ignore_errors=ignore_errors)
+    def rmdirs(self, path, ignore_errors=True, retries=3, delay=1):
+        """
+        Removes a directory and its contents, retrying on failure.
+
+        :param path: Path to the directory.
+        :param ignore_errors: If True, ignore errors.
+        :param retries: Number of attempts to remove the directory.
+        :param delay: Delay between attempts in seconds.
+        """
+        for attempt in range(retries):
+            try:
+                rmtree(path, ignore_errors=ignore_errors)
+                if not os.path.exists(path):
+                    return True
+            except FileNotFoundError:
+                return True
+            except Exception as e:
+                logging.error(f"Error: Failed to remove directory {path} on attempt {attempt + 1}: {e}")
+            time.sleep(delay)
+        logging.error(f"Error: Failed to remove directory {path} after {retries} attempts.")
+        return False
 
     def listdir(self, path):
         return os.listdir(path)
@@ -271,10 +294,10 @@ class LocalOperations(OsOperations):
         return os.remove(filename)
 
     # Processes control
-    def kill(self, pid, signal):
+    def kill(self, pid, signal, expect_error=False):
         # Kill the process
         cmd = "kill -{} {}".format(signal, pid)
-        return self.exec_command(cmd)
+        return self.exec_command(cmd, expect_error=expect_error)
 
     def get_pid(self):
         # Get current process id

{testgres-1.10.1 → testgres-1.10.2}/testgres/operations/os_ops.py

@@ -1,4 +1,6 @@
+import getpass
 import locale
+import sys
 
 try:
     import psycopg2 as pglib  # noqa: F401
@@ -10,20 +12,23 @@ except ImportError:
 
 
 class ConnectionParams:
-    def __init__(self, host='127.0.0.1', ssh_key=None, username=None):
+    def __init__(self, host='127.0.0.1', port=None, ssh_key=None, username=None):
         self.host = host
+        self.port = port
         self.ssh_key = ssh_key
         self.username = username
 
 
 def get_default_encoding():
-    return locale.getdefaultlocale()[1] or 'UTF-8'
+    if not hasattr(locale, 'getencoding'):
+        locale.getencoding = locale.getpreferredencoding
+    return locale.getencoding() or 'UTF-8'
 
 
 class OsOperations:
     def __init__(self, username=None):
         self.ssh_key = None
-        self.username = username
+        self.username = username or getpass.getuser()
 
     # Command execution
     def exec_command(self, cmd, **kwargs):
@@ -33,6 +38,13 @@ class OsOperations:
     def environ(self, var_name):
         raise NotImplementedError()
 
+    def cwd(self):
+        if sys.platform == 'linux':
+            cmd = 'pwd'
+        elif sys.platform == 'win32':
+            cmd = 'cd'
+        return self.exec_command(cmd).decode().rstrip()
+
     def find_executable(self, executable):
         raise NotImplementedError()
 
@@ -44,9 +56,8 @@ class OsOperations:
         # Check if the directory is already in PATH
         raise NotImplementedError()
 
-    # Get environment variables
     def get_user(self):
-        raise NotImplementedError()
+        return self.username
 
     def get_name(self):
         raise NotImplementedError()