testgres 1.16.4__tar.gz → 1.17.0__tar.gz

{testgres-1.16.4/testgres.egg-info → testgres-1.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: testgres
-Version: 1.16.4
+Version: 1.17.0
 Summary: Testing utility for PostgreSQL and its extensions
 Author-email: Postgres Professional <testgres@postgrespro.ru>
 License: PostgreSQL

{testgres-1.16.4 → testgres-1.17.0}/src/__init__.py

@@ -24,6 +24,7 @@ from .exceptions import \
     TimeoutException, \
     CatchUpException, \
     StartNodeException, \
+    StopNodeException, \
     InitNodeException, \
     BackupException, \
     InvalidOperationException
@@ -60,7 +61,7 @@ from testgres.operations.os_ops import OsProcessController
 from testgres.operations.local_ops import LocalOperations
 from testgres.operations.remote_ops import RemoteOperations
 
-__version__ = "1.16.4"
+__version__ = "1.17.0"
 
 __all__ = [
     "get_new_node",
@@ -73,6 +74,7 @@ __all__ = [
     "QueryTimeoutException",
     "TimeoutException", "CatchUpException", "StartNodeException", "InitNodeException", "BackupException", "InvalidOperationException",
     "XLogMethod", "IsolationLevel", "NodeStatus", "ProcessType", "DumpFormat",
+    "StopNodeException",
     "NodeApp",
     "PostgresNode",
     "PortManager",

{testgres-1.16.4 → testgres-1.17.0}/src/backup.py

@@ -35,7 +35,8 @@ class NodeBackup:
                  base_dir=None,
                  username=None,
                  xlog_method=XLogMethod.fetch,
-                 options=None):
+                 options=None,
+                 os_ops=None):
         """
         Create a new backup.
 
@@ -44,13 +45,18 @@ class NodeBackup:
             base_dir: where should we store it?
             username: database user name.
             xlog_method: none | fetch | stream (see docs)
+            os_ops: execution context where pg_basebackup is run and the
+                backup is stored. If omitted, use the source node's os_ops.
+                The source node must be reachable from this context using its
+                configured host and port.
         """
         assert node.os_ops is not None
         assert isinstance(node.os_ops, OsOperations)
+        assert os_ops is None or isinstance(os_ops, OsOperations)
 
         if not options:
             options = []
-        self.os_ops = node.os_ops
+        self.os_ops = os_ops if os_ops is not None else node.os_ops
         if not node.status():
             raise BackupException('Node must be running')
 
@@ -63,7 +69,9 @@ class NodeBackup:
                 raise BackupException(msg)
 
         # Set default arguments
-        username = username or self.os_ops.get_user()
+        # The username belongs to the source PostgreSQL instance, not to the
+        # environment where pg_basebackup happens to run.
+        username = username or node.os_ops.get_user()
         base_dir = base_dir or self.os_ops.mkdtemp(prefix=TMP_BACKUP)
 
         # public
@@ -139,13 +147,24 @@ class NodeBackup:
         # Return path to new node
         return dest_base_dir
 
-    def spawn_primary(self, name=None, destroy=True):
+    def spawn_primary(
+        self,
+        name=None,
+        destroy=True,
+        target_host=None,
+        target_port=None,
+    ):
         """
         Create a primary node from a backup.
 
         Args:
             name: primary's application name.
             destroy: should we convert this backup into a node?
+            target_host: host/address for the new node. If omitted, preserve
+                the source host when the backup uses the source os_ops,
+                otherwise use the backup os_ops host.
+            target_port: explicit port for the new node. If omitted, reserve a
+                port in the backup's execution context.
 
         Returns:
             New instance of :class:`.PostgresNode`.
@@ -157,19 +176,37 @@ class NodeBackup:
         # Build a new PostgresNode
         assert self.original_node is not None
 
-        node = self.original_node.clone_with_new_name_and_base_dir(name=name, base_dir=base_dir)
+        node = self.original_node.clone_with_new_name_and_base_dir(
+            name=name,
+            base_dir=base_dir,
+            target_os_ops=self.os_ops,
+            target_host=target_host,
+            target_port=target_port,
+        )
 
         assert node is not None
         assert type(node) is self.original_node.__class__
 
         with clean_on_error(node) as node:
-            # Set a new port
+            # The backup contains the source node's connection settings.
+            # Override settings that belong to the spawned node.
             node.append_conf(filename=PG_CONF_FILE, line='\n')
-            node.append_conf(filename=PG_CONF_FILE, port=node.port)
+            node.append_conf(
+                filename=PG_CONF_FILE,
+                listen_addresses=node.host,
+                port=node.port,
+            )
 
             return node
 
-    def spawn_replica(self, name=None, destroy=True, slot=None):
+    def spawn_replica(
+        self,
+        name=None,
+        destroy=True,
+        slot=None,
+        target_host=None,
+        target_port=None,
+    ):
         """
         Create a replica of the original node from a backup.
 
@@ -177,13 +214,23 @@ class NodeBackup:
             name: replica's application name.
             slot: create a replication slot with the specified name.
             destroy: should we convert this backup into a node?
+            target_host: host/address for the replica. If omitted, preserve
+                the source host when the backup uses the source os_ops,
+                otherwise use the backup os_ops host.
+            target_port: explicit port for the replica. If omitted, reserve a
+                port in the backup's execution context.
 
         Returns:
             New instance of :class:`.PostgresNode`.
         """
 
         # Build a new PostgresNode
-        node = self.spawn_primary(name=name, destroy=destroy)
+        node = self.spawn_primary(
+            name=name,
+            destroy=destroy,
+            target_host=target_host,
+            target_port=target_port,
+        )
         assert node is not None
 
         try:

testgres-1.17.0/src/connection.py

@@ -0,0 +1,23 @@
+# coding: utf-8
+# we support both pg8000 and psycopg2
+
+# to backward compatibility
+from .exceptions import NodeConnectionDatabaseException as DatabaseError
+from .exceptions import NodeConnectionOperationalException as OperationalError
+from .exceptions import NodeConnectionInternalException as InternalError
+from .exceptions import NodeConnectionProgrammingException as ProgrammingError
+
+try:
+    from .impl.connection_psycopg2 import NodeConnection, pglib
+except ImportError:
+    try:
+        from .impl.connection_pg8000 import NodeConnection, pglib
+    except ImportError:
+        raise ImportError("You must have psycopg2 or pg8000 modules installed")
+
+assert DatabaseError.__name__ == "NodeConnectionDatabaseException"
+assert OperationalError.__name__ == "NodeConnectionOperationalException"
+assert InternalError.__name__ == "NodeConnectionInternalException"
+assert ProgrammingError.__name__ == "NodeConnectionProgrammingException"
+assert NodeConnection.__name__ == "NodeConnection"
+assert pglib is not None

{testgres-1.16.4 → testgres-1.17.0}/src/exceptions.py

@@ -160,7 +160,7 @@ class CatchUpException(TestgresException):
         return result
 
 
-class StartNodeException(TestgresException):
+class NodeException(TestgresException):
     _description: typing.Optional[str]
     _files: typing.Optional[typing.Iterable]
 
@@ -224,6 +224,14 @@ class StartNodeException(TestgresException):
         return result
 
 
+class StartNodeException(NodeException):
+    pass
+
+
+class StopNodeException(NodeException):
+    pass
+
+
 class InitNodeException(TestgresException):
     _message: typing.Optional[str]
 
@@ -294,6 +302,85 @@ class BackupException(TestgresException):
         return result
 
 
+class NodeConnectionException(TestgresException):
+    _message: typing.Optional[str]
+
+    def __init__(
+        self,
+        message: typing.Optional[str] = None,
+    ):
+        assert message is None or type(message) is str
+        super().__init__(message)
+        self._message = message
+        return
+
+    @property
+    def message(self) -> str:
+        assert self._message is None or type(self._message) is str
+        if self._message is None:
+            return ""
+        return self._message
+
+    def __repr__(self) -> str:
+        args = []
+
+        if self._message is not None:
+            args.append(("message", self._message))
+
+        result = "{}(".format(type(self).__name__)
+        sep = ""
+        for a in args:
+            result += sep + a[0] + "=" + repr(a[1])
+            sep = ", "
+            continue
+        result += ")"
+        return result
+
+
+class NodeConnectionInterfaceException(NodeConnectionException):
+    pass
+
+
+class NodeConnectionDatabaseException(NodeConnectionException):
+    pass
+
+
+class NodeConnectionDataException(NodeConnectionDatabaseException):
+    pass
+
+
+class NodeConnectionOperationalException(NodeConnectionDatabaseException):
+    pass
+
+
+class NodeConnectionIntegrityException(NodeConnectionDatabaseException):
+    pass
+
+
+class NodeConnectionInternalException(NodeConnectionDatabaseException):
+    pass
+
+
+class NodeConnectionProgrammingException(NodeConnectionDatabaseException):
+    pass
+
+
+class NodeConnectionNotSupportedException(NodeConnectionDatabaseException):
+    pass
+
+
+class NodeConnectionQueryCanceledException(NodeConnectionOperationalException):
+    pass
+
+
+class NodeConnectionTransactionRollbackException(NodeConnectionOperationalException):
+    pass
+
+
+class NodeConnectionArrayContentNotSupportedException(NodeConnectionNotSupportedException):
+    pass
+
+
 assert ExecUtilException.__name__ == "ExecUtilException"
 assert InvalidOperationException.__name__ == "InvalidOperationException"
 assert ExecTimeoutException.__name__ == "ExecTimeoutException"

testgres-1.17.0/src/impl/connection_pg8000.py

@@ -0,0 +1,197 @@
+# coding: utf-8
+# we support both pg8000 and psycopg2
+
+from ..enums import IsolationLevel
+
+from ..defaults import \
+    default_dbname, \
+    default_username2
+
+from ..exceptions import QueryException
+from ..exceptions import NodeConnectionException
+from ..exceptions import NodeConnectionInterfaceException
+from ..exceptions import NodeConnectionDatabaseException
+from ..exceptions import NodeConnectionDataException
+from ..exceptions import NodeConnectionOperationalException
+from ..exceptions import NodeConnectionIntegrityException
+from ..exceptions import NodeConnectionInternalException
+from ..exceptions import NodeConnectionProgrammingException
+from ..exceptions import NodeConnectionNotSupportedException
+from ..exceptions import NodeConnectionArrayContentNotSupportedException
+
+import typing
+
+import pg8000 as pglib
+import pg8000.dbapi as pglib_dbapi
+
+
+class NodeConnection:
+    """
+    Transaction wrapper returned by Node
+    """
+    def __init__(
+        self,
+        node,
+        dbname=None,
+        username=None,
+        password=None,
+        autocommit=False,
+    ):
+        # Set default arguments
+        dbname = dbname or default_dbname()
+        username = username or default_username2(node.os_ops)
+
+        self._node = node
+
+        c = __class__._exec_dblib(
+            pglib.connect,
+            database=dbname,
+            user=username,
+            password=password,
+            host=node.host,
+            port=node.port,
+        )
+        assert isinstance(c, pglib.Connection)
+
+        self._connection = c
+
+        self._connection.autocommit = autocommit
+
+        c = self.connection.cursor()
+        assert isinstance(c, pglib.Cursor)
+        self._cursor = c
+
+    @property
+    def node(self):
+        return self._node
+
+    @property
+    def connection(self):
+        return self._connection
+
+    @property
+    def pid(self):
+        r = self.execute("select pg_catalog.pg_backend_pid()")
+        assert r is not None
+        return r[0][0]
+
+    @property
+    def cursor(self):
+        return self._cursor
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+        return
+
+    def begin(self, isolation_level=IsolationLevel.ReadCommitted):
+        # Check if level isn't an IsolationLevel
+        if not isinstance(isolation_level, IsolationLevel):
+            # Get name of isolation level
+            level_str = str(isolation_level).lower()
+
+            # Validate level string
+            try:
+                isolation_level = IsolationLevel(level_str)
+            except ValueError:
+                error = 'Invalid isolation level "{}"'
+                raise QueryException(error.format(level_str))
+
+        # Set isolation level
+        cmd_f = 'SET TRANSACTION ISOLATION LEVEL {}'
+        cmd = cmd_f.format(isolation_level.value)
+
+        __class__._exec_dblib(
+            self.cursor.execute,
+            cmd,
+        )
+        return self
+
+    def commit(self):
+        __class__._exec_dblib(
+            self.connection.commit,
+        )
+        return self
+
+    def rollback(self):
+        __class__._exec_dblib(
+            self.connection.rollback,
+        )
+
+        return self
+
+    def execute(self, query, *args):
+        __class__._exec_dblib(
+            self.cursor.execute,
+            query,
+            args,
+        )
+
+        if self.cursor.description is None:
+            return None
+
+        # pg8000 might return tuples
+        r = __class__._exec_dblib(
+            self.cursor.fetchall,
+        )
+        assert type(r) is tuple
+        res = [tuple(t) for t in r]
+        return res
+
+    def close(self):
+        __class__._exec_dblib(
+            self.cursor.close,
+        )
+        __class__._exec_dblib(
+            self.connection.close,
+        )
+        return
+
+    @staticmethod
+    def _exec_dblib(
+        func: typing.Callable,
+        *args,
+        **kwargs,
+    ) -> typing.Any:
+        try:
+            r = func(*args, **kwargs)
+        except Exception as e:
+            __class__._map_pglib_exception(e)
+
+        return r
+
+    @staticmethod
+    def _map_pglib_exception(exc: Exception) -> typing.NoReturn:
+        if isinstance(exc, pglib_dbapi.ArrayContentNotSupportedError):
+            raise NodeConnectionArrayContentNotSupportedException(str(exc)) from exc
+
+        if isinstance(exc, pglib.NotSupportedError):
+            raise NodeConnectionNotSupportedException(str(exc)) from exc
+
+        if isinstance(exc, pglib.InternalError):
+            raise NodeConnectionInternalException(str(exc)) from exc
+
+        if isinstance(exc, pglib.ProgrammingError):
+            raise NodeConnectionProgrammingException(str(exc)) from exc
+
+        if isinstance(exc, pglib.IntegrityError):
+            raise NodeConnectionIntegrityException(str(exc)) from exc
+
+        if isinstance(exc, pglib.OperationalError):
+            raise NodeConnectionOperationalException(str(exc)) from exc
+
+        if isinstance(exc, pglib.DataError):
+            raise NodeConnectionDataException(str(exc)) from exc
+
+        if isinstance(exc, pglib.DatabaseError):
+            raise NodeConnectionDatabaseException(str(exc)) from exc
+
+        if isinstance(exc, pglib.InterfaceError):
+            raise NodeConnectionInterfaceException(str(exc)) from exc
+
+        if isinstance(exc, pglib.Error):
+            raise NodeConnectionException(str(exc)) from exc
+
+        raise exc

testgres-1.17.0/src/impl/connection_psycopg2.py

@@ -0,0 +1,203 @@
+# coding: utf-8
+# we support both pg8000 and psycopg2
+
+from ..enums import IsolationLevel
+
+from ..defaults import \
+    default_dbname, \
+    default_username2
+
+from ..exceptions import QueryException
+from ..exceptions import NodeConnectionException
+from ..exceptions import NodeConnectionInterfaceException
+from ..exceptions import NodeConnectionDatabaseException
+from ..exceptions import NodeConnectionDataException
+from ..exceptions import NodeConnectionOperationalException
+from ..exceptions import NodeConnectionIntegrityException
+from ..exceptions import NodeConnectionInternalException
+from ..exceptions import NodeConnectionProgrammingException
+from ..exceptions import NodeConnectionNotSupportedException
+from ..exceptions import NodeConnectionTransactionRollbackException
+from ..exceptions import NodeConnectionQueryCanceledException
+
+import typing
+
+import psycopg2 as pglib
+from psycopg2._psycopg import connection as PSYCOPG2__Connection
+from psycopg2._psycopg import cursor as PSYCOPG2__Cursor
+from psycopg2._psycopg import QueryCanceledError as PSYCOPG2__QueryCanceledError
+from psycopg2._psycopg import TransactionRollbackError as PSYCOPG2__TransactionRollbackError
+
+
+class NodeConnection:
+    """
+    Transaction wrapper returned by Node
+    """
+    def __init__(
+        self,
+        node,
+        dbname=None,
+        username=None,
+        password=None,
+        autocommit=False,
+    ):
+        # Set default arguments
+        dbname = dbname or default_dbname()
+        username = username or default_username2(node.os_ops)
+
+        self._node = node
+
+        c = __class__._exec_dblib(
+            pglib.connect,
+            database=dbname,
+            user=username,
+            password=password,
+            host=node.host,
+            port=node.port,
+        )
+        assert isinstance(c, PSYCOPG2__Connection)
+
+        self._connection = c
+
+        self._connection.autocommit = autocommit
+
+        c = self.connection.cursor()
+        assert isinstance(c, PSYCOPG2__Cursor)
+        self._cursor = c
+
+    @property
+    def node(self):
+        return self._node
+
+    @property
+    def connection(self):
+        return self._connection
+
+    @property
+    def pid(self):
+        r = self.execute("select pg_catalog.pg_backend_pid()")
+        assert r is not None
+        return r[0][0]
+
+    @property
+    def cursor(self):
+        return self._cursor
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+        return
+
+    def begin(self, isolation_level=IsolationLevel.ReadCommitted):
+        # Check if level isn't an IsolationLevel
+        if not isinstance(isolation_level, IsolationLevel):
+            # Get name of isolation level
+            level_str = str(isolation_level).lower()
+
+            # Validate level string
+            try:
+                isolation_level = IsolationLevel(level_str)
+            except ValueError:
+                error = 'Invalid isolation level "{}"'
+                raise QueryException(error.format(level_str))
+
+        # Set isolation level
+        cmd_f = 'SET TRANSACTION ISOLATION LEVEL {}'
+        cmd = cmd_f.format(isolation_level.value)
+
+        __class__._exec_dblib(
+            self.cursor.execute,
+            cmd,
+        )
+        return self
+
+    def commit(self):
+        __class__._exec_dblib(
+            self.connection.commit,
+        )
+        return self
+
+    def rollback(self):
+        __class__._exec_dblib(
+            self.connection.rollback,
+        )
+
+        return self
+
+    def execute(self, query, *args):
+        __class__._exec_dblib(
+            self.cursor.execute,
+            query,
+            args,
+        )
+
+        if self.cursor.description is None:
+            return None
+
+        # note: pg8000 might return tuples
+        res = __class__._exec_dblib(
+            self.cursor.fetchall,
+        )
+        assert type(res) is list
+        return res
+
+    def close(self):
+        __class__._exec_dblib(
+            self.cursor.close,
+        )
+        __class__._exec_dblib(
+            self.connection.close,
+        )
+        return
+
+    @staticmethod
+    def _exec_dblib(
+        func: typing.Callable,
+        *args,
+        **kwargs,
+    ) -> typing.Any:
+        try:
+            r = func(*args, **kwargs)
+        except Exception as e:
+            __class__._map_pglib_exception(e)
+
+        return r
+
+    @staticmethod
+    def _map_pglib_exception(exc: Exception) -> typing.NoReturn:
+        if isinstance(exc, PSYCOPG2__QueryCanceledError):
+            raise NodeConnectionQueryCanceledException(str(exc)) from exc
+
+        if isinstance(exc, PSYCOPG2__TransactionRollbackError):
+            raise NodeConnectionTransactionRollbackException(str(exc)) from exc
+
+        if isinstance(exc, pglib.NotSupportedError):
+            raise NodeConnectionNotSupportedException(str(exc)) from exc
+
+        if isinstance(exc, pglib.InternalError):
+            raise NodeConnectionInternalException(str(exc)) from exc
+
+        if isinstance(exc, pglib.ProgrammingError):
+            raise NodeConnectionProgrammingException(str(exc)) from exc
+
+        if isinstance(exc, pglib.IntegrityError):
+            raise NodeConnectionIntegrityException(str(exc)) from exc
+
+        if isinstance(exc, pglib.OperationalError):
+            raise NodeConnectionOperationalException(str(exc)) from exc
+
+        if isinstance(exc, pglib.DataError):
+            raise NodeConnectionDataException(str(exc)) from exc
+
+        if isinstance(exc, pglib.DatabaseError):
+            raise NodeConnectionDatabaseException(str(exc)) from exc
+
+        if isinstance(exc, pglib.InterfaceError):
+            raise NodeConnectionInterfaceException(str(exc)) from exc
+
+        if isinstance(exc, pglib.Error):
+            raise NodeConnectionException(str(exc)) from exc
+
+        raise exc