[BUG FIX] A problem in psql/safe_psql and 'input' data was fixed [local_op]

dmitry-lipetsk · dmitry-lipetsk · commit 45b8dc024dc3 · 2024-12-09T15:08:00.000+03:00
Both LocalOperations::exec_command and RemoteOperations::exec_command were updated.
diff --git a/testgres/node.py b/testgres/node.py
@@ -1005,6 +1005,18 @@ def _psql(
             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 = [
diff --git a/testgres/operations/helpers.py b/testgres/operations/helpers.py
@@ -0,0 +1,15 @@
+class Helpers:
+    def PrepareProcessInput(input, encoding):
+        if not input:
+            return None
+
+        if type(input) == str:  # noqa: E721
+            if encoding is None:
+                return input.encode()
+
+            assert type(encoding) == str  # noqa: E721
+            return input.encode(encoding)
+
+        # It is expected!
+        assert type(input) == bytes  # noqa: E721
+        return input
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -11,6 +11,7 @@
 
 from ..exceptions import ExecUtilException
 from .os_ops import ConnectionParams, OsOperations, pglib, get_default_encoding
+from .helpers import Helpers
 from ..helpers.raise_error import RaiseError
 
 try:
@@ -58,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
@@ -79,6 +82,10 @@ 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
+
         process = subprocess.Popen(
             cmd,
             shell=shell,
@@ -89,7 +96,7 @@ def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_pr
         if get_process:
             return process, None, None
         try:
-            output, error = process.communicate(input=input.encode(encoding) if input else None, timeout=timeout)
+            output, error = process.communicate(input=input_prepared, timeout=timeout)
             if encoding:
                 output = output.decode(encoding)
                 error = error.decode(encoding)
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -16,6 +16,7 @@
 from ..exceptions import ExecUtilException
 from ..helpers.raise_error import RaiseError
 from .os_ops import OsOperations, ConnectionParams, get_default_encoding
+from .helpers import Helpers
 
 error_markers = [b'error', b'Permission denied', b'fatal', b'No such file or directory']
 
@@ -70,6 +71,10 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
         assert type(expect_error) == bool  # noqa: E721
         assert type(ignore_errors) == bool  # noqa: E721
 
+        input_prepared = None
+        if not get_process:
+            input_prepared = Helpers.PrepareProcessInput(input, encoding)  # throw
+
         ssh_cmd = []
         if isinstance(cmd, str):
             ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_args + [cmd]
@@ -80,7 +85,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
             return process
 
         try:
-            result, error = process.communicate(input, timeout=timeout)
+            result, error = process.communicate(input=input_prepared, timeout=timeout)
         except subprocess.TimeoutExpired:
             process.kill()
             raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
