Aug 4, 2023 · Apr 10, 2023 · Apr 11, 2023 · Apr 11, 2023 · Apr 17, 2023 · May 2, 2023
diff --git a/setup.py b/setup.py
    from distutils.core import setup

 # Basic dependencies
 install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil"]
 install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil", "fabric"]

 # Add compatibility enum class
 if sys.version_info < (3, 4):
    readme = f.read()

 setup(
    version='1.8.6',
    version='1.9.0',
    name='testgres',
    packages=['testgres'],
    description='Testing utility for PostgreSQL and its extensions',
diff --git a/testgres/api.py b/testgres/api.py
 PostgresNode(name='...', port=..., base_dir='...')
 [(3,)]
 """
 from defaults import default_username
 from .node import PostgresNode


 def get_new_node(name=None, base_dir=None, **kwargs):
    """
    Simply a wrapper around :class:`.PostgresNode` constructor.
    See :meth:`.PostgresNode.__init__` for details.
    For remote connection you can add next parameters:
    host='127.0.0.1',
    hostname='localhost',
    ssh_key=None,
    username=default_username()
    """
    # NOTE: leave explicit 'name' and 'base_dir' for compatibility
    return PostgresNode(name=name, base_dir=base_dir, **kwargs)
diff --git a/testgres/backup.py b/testgres/backup.py
 from six import raise_from
 from tempfile import mkdtemp

 from os_ops import OsOperations
 from .enums import XLogMethod

 from .consts import \
            username: database user name.
            xlog_method: none | fetch | stream (see docs)
        """

        self.os_ops = node.os_ops
        if not node.status():
            raise BackupException('Node must be running')

                raise BackupException(msg)

        # Set default arguments
        username = username ordefault_username()
        base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP)
        username = username orself.os_ops.get_user()
        base_dir = base_dir orself.os_ops.mkdtemp(prefix=TMP_BACKUP)

        # public
        self.original_node = node
            "-D", data_dir,
            "-X", xlog_method.value
        ]  # yapf: disable
        execute_utility(_params, self.log_file)
        execute_utility(_params, self.log_file, hostname=node.hostname, ssh_key=node.ssh_key)

    def __enter__(self):
        return self
        available = not destroy

        if available:
            dest_base_dir = mkdtemp(prefix=TMP_NODE)
            dest_base_dir =self.os_ops.mkdtemp(prefix=TMP_NODE)

            data1 = os.path.join(self.base_dir, DATA_DIR)
            data2 = os.path.join(dest_base_dir, DATA_DIR)

        if self._available:
            self._available = False
 rmtree(self.base_dir, ignore_errors=True)
 self.os_ops.rmdirs(self.base_dir, ignore_errors=True)
diff --git a/testgres/cache.py b/testgres/cache.py
 from shutil import copytree
 from six import raise_from

 from os_ops import OsOperations
 from .config import testgres_config

 from .consts import XLOG_CONTROL_FILE
    execute_utility


 def cached_initdb(data_dir, logfile=None, params=None):
 def cached_initdb(data_dir, logfile=None,hostname='localhost', ssh_key=None,params=None):
    """
    Perform initdb or use cached node files.
    """
    os_ops = OsOperations(hostname=hostname, ssh_key=ssh_key)

    def call_initdb(initdb_dir, log=None):
        try:
            _params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
            execute_utility(_params + (params or []), log)
            execute_utility(_params + (params or []), log, hostname=hostname, ssh_key=ssh_key)
        except ExecUtilException as e:
            raise_from(InitNodeException("Failed to run initdb"), e)

        cached_data_dir = testgres_config.cached_initdb_dir

        # Initialize cached initdb
        if not os.path.exists(cached_data_dir) or \
           not os.listdir(cached_data_dir):

        if not os_ops.path_exists(cached_data_dir) or \
           not os_ops.listdir(cached_data_dir):
            call_initdb(cached_data_dir)

        try:
            # Copy cached initdb to current data dir
            copytree(cached_data_dir, data_dir)
 os_ops.copytree(cached_data_dir, data_dir)

            # Assign this node a unique system id if asked to
            if testgres_config.cached_initdb_unique:
                # XXX: write new unique system id to control file
                # Some users might rely upon unique system ids, but
                # our initdb caching mechanism breaks this contract.
                pg_control = os.path.join(data_dir, XLOG_CONTROL_FILE)
 with io.open(pg_control, "r+b") as f:
    f.write(generate_system_id())    # overwrite id
 system_id = generate_system_id()
 os_ops.write(pg_control, system_id, truncate=True, binary=True, read_and_write=True)

                # XXX: build new WAL segment with our system id
                _params = [get_bin_path("pg_resetwal"), "-D", data_dir, "-f"]
                execute_utility(_params, logfile)
                execute_utility(_params, logfile, hostname=hostname, ssh_key=ssh_key)

        except ExecUtilException as e:
            msg = "Failed to reset WAL for system id"
diff --git a/testgres/config.py b/testgres/config.py

    _cached_initdb_dir = None
    """ underlying class attribute for cached_initdb_dir property """

    os_ops = None
    """ OsOperation object that allows work on remote host """
    @property
    def cached_initdb_dir(self):
        """ path to a temp directory for cached initdb. """
    def cached_initdb_dir(self, value):
        self._cached_initdb_dir = value

        # NOTE: assign initial cached dir for initdb
        if self.os_ops:
            testgres_config.cached_initdb_dir = self.os_ops.mkdtemp(prefix=TMP_CACHE)
        else:
            testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)

        if value:
            cached_initdb_dirs.add(value)
        return testgres_config.cached_initdb_dir

    @property
    def temp_dir(self):


 @atexit.register
 def _rm_cached_initdb_dirs():
 def _rm_cached_initdb_dirs(os_ops=None):
    for d in cached_initdb_dirs:
        rmtree(d, ignore_errors=True)
        if os_ops:
            os_ops.rmtree(d, ignore_errors=True)
        else:
            rmtree(d, ignore_errors=True)


 def push_config(**options):
    """

    testgres_config.update(options)


 # NOTE: assign initial cached dir for initdb
 testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
diff --git a/testgres/connection.py b/testgres/connection.py
 # coding: utf-8
 from os_ops import OsOperations

 # we support both pg8000 and psycopg2
 try:

        self._node = node

        self._connection = pglib.connect(database=dbname,
                                         user=username,
                                         password=password,
                                         host=node.host,
                                         port=node.port)
        self.os_ops = OsOperations(node.host, node.hostname, node.ssh_key, node.username)
        self._connection = self.os_ops.db_connect(dbname=dbname,
                                                  user=username,
                                                  password=password,
                                                  host=node.host,
                                                  port=node.port)

        self._connection.autocommit = autocommit
        self._cursor = self.connection.cursor()
        return self

    def execute(self, query, *args):
        self.cursor.execute(query, args)

        try:
            res = self.cursor.fetchall()

            # pg8000 might return tuples
            if isinstance(res, tuple):
                res = [tuple(t) for t in res]

            return res
        except Exception:
            with self.connection.cursor() as cursor:
                cursor.execute(query, args)
                try:
                    res = cursor.fetchall()

                    # pg8000 might return tuples
                    if isinstance(res, tuple):
                        res = [tuple(t) for t in res]

                    return res
                except (pglib.ProgrammingError, pglib.InternalError) as e:
                    # An error occurred while trying to fetch results (e.g., no results to fetch)
                    print(f"Error fetching results: {e}")
                    return None
        except (pglib.Error, Exception) as e:
            # Handle other database errors
            print(f"Error executing query: {e}")
            return None

    def close(self):
diff --git a/testgres/defaults.py b/testgres/defaults.py
    return 'postgres'


 def default_username():
 def default_username(os_ops=None):
    """
    Return default username (current user).
    """

    return getpass.getuser()
    if os_ops:
        user = os_ops.get_user()
    else:
        user = getpass.getuser()
    return user


 def generate_app_name():
    return 'testgres-{}'.format(str(uuid.uuid4()))


 def generate_system_id():
 def generate_system_id(os_ops=None):
    """
    Generate a new 64-bit unique system identifier for node.
    """
    system_id = 0
    system_id |= (secs << 32)
    system_id |= (usecs << 12)
    system_id |= (os.getpid() & 0xFFF)
    if os_ops:
        system_id |= (os_ops.get_pid() & 0xFFF)
    else:
        system_id |= (os.getpid() & 0xFFF)

    # pack ULL in native byte order
    return struct.pack('=Q', system_id)
diff --git a/testgres/logger.py b/testgres/logger.py
 import time


 # create logger
 log = logging.getLogger('Testgres')
 log.setLevel(logging.DEBUG)
 # create console handler and set level to debug
 ch = logging.StreamHandler()
 ch.setLevel(logging.DEBUG)
 # create formatter
 formatter = logging.Formatter('\n%(asctime)s - %(name)s[%(levelname)s]: %(message)s')
 # add formatter to ch
 ch.setFormatter(formatter)
 # add ch to logger
 log.addHandler(ch)


 class TestgresLogger(threading.Thread):
    """
    Helper class to implement reading from log files.
diff --git a/testgres/node.py b/testgres/node.py
 import subprocess
 import time

 from os_ops import OsOperations

 try:
    from collections.abc import Iterable


 class PostgresNode(object):
    def __init__(self, name=None, port=None, base_dir=None):
    def __init__(self, name=None, port=None, base_dir=None,
                 host='127.0.0.1', hostname='localhost', ssh_key=None, username=default_username()):
        """
        PostgresNode constructor.

        self._master = None

        # basic
        self.host = '127.0.0.1'
        self.name = name or generate_app_name()
        self.port = port or reserve_port()

        self.host = host
        self.hostname = hostname
        self.ssh_key = ssh_key
        self.os_ops = OsOperations(host, hostname, ssh_key, username=username)

        # defaults for __exit__()
        self.cleanup_on_good_exit = testgres_config.node_cleanup_on_good_exit
        self.cleanup_on_bad_exit = testgres_config.node_cleanup_on_bad_exit
        """

        # initialize this PostgreSQL node
        cached_initdb(data_dir=self.data_dir,
                      logfile=self.utils_log_file,
                      params=initdb_params)
        cached_initdb(
            data_dir=self.data_dir,
            logfile=self.utils_log_file,
            hostname=self.hostname,
            ssh_key=self.ssh_key,
            params=initdb_params)

        # initialize default config files
        self.default_conf(**kwargs)
                new_lines = [
                    u"local\treplication\tall\t\t\t{}\n".format(auth_local),
                    u"host\treplication\tall\t127.0.0.1/32\t{}\n".format(auth_host),

                    u"host\treplication\tall\t0.0.0.0/0\t{}\n".format(auth_host),
                    u"host\tall\tall\t0.0.0.0/0\t{}\n".format(auth_host),

                    u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host)
                ]  # yapf: disable
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,7 +6,7 @@
		from distutils.core import setup

		# Basic dependencies
		install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil"]
		install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil", "fabric"]

		# Add compatibility enum class
		if sys.version_info < (3, 4):
Expand All		@@ -21,7 +21,7 @@
		readme = f.read()

		setup(
		version='1.8.6',
		version='1.9.0',
		name='testgres',
		packages=['testgres'],
		description='Testing utility for PostgreSQL and its extensions',
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -30,13 +30,19 @@
		PostgresNode(name='...', port=..., base_dir='...')
		[(3,)]
		"""
		from defaults import default_username
		from .node import PostgresNode


		def get_new_node(name=None, base_dir=None, **kwargs):
		"""
		Simply a wrapper around :class:`.PostgresNode` constructor.
		See :meth:`.PostgresNode.__init__` for details.
		For remote connection you can add next parameters:
		host='127.0.0.1',
		hostname='localhost',
		ssh_key=None,
		username=default_username()
		"""
		# NOTE: leave explicit 'name' and 'base_dir' for compatibility
		return PostgresNode(name=name, base_dir=base_dir, **kwargs)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,6 +6,7 @@
		from six import raise_from
		from tempfile import mkdtemp

		from os_ops import OsOperations
		from .enums import XLogMethod

		from .consts import \
Expand DownExpand Up		@@ -47,7 +48,7 @@ def __init__(self,
		username: database user name.
		xlog_method: none \| fetch \| stream (see docs)
		"""

		self.os_ops = node.os_ops
		if not node.status():
		raise BackupException('Node must be running')

Expand All		@@ -60,8 +61,8 @@ def __init__(self,
		raise BackupException(msg)

		# Set default arguments
		username = username ordefault_username()
		base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP)
		username = username orself.os_ops.get_user()
		base_dir = base_dir orself.os_ops.mkdtemp(prefix=TMP_BACKUP)

		# public
		self.original_node = node
Expand All		@@ -81,7 +82,7 @@ def __init__(self,
		"-D", data_dir,
		"-X", xlog_method.value
		] # yapf: disable
		execute_utility(_params, self.log_file)
		execute_utility(_params, self.log_file, hostname=node.hostname, ssh_key=node.ssh_key)

		def __enter__(self):
		return self
Expand All		@@ -107,7 +108,7 @@ def _prepare_dir(self, destroy):
		available = not destroy

		if available:
		dest_base_dir = mkdtemp(prefix=TMP_NODE)
		dest_base_dir =self.os_ops.mkdtemp(prefix=TMP_NODE)

		data1 = os.path.join(self.base_dir, DATA_DIR)
		data2 = os.path.join(dest_base_dir, DATA_DIR)
Expand DownExpand Up		@@ -185,4 +186,4 @@ def cleanup(self):

		if self._available:
		self._available = False
		rmtree(self.base_dir, ignore_errors=True)
		self.os_ops.rmdirs(self.base_dir, ignore_errors=True)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,6 +6,7 @@
		from shutil import copytree
		from six import raise_from

		from os_ops import OsOperations
		from .config import testgres_config

		from .consts import XLOG_CONTROL_FILE
Expand All		@@ -21,14 +22,16 @@
		execute_utility


		def cached_initdb(data_dir, logfile=None, params=None):
		def cached_initdb(data_dir, logfile=None,hostname='localhost', ssh_key=None,params=None):
		"""
		Perform initdb or use cached node files.
		"""
		os_ops = OsOperations(hostname=hostname, ssh_key=ssh_key)

		def call_initdb(initdb_dir, log=None):
		try:
		_params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
		execute_utility(_params + (params or []), log)
		execute_utility(_params + (params or []), log, hostname=hostname, ssh_key=ssh_key)
		except ExecUtilException as e:
		raise_from(InitNodeException("Failed to run initdb"), e)

Expand All		@@ -39,26 +42,27 @@ def call_initdb(initdb_dir, log=None):
		cached_data_dir = testgres_config.cached_initdb_dir

		# Initialize cached initdb
		if not os.path.exists(cached_data_dir) or \
		not os.listdir(cached_data_dir):

		if not os_ops.path_exists(cached_data_dir) or \
		not os_ops.listdir(cached_data_dir):
		call_initdb(cached_data_dir)

		try:
		# Copy cached initdb to current data dir
		copytree(cached_data_dir, data_dir)
		os_ops.copytree(cached_data_dir, data_dir)

		# Assign this node a unique system id if asked to
		if testgres_config.cached_initdb_unique:
		# XXX: write new unique system id to control file
		# Some users might rely upon unique system ids, but
		# our initdb caching mechanism breaks this contract.
		pg_control = os.path.join(data_dir, XLOG_CONTROL_FILE)
		with io.open(pg_control, "r+b") as f:
		f.write(generate_system_id()) # overwrite id
		system_id = generate_system_id()
		os_ops.write(pg_control, system_id, truncate=True, binary=True, read_and_write=True)

		# XXX: build new WAL segment with our system id
		_params = [get_bin_path("pg_resetwal"), "-D", data_dir, "-f"]
		execute_utility(_params, logfile)
		execute_utility(_params, logfile, hostname=hostname, ssh_key=ssh_key)

		except ExecUtilException as e:
		msg = "Failed to reset WAL for system id"
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -43,6 +43,9 @@ class GlobalConfig(object):

		_cached_initdb_dir = None
		""" underlying class attribute for cached_initdb_dir property """

		os_ops = None
		""" OsOperation object that allows work on remote host """
		@property
		def cached_initdb_dir(self):
		""" path to a temp directory for cached initdb. """
Expand All		@@ -52,8 +55,15 @@ def cached_initdb_dir(self):
		def cached_initdb_dir(self, value):
		self._cached_initdb_dir = value

		# NOTE: assign initial cached dir for initdb
		if self.os_ops:
		testgres_config.cached_initdb_dir = self.os_ops.mkdtemp(prefix=TMP_CACHE)
		else:
		testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)

		if value:
		cached_initdb_dirs.add(value)
		return testgres_config.cached_initdb_dir

		@property
		def temp_dir(self):
Expand DownExpand Up		@@ -133,9 +143,12 @@ def copy(self):


		@atexit.register
		def _rm_cached_initdb_dirs():
		def _rm_cached_initdb_dirs(os_ops=None):
		for d in cached_initdb_dirs:
		rmtree(d, ignore_errors=True)
		if os_ops:
		os_ops.rmtree(d, ignore_errors=True)
		else:
		rmtree(d, ignore_errors=True)


		def push_config(**options):
Expand DownExpand Up		@@ -195,7 +208,3 @@ def configure_testgres(**options):
		"""

		testgres_config.update(options)


		# NOTE: assign initial cached dir for initdb
		testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
Original file line number	Diff line number	Diff line change
		@@ -1,4 +1,5 @@
		# coding: utf-8
		from os_ops import OsOperations

		# we support both pg8000 and psycopg2
		try:
Expand DownExpand Up		@@ -41,11 +42,12 @@ def __init__(self,

		self._node = node

		self._connection = pglib.connect(database=dbname,
		user=username,
		password=password,
		host=node.host,
		port=node.port)
		self.os_ops = OsOperations(node.host, node.hostname, node.ssh_key, node.username)
		self._connection = self.os_ops.db_connect(dbname=dbname,
		user=username,
		password=password,
		host=node.host,
		port=node.port)

		self._connection.autocommit = autocommit
		self._cursor = self.connection.cursor()
Expand DownExpand Up		@@ -102,17 +104,24 @@ def rollback(self):
		return self

		def execute(self, query, *args):
		self.cursor.execute(query, args)

		try:
		res = self.cursor.fetchall()

		# pg8000 might return tuples
		if isinstance(res, tuple):
		res = [tuple(t) for t in res]

		return res
		except Exception:
		with self.connection.cursor() as cursor:
		cursor.execute(query, args)
		try:
		res = cursor.fetchall()

		# pg8000 might return tuples
		if isinstance(res, tuple):
		res = [tuple(t) for t in res]

		return res
		except (pglib.ProgrammingError, pglib.InternalError) as e:
		# An error occurred while trying to fetch results (e.g., no results to fetch)
		print(f"Error fetching results: {e}")
		return None
		except (pglib.Error, Exception) as e:
		# Handle other database errors
		print(f"Error executing query: {e}")
		return None

		def close(self):
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -13,12 +13,15 @@ def default_dbname():
		return 'postgres'


		def default_username():
		def default_username(os_ops=None):
		"""
		Return default username (current user).
		"""

		return getpass.getuser()
		if os_ops:
		user = os_ops.get_user()
		else:
		user = getpass.getuser()
		return user


		def generate_app_name():
Expand All		@@ -29,7 +32,7 @@ def generate_app_name():
		return 'testgres-{}'.format(str(uuid.uuid4()))


		def generate_system_id():
		def generate_system_id(os_ops=None):
		"""
		Generate a new 64-bit unique system identifier for node.
		"""
Expand All		@@ -44,7 +47,10 @@ def generate_system_id():
		system_id = 0
		system_id \|= (secs << 32)
		system_id \|= (usecs << 12)
		system_id \|= (os.getpid() & 0xFFF)
		if os_ops:
		system_id \|= (os_ops.get_pid() & 0xFFF)
		else:
		system_id \|= (os.getpid() & 0xFFF)

		# pack ULL in native byte order
		return struct.pack('=Q', system_id)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,6 +6,20 @@
		import time


		# create logger
		log = logging.getLogger('Testgres')
		log.setLevel(logging.DEBUG)
		# create console handler and set level to debug
		ch = logging.StreamHandler()
		ch.setLevel(logging.DEBUG)
		# create formatter
		formatter = logging.Formatter('\n%(asctime)s - %(name)s[%(levelname)s]: %(message)s')
		# add formatter to ch
		ch.setFormatter(formatter)
		# add ch to logger
		log.addHandler(ch)


		class TestgresLogger(threading.Thread):
		"""
		Helper class to implement reading from log files.
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -12,6 +12,7 @@
		import subprocess
		import time

		from os_ops import OsOperations

		try:
		from collections.abc import Iterable
Expand DownExpand Up		@@ -129,7 +130,8 @@ def __repr__(self):


		class PostgresNode(object):
		def __init__(self, name=None, port=None, base_dir=None):
		def __init__(self, name=None, port=None, base_dir=None,
		host='127.0.0.1', hostname='localhost', ssh_key=None, username=default_username()):
		"""
		PostgresNode constructor.

Expand All		@@ -147,10 +149,14 @@ def __init__(self, name=None, port=None, base_dir=None):
		self._master = None

		# basic
		self.host = '127.0.0.1'
		self.name = name or generate_app_name()
		self.port = port or reserve_port()

		self.host = host
		self.hostname = hostname
		self.ssh_key = ssh_key
		self.os_ops = OsOperations(host, hostname, ssh_key, username=username)

		# defaults for __exit__()
		self.cleanup_on_good_exit = testgres_config.node_cleanup_on_good_exit
		self.cleanup_on_bad_exit = testgres_config.node_cleanup_on_bad_exit
Expand DownExpand Up		@@ -455,9 +461,12 @@ def init(self, initdb_params=None, **kwargs):
		"""

		# initialize this PostgreSQL node
		cached_initdb(data_dir=self.data_dir,
		logfile=self.utils_log_file,
		params=initdb_params)
		cached_initdb(
		data_dir=self.data_dir,
		logfile=self.utils_log_file,
		hostname=self.hostname,
		ssh_key=self.ssh_key,
		params=initdb_params)

		# initialize default config files
		self.default_conf(**kwargs)
Expand DownExpand Up		@@ -514,6 +523,10 @@ def get_auth_method(t):
		new_lines = [
		u"local\treplication\tall\t\t\t{}\n".format(auth_local),
		u"host\treplication\tall\t127.0.0.1/32\t{}\n".format(auth_host),

		u"host\treplication\tall\t0.0.0.0/0\t{}\n".format(auth_host),
		u"host\tall\tall\t0.0.0.0/0\t{}\n".format(auth_host),

		u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host)
		] # yapf: disable

Expand Down