PBCKP-152 fix failed tests

v.shepard · v.shepard · commit f6bfe17af893 · 2023-06-13T20:17:43.000+02:00
diff --git a/testgres/__init__.py b/testgres/__init__.py
@@ -46,6 +46,11 @@
     First, \
     Any
 
+import operations
+from .operations.os_ops import OsOperations
+from .operations.local_ops import LocalOperations
+from .operations.remote_ops import RemoteOperations
+
 __all__ = [
     "get_new_node",
     "NodeBackup",
@@ -56,4 +61,5 @@
     "PostgresNode", "NodeApp",
     "reserve_port", "release_port", "bound_ports", "get_bin_path", "get_pg_config", "get_pg_version",
     "First", "Any",
+    "OsOperations", "LocalOperations", "RemoteOperations", "operations"
 ]
diff --git a/testgres/cache.py b/testgres/cache.py
@@ -4,8 +4,6 @@
 
 from six import raise_from
 
-from .os_ops.local_ops import LocalOperations
-from .os_ops.os_ops import OsOperations
 from .config import testgres_config
 
 from .consts import XLOG_CONTROL_FILE
@@ -20,6 +18,9 @@
     get_bin_path, \
     execute_utility
 
+from .operations.local_ops import LocalOperations
+from .operations.os_ops import OsOperations
+
 
 def cached_initdb(data_dir, logfile=None, params=None, os_ops: OsOperations = LocalOperations()):
     """
@@ -38,7 +39,7 @@ def call_initdb(initdb_dir, log=None):
         call_initdb(data_dir, logfile)
     else:
         # Fetch cached initdb dir
-        cached_data_dir = testgres_config.cached_initdb_dir()
+        cached_data_dir = testgres_config.cached_initdb_dir
 
         # Initialize cached initdb
 
diff --git a/testgres/config.py b/testgres/config.py
@@ -6,8 +6,8 @@
 
 from contextlib import contextmanager
 
-from .os_ops.local_ops import LocalOperations
 from .consts import TMP_CACHE
+from .operations.local_ops import LocalOperations
 
 
 class GlobalConfig(object):
diff --git a/testgres/connection.py b/testgres/connection.py
@@ -102,23 +102,15 @@ def rollback(self):
         return self
 
     def execute(self, query, *args):
+        self.cursor.execute(query, args)
         try:
-            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
+            res = self.cursor.fetchall()
+            # pg8000 might return tuples
+            if isinstance(res, tuple):
+                res = [tuple(t) for t in res]
+
+            return res
+        except Exception as e:
             print(f"Error executing query: {e}")
             return None
 
diff --git a/testgres/defaults.py b/testgres/defaults.py
@@ -2,7 +2,7 @@
 import struct
 import uuid
 
-from .os_ops.local_ops import LocalOperations
+from .operations.local_ops import LocalOperations
 
 
 def default_dbname():
diff --git a/testgres/node.py b/testgres/node.py
@@ -9,9 +9,6 @@
 import psutil
 import time
 
-from .os_ops.local_ops import LocalOperations
-from .os_ops.remote_ops import RemoteOperations
-
 try:
     from collections.abc import Iterable
 except ImportError:
@@ -99,6 +96,9 @@
 
 from .backup import NodeBackup
 
+from .operations.local_ops import LocalOperations
+from .operations.remote_ops import RemoteOperations
+
 InternalError = pglib.InternalError
 ProgrammingError = pglib.ProgrammingError
 OperationalError = pglib.OperationalError
@@ -201,7 +201,7 @@ def pid(self):
 
         if self.status():
             pid_file = os.path.join(self.data_dir, PG_PID_FILE)
-            lines = self.os_ops.readlines(pid_file, num_lines=1)
+            lines = self.os_ops.readlines(pid_file)
             pid = int(lines[0]) if lines else None
             return pid
 
@@ -433,7 +433,8 @@ def _collect_special_files(self):
             if not self.os_ops.path_exists(f):
                 continue
 
-            lines = b''.join(self.os_ops.readlines(f, num_lines, encoding='utf-8'))
+            file_lines = self.os_ops.readlines(f, num_lines, binary=True, encoding=None)
+            lines = b''.join(file_lines)
 
             # fill list
             result.append((f, lines))
@@ -498,7 +499,7 @@ def default_conf(self,
         ]
 
         # write filtered lines
-        self.os_ops.write(hba_conf_file, lines, truncate=True)
+        self.os_ops.write(hba_conf, lines, truncate=True)
 
         # replication-related settings
         if allow_streaming:
@@ -960,11 +961,9 @@ def psql(self,
         psql_params.append(dbname)
 
         # start psql process
-        process = self.os_ops.exec_command(psql_params)
+        status_code, out, err = self.os_ops.exec_command(psql_params, shell=False, verbose=True, input=input)
 
-        # wait until it finishes and get stdout and stderr
-        out, err = process.communicate(input=input)
-        return process.returncode, out, err
+        return status_code, out, err
 
     @method_decorator(positional_args_hack(['dbname', 'query']))
     def safe_psql(self, query=None, expect_error=False, **kwargs):
@@ -1348,7 +1347,7 @@ def pgbench(self,
         # should be the last one
         _params.append(dbname)
 
-        proc = self.os_ops.exec_command(_params, wait_exit=True)
+        proc = self.os_ops.exec_command(_params, stdout=stdout, stderr=stderr, wait_exit=True, shell=False, proc=True)
 
         return proc
 
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -19,18 +19,25 @@ def __init__(self, username=None):
         self.username = username or self.get_user()
 
     # Command execution
-    def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False):
-        if isinstance(cmd, list):
-            cmd = " ".join(cmd)
+    def exec_command(self, cmd, wait_exit=False, verbose=False,
+                     expect_error=False, encoding=None, shell=True, text=False,
+                     input=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, proc=None):
         log.debug(f"os_ops.exec_command: `{cmd}`; remote={self.remote}")
         # Source global profile file + execute command
         try:
+            if proc:
+                return subprocess.Popen(cmd,
+                                        shell=shell,
+                                        stdin=input or subprocess.PIPE,
+                                        stdout=stdout,
+                                        stderr=stderr)
             process = subprocess.run(
                 cmd,
-                shell=True,
-                text=True,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
+                input=input,
+                shell=shell,
+                text=text,
+                stdout=stdout,
+                stderr=stderr,
                 timeout=CMD_TIMEOUT_SEC,
             )
             exit_status = process.returncode
@@ -39,11 +46,11 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False):
 
             if expect_error:
                 raise Exception(result, error)
-            if exit_status != 0 or "error" in error.lower():
+            if exit_status != 0 or "error" in error.lower().decode(encoding or 'utf-8'):  # Decode error for comparison
                 log.error(
-                    f"Problem in executing command: `{cmd}`\nerror: {error}\nexit_code: {exit_status}"
+                    f"Problem in executing command: `{cmd}`\nerror: {error.decode(encoding or 'utf-8')}\nexit_code: {exit_status}"
+                    # Decode for logging
                 )
-                exit(1)
 
             if verbose:
                 return exit_status, result, error
@@ -152,9 +159,9 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
         """
         mode = "wb" if binary else "w"
         if not truncate:
-            mode = "a" + mode
+            mode = "ab" if binary else "a"
         if read_and_write:
-            mode = "r+" + mode
+            mode = "r+b" if binary else "r+"
 
         with open(filename, mode) as file:
             if isinstance(data, list):
@@ -174,26 +181,26 @@ def touch(self, filename):
         with open(filename, "a"):
             os.utime(filename, None)
 
-    def read(self, filename):
-        with open(filename, "r") as file:
+    def read(self, filename, encoding=None):
+        with open(filename, "r", encoding=encoding) as file:
             return file.read()
 
-    def readlines(self, filename, num_lines=0, encoding=None):
+    def readlines(self, filename, num_lines=0, binary=False, encoding=None):
         """
         Read lines from a local file.
         If num_lines is greater than 0, only the last num_lines lines will be read.
         """
         assert num_lines >= 0
-
+        mode = 'rb' if binary else 'r'
         if num_lines == 0:
-            with open(filename, "r", encoding=encoding) as file:
+            with open(filename, mode, encoding=encoding) as file:  # open in binary mode
                 return file.readlines()
 
         else:
             bufsize = 8192
             buffers = 1
 
-            with open(filename, "r", encoding=encoding) as file:
+            with open(filename, mode, encoding=encoding) as file:  # open in binary mode
                 file.seek(0, os.SEEK_END)
                 end_pos = file.tell()
 
@@ -205,7 +212,7 @@ def readlines(self, filename, num_lines=0, encoding=None):
                     cur_lines = len(lines)
 
                     if cur_lines >= num_lines or pos == 0:
-                        return lines[-num_lines:]
+                        return lines[-num_lines:]  # get last num_lines from lines
 
                     buffers = int(
                         buffers * max(2, int(num_lines / max(cur_lines, 1)))
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -1,3 +1,4 @@
+import io
 import os
 import tempfile
 from contextlib import contextmanager
@@ -65,30 +66,41 @@ def connect(self):
             return ssh
 
     # Command execution
-    def exec_command(
-        self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding="utf-8"
-    ):
+    def exec_command(self, cmd, wait_exit=False, verbose=False,
+                     expect_error=False, encoding=None, shell=True, text=False,
+                     input=None, stdout=None, stderr=None, proc=None):
         if isinstance(cmd, list):
             cmd = " ".join(cmd)
         log.debug(f"os_ops.exec_command: `{cmd}`; remote={self.remote}")
         # Source global profile file + execute command
         try:
             cmd = f"source /etc/profile.d/custom.sh; {cmd}"
             with self.ssh_connect() as ssh:
-                stdin, stdout, stderr = ssh.exec_command(cmd)
+                if input:
+                    # encode input and feed it to stdin
+                    stdin, stdout, stderr = ssh.exec_command(cmd)
+                    stdin.write(input)
+                    stdin.flush()
+                else:
+                    stdin, stdout, stderr = ssh.exec_command(cmd)
                 exit_status = 0
                 if wait_exit:
                     exit_status = stdout.channel.recv_exit_status()
-                result = stdout.read().decode(encoding)
-                error = stderr.read().decode(encoding)
+                if encoding:
+                    result = stdout.read().decode(encoding)
+                    error = stderr.read().decode(encoding)
+                else:
+                    # Save as binary string
+                    result = io.BytesIO(stdout.read()).getvalue()
+                    error = io.BytesIO(stderr.read()).getvalue()
+                error_str = stderr.read()
 
             if expect_error:
                 raise Exception(result, error)
-            if exit_status != 0 or "error" in error.lower():
+            if exit_status != 0 or 'error' in error_str:
                 log.error(
                     f"Problem in executing command: `{cmd}`\nerror: {error}\nexit_code: {exit_status}"
                 )
-                exit(1)
 
             if verbose:
                 return exit_status, result, error
@@ -203,9 +215,9 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
         """
         mode = "wb" if binary else "w"
         if not truncate:
-            mode = "a" + mode
+            mode = "ab" if binary else "a"
         if read_and_write:
-            mode = "r+" + mode
+            mode = "r+b" if binary else "r+"
 
         with tempfile.NamedTemporaryFile(mode=mode) as tmp_file:
             if isinstance(data, list):
@@ -229,17 +241,28 @@ def touch(self, filename):
         """
         self.exec_command(f"touch {filename}")
 
-    def read(self, filename, encoding="utf-8"):
+    def read(self, filename, binary=False, encoding=None):
         cmd = f"cat {filename}"
-        return self.exec_command(cmd, encoding=encoding)
+        result = self.exec_command(cmd, encoding=encoding)
+
+        if not binary and result:
+            result = result.decode(encoding or 'utf-8')
+
+        return result
 
-    def readlines(self, filename, num_lines=0, encoding=None):
-        encoding = encoding or "utf-8"
+    def readlines(self, filename, num_lines=0, binary=False, encoding=None):
         if num_lines > 0:
             cmd = f"tail -n {num_lines} {filename}"
-            lines = self.exec_command(cmd, encoding)
         else:
-            lines = self.read(filename, encoding=encoding).splitlines()
+            cmd = f"cat {filename}"
+
+        result = self.exec_command(cmd, encoding=encoding)
+
+        if not binary and result:
+            lines = result.decode(encoding or 'utf-8').splitlines()
+        else:
+            lines = result.splitlines()
+
         return lines
 
     def isfile(self, remote_file):
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -13,17 +13,17 @@
 from contextlib import contextmanager
 from packaging.version import Version
 
-from .os_ops.remote_ops import RemoteOperations
-
 try:
     from shutil import which as find_executable
 except ImportError:
     from distutils.spawn import find_executable
 from six import iteritems
 
 from fabric import Connection
-from .os_ops.local_ops import LocalOperations
-from .os_ops.os_ops import OsOperations
+
+from .operations.remote_ops import RemoteOperations
+from .operations.local_ops import LocalOperations
+from .operations.os_ops import OsOperations
 
 from .config import testgres_config
 from .exceptions import ExecUtilException
diff --git a/tests/test_remote.py b/tests/test_remote.py
diff --git a/tests/test_simple.py b/tests/test_simple.py
