Merge pull request #161 from dmitry-lipetsk/master-fix154--v02

Proposal to fix #154 (v2)

Author: dmitry-lipetsk

testgres/__init__.py

@@ -23,7 +23,8 @@
     CatchUpException, \
     StartNodeException, \
     InitNodeException, \
-    BackupException
+    BackupException, \
+    InvalidOperationException
 
 from .enums import \
     XLogMethod, \
@@ -60,7 +61,7 @@
     "NodeBackup", "testgres_config",
     "TestgresConfig", "configure_testgres", "scoped_config", "push_config", "pop_config",
     "NodeConnection", "DatabaseError", "InternalError", "ProgrammingError", "OperationalError",
-    "TestgresException", "ExecUtilException", "QueryException", "TimeoutException", "CatchUpException", "StartNodeException", "InitNodeException", "BackupException",
+    "TestgresException", "ExecUtilException", "QueryException", "TimeoutException", "CatchUpException", "StartNodeException", "InitNodeException", "BackupException", "InvalidOperationException",
     "XLogMethod", "IsolationLevel", "NodeStatus", "ProcessType", "DumpFormat",
     "PostgresNode", "NodeApp",
     "reserve_port", "release_port", "bound_ports", "get_bin_path", "get_pg_config", "get_pg_version",

testgres/exceptions.py

@@ -9,13 +9,14 @@ class TestgresException(Exception):
 
 @six.python_2_unicode_compatible
 class ExecUtilException(TestgresException):
-    def __init__(self, message=None, command=None, exit_code=0, out=None):
+    def __init__(self, message=None, command=None, exit_code=0, out=None, error=None):
         super(ExecUtilException, self).__init__(message)
 
         self.message = message
         self.command = command
         self.exit_code = exit_code
         self.out = out
+        self.error = error
 
     def __str__(self):
         msg = []
@@ -24,13 +25,17 @@ def __str__(self):
             msg.append(self.message)
 
         if self.command:
-            msg.append(u'Command: {}'.format(self.command))
+            command_s = ' '.join(self.command) if isinstance(self.command, list) else self.command,
+            msg.append(u'Command: {}'.format(command_s))
 
         if self.exit_code:
             msg.append(u'Exit code: {}'.format(self.exit_code))
 
+        if self.error:
+            msg.append(u'---- Error:\n{}'.format(self.error))
+
         if self.out:
-            msg.append(u'----\n{}'.format(self.out))
+            msg.append(u'---- Out:\n{}'.format(self.out))
 
         return self.convert_and_join(msg)
 
@@ -98,3 +103,7 @@ class InitNodeException(TestgresException):
 
 class BackupException(TestgresException):
     pass
+
+
+class InvalidOperationException(TestgresException):
+    pass

testgres/node.py

@@ -74,7 +74,8 @@
     TimeoutException,   \
     InitNodeException,  \
     TestgresException,  \
-    BackupException
+    BackupException,    \
+    InvalidOperationException
 
 from .logger import TestgresLogger
 
@@ -987,6 +988,37 @@ def psql(self,
             >>> psql(query='select 3', ON_ERROR_STOP=1)
         """
 
+        return self._psql(
+            ignore_errors=True,
+            query=query,
+            filename=filename,
+            dbname=dbname,
+            username=username,
+            input=input,
+            **variables
+        )
+
+    def _psql(
+            self,
+            ignore_errors,
+            query=None,
+            filename=None,
+            dbname=None,
+            username=None,
+            input=None,
+            **variables):
+        assert type(variables) == dict  # noqa: E721
+
+        #
+        # We do not support encoding. It may be added later. Ok?
+        #
+        if input is None:
+            pass
+        elif type(input) == bytes:  # noqa: E721
+            pass
+        else:
+            raise Exception("Input data must be None or bytes.")
+
         dbname = dbname or default_dbname()
 
         psql_params = [
@@ -1017,20 +1049,14 @@ def psql(self,
 
         # should be the last one
         psql_params.append(dbname)
-        if not self.os_ops.remote:
-            # start psql process
-            process = subprocess.Popen(psql_params,
-                                       stdin=subprocess.PIPE,
-                                       stdout=subprocess.PIPE,
-                                       stderr=subprocess.PIPE)
-
-            # wait until it finishes and get stdout and stderr
-            out, err = process.communicate(input=input)
-            return process.returncode, out, err
-        else:
-            status_code, out, err = self.os_ops.exec_command(psql_params, verbose=True, input=input)
 
-            return status_code, out, err
+        return self.os_ops.exec_command(
+            psql_params,
+            verbose=True,
+            input=input,
+            stderr=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            ignore_errors=ignore_errors)
 
     @method_decorator(positional_args_hack(['dbname', 'query']))
     def safe_psql(self, query=None, expect_error=False, **kwargs):
@@ -1051,22 +1077,27 @@ def safe_psql(self, query=None, expect_error=False, **kwargs):
         Returns:
             psql's output as str.
         """
+        assert type(kwargs) == dict  # noqa: E721
+        assert not ("ignore_errors" in kwargs.keys())
+        assert not ("expect_error" in kwargs.keys())
 
         # force this setting
         kwargs['ON_ERROR_STOP'] = 1
         try:
-            ret, out, err = self.psql(query=query, **kwargs)
+            ret, out, err = self._psql(ignore_errors=False, query=query, **kwargs)
         except ExecUtilException as e:
-            ret = e.exit_code
-            out = e.out
-            err = e.message
-        if ret:
-            if expect_error:
-                out = (err or b'').decode('utf-8')
-            else:
-                raise QueryException((err or b'').decode('utf-8'), query)
-        elif expect_error:
-            assert False, "Exception was expected, but query finished successfully: `{}` ".format(query)
+            if not expect_error:
+                raise QueryException(e.message, query)
+
+            if type(e.error) == bytes:  # noqa: E721
+                return e.error.decode("utf-8")  # throw
+
+            # [2024-12-09] This situation is not expected
+            assert False
+            return e.error
+
+        if expect_error:
+            raise InvalidOperationException("Exception was expected, but query finished successfully: `{}`.".format(query))
 
         return out
 

testgres/operations/helpers.py

@@ -0,0 +1,52 @@
+import locale
+
+
+class Helpers:
+    def _make_get_default_encoding_func():
+        # locale.getencoding is added in Python 3.11
+        if hasattr(locale, 'getencoding'):
+            return locale.getencoding
+
+        # It must exist
+        return locale.getpreferredencoding
+
+    # Prepared pointer on function to get a name of system codepage
+    _get_default_encoding_func = _make_get_default_encoding_func()
+
+    def GetDefaultEncoding():
+        #
+        #   Original idea/source was:
+        #
+        #   def os_ops.get_default_encoding():
+        #       if not hasattr(locale, 'getencoding'):
+        #       locale.getencoding = locale.getpreferredencoding
+        #       return locale.getencoding() or 'UTF-8'
+        #
+
+        assert __class__._get_default_encoding_func is not None
+
+        r = __class__._get_default_encoding_func()
+
+        if r:
+            assert r is not None
+            assert type(r) == str  # noqa: E721
+            assert r != ""
+            return r
+
+        # Is it an unexpected situation?
+        return 'UTF-8'
+
+    def PrepareProcessInput(input, encoding):
+        if not input:
+            return None
+
+        if type(input) == str:  # noqa: E721
+            if encoding is None:
+                return input.encode(__class__.GetDefaultEncoding())
+
+            assert type(encoding) == str  # noqa: E721
+            return input.encode(encoding)
+
+        # It is expected!
+        assert type(input) == bytes  # noqa: E721
+        return input

testgres/operations/local_ops.py

@@ -11,6 +11,8 @@
 
 from ..exceptions import ExecUtilException
 from .os_ops import ConnectionParams, OsOperations, pglib, get_default_encoding
+from .raise_error import RaiseError
+from .helpers import Helpers
 
 try:
     from shutil import which as find_executable
@@ -47,14 +49,6 @@ def __init__(self, conn_params=None):
         self.remote = False
         self.username = conn_params.username or getpass.getuser()
 
-    @staticmethod
-    def _raise_exec_exception(message, command, exit_code, output):
-        """Raise an ExecUtilException."""
-        raise ExecUtilException(message=message.format(output),
-                                command=' '.join(command) if isinstance(command, list) else command,
-                                exit_code=exit_code,
-                                out=output)
-
     @staticmethod
     def _process_output(encoding, temp_file_path):
         """Process the output of a command from a temporary file."""
@@ -65,6 +59,8 @@ def _process_output(encoding, temp_file_path):
             return output, None  # In Windows stderr writing in stdout
 
     def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
+        # TODO: why don't we use the data from input?
+
         with tempfile.NamedTemporaryFile(mode='w+b', delete=False) as temp_file:
             stdout = temp_file
             stderr = subprocess.STDOUT
@@ -86,25 +82,36 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process
         return process, output, error
 
     def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
+        input_prepared = None
+        if not get_process:
+            input_prepared = Helpers.PrepareProcessInput(input, encoding)  # throw
+
+        assert input_prepared is None or (type(input_prepared) == bytes)  # noqa: E721
+
         process = subprocess.Popen(
             cmd,
             shell=shell,
             stdin=stdin or subprocess.PIPE if input is not None else None,
             stdout=stdout or subprocess.PIPE,
             stderr=stderr or subprocess.PIPE,
         )
+        assert not (process is None)
         if get_process:
             return process, None, None
         try:
-            output, error = process.communicate(input=input.encode(encoding) if input else None, timeout=timeout)
-            if encoding:
-                output = output.decode(encoding)
-                error = error.decode(encoding)
-            return process, output, error
+            output, error = process.communicate(input=input_prepared, timeout=timeout)
         except subprocess.TimeoutExpired:
             process.kill()
             raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
 
+        assert type(output) == bytes  # noqa: E721
+        assert type(error) == bytes  # noqa: E721
+
+        if encoding:
+            output = output.decode(encoding)
+            error = error.decode(encoding)
+        return process, output, error
+
     def _run_command(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
         """Execute a command and return the process and its output."""
         if os.name == 'nt' and stdout is None:  # Windows
@@ -120,11 +127,20 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
         """
         Execute a command in a subprocess and handle the output based on the provided parameters.
         """
+        assert type(expect_error) == bool  # noqa: E721
+        assert type(ignore_errors) == bool  # noqa: E721
+
         process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
         if get_process:
             return process
         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)
+            RaiseError.UtilityExitedWithNonZeroCode(
+                cmd=cmd,
+                exit_code=process.returncode,
+                msg_arg=error or output,
+                error=error,
+                out=output
+            )
 
         if verbose:
             return process.returncode, output, error

testgres/operations/raise_error.py

@@ -0,0 +1,46 @@
+from ..exceptions import ExecUtilException
+from .helpers import Helpers
+
+
+class RaiseError:
+    def UtilityExitedWithNonZeroCode(cmd, exit_code, msg_arg, error, out):
+        assert type(exit_code) == int  # noqa: E721
+
+        msg_arg_s = __class__._TranslateDataIntoString(msg_arg).strip()
+        assert type(msg_arg_s) == str  # noqa: E721
+
+        if msg_arg_s == "":
+            msg_arg_s = "#no_error_message"
+
+        message = "Utility exited with non-zero code. Error: `" + msg_arg_s + "`"
+        raise ExecUtilException(
+            message=message,
+            command=cmd,
+            exit_code=exit_code,
+            out=out,
+            error=error)
+
+    def _TranslateDataIntoString(data):
+        if type(data) == bytes:  # noqa: E721
+            return __class__._TranslateDataIntoString__FromBinary(data)
+
+        return str(data)
+
+    def _TranslateDataIntoString__FromBinary(data):
+        assert type(data) == bytes  # noqa: E721
+
+        try:
+            return data.decode(Helpers.GetDefaultEncoding())
+        except UnicodeDecodeError:
+            pass
+
+        return "#cannot_decode_text"
+
+    def _BinaryIsASCII(data):
+        assert type(data) == bytes  # noqa: E721
+
+        for b in data:
+            if not (b >= 0 and b <= 127):
+                return False
+
+        return True