Add ability to skip ssl when connect to PostgresNode

vshepard · vshepard · commit 1b10a5499b26 · 2024-11-18T12:38:37.000+01:00
diff --git a/testgres/api.py b/testgres/api.py
@@ -42,7 +42,7 @@ def get_new_node(name=None, base_dir=None, **kwargs):
     return PostgresNode(name=name, base_dir=base_dir, **kwargs)
 
 
-def get_remote_node(name=None, conn_params=None):
+def get_remote_node(name=None):
     """
     Simply a wrapper around :class:`.PostgresNode` constructor for remote node.
     See :meth:`.PostgresNode.__init__` for details.
@@ -51,4 +51,4 @@ def get_remote_node(name=None, conn_params=None):
                                    ssh_key=None,
                                    username=default_username())
     """
-    return get_new_node(name=name, conn_params=conn_params)
+    return get_new_node(name=name)
diff --git a/testgres/node.py b/testgres/node.py
@@ -126,7 +126,8 @@ def __repr__(self):
 
 
 class PostgresNode(object):
-    def __init__(self, name=None, base_dir=None, port=None, conn_params: ConnectionParams = ConnectionParams(), bin_dir=None, prefix=None):
+    def __init__(self, name=None, base_dir=None, port=None, conn_params: ConnectionParams = ConnectionParams(),
+                 bin_dir=None, prefix=None):
         """
         PostgresNode constructor.
 
@@ -150,13 +151,9 @@ def __init__(self, name=None, base_dir=None, port=None, conn_params: ConnectionP
         self.name = name or generate_app_name()
         if testgres_config.os_ops:
             self.os_ops = testgres_config.os_ops
-        elif conn_params.ssh_key:
-            self.os_ops = RemoteOperations(conn_params)
-        else:
-            self.os_ops = LocalOperations(conn_params)
 
         self.host = self.os_ops.host
-        self.port = port or reserve_port()
+        self.port = port or self.os_ops.port or reserve_port()
 
         self.ssh_key = self.os_ops.ssh_key
 
@@ -1005,7 +1002,7 @@ def psql(self,
 
         # select query source
         if query:
-            if self.os_ops.remote:
+            if self.os_ops.conn_params.remote:
                 psql_params.extend(("-c", '"{}"'.format(query)))
             else:
                 psql_params.extend(("-c", query))
@@ -1016,7 +1013,7 @@ def psql(self,
 
         # should be the last one
         psql_params.append(dbname)
-        if not self.os_ops.remote:
+        if not self.os_ops.conn_params.remote:
             # start psql process
             process = subprocess.Popen(psql_params,
                                        stdin=subprocess.PIPE,
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -40,12 +40,7 @@ class LocalOperations(OsOperations):
     def __init__(self, conn_params=None):
         if conn_params is None:
             conn_params = ConnectionParams()
-        super(LocalOperations, self).__init__(conn_params.username)
-        self.conn_params = conn_params
-        self.host = conn_params.host
-        self.ssh_key = None
-        self.remote = False
-        self.username = conn_params.username or getpass.getuser()
+        super(LocalOperations, self).__init__(conn_params)
 
     @staticmethod
     def _raise_exec_exception(message, command, exit_code, output):
@@ -305,14 +300,3 @@ def get_pid(self):
 
     def get_process_children(self, pid):
         return psutil.Process(pid).children()
-
-    # Database control
-    def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
-        conn = pglib.connect(
-            host=host,
-            port=port,
-            database=dbname,
-            user=user,
-            password=password,
-        )
-        return conn
diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py
@@ -12,11 +12,16 @@
 
 
 class ConnectionParams:
-    def __init__(self, host='127.0.0.1', port=None, ssh_key=None, username=None):
+    def __init__(self, host='127.0.0.1', port=None, ssh_key=None, username=None, remote=False, skip_ssl=False):
+        """
+            skip_ssl: if is True, the connection is established without SSL.
+        """
+        self.remote = remote
         self.host = host
         self.port = port
         self.ssh_key = ssh_key
         self.username = username
+        self.skip_ssl = skip_ssl
 
 
 def get_default_encoding():
@@ -26,9 +31,12 @@ def get_default_encoding():
 
 
 class OsOperations:
-    def __init__(self, username=None):
-        self.ssh_key = None
-        self.username = username or getpass.getuser()
+    def __init__(self, conn_params=ConnectionParams()):
+        self.ssh_key = conn_params.ssh_key
+        self.username = conn_params.username or getpass.getuser()
+        self.host = conn_params.host
+        self.port = conn_params.port
+        self.conn_params = conn_params
 
     # Command execution
     def exec_command(self, cmd, **kwargs):
@@ -115,4 +123,14 @@ def get_process_children(self, pid):
 
     # Database control
     def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
-        raise NotImplementedError()
+        ssl_options = {"sslmode": "disable"} if self.conn_params.skip_ssl and 'psycopg2' in globals() else {}
+        conn = pglib.connect(
+            host=host,
+            port=port,
+            database=dbname,
+            user=user,
+            password=password,
+            **({"ssl_context": None} if self.conn_params.skip_ssl and 'pg8000' in globals() else ssl_options)
+        )
+
+        return conn
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -37,23 +37,17 @@ def cmdline(self):
 
 class RemoteOperations(OsOperations):
     def __init__(self, conn_params: ConnectionParams):
-
         if not platform.system().lower() == "linux":
             raise EnvironmentError("Remote operations are supported only on Linux!")
+        super().__init__(conn_params)
 
-        super().__init__(conn_params.username)
-        self.conn_params = conn_params
-        self.host = conn_params.host
-        self.port = conn_params.port
-        self.ssh_key = conn_params.ssh_key
         self.ssh_args = []
         if self.ssh_key:
             self.ssh_args += ["-i", self.ssh_key]
         if self.port:
             self.ssh_args += ["-p", self.port]
-        self.remote = True
         self.username = conn_params.username or getpass.getuser()
-        self.ssh_dest = f"{self.username}@{self.host}" if conn_params.username else self.host
+        self.ssh_dest = f"{self.username}@{self.host}" if self.username else self.host
 
     def __enter__(self):
         return self
@@ -361,17 +355,6 @@ def get_process_children(self, pid):
         else:
             raise ExecUtilException(f"Error in getting process children. Error: {result.stderr}")
 
-    # Database control
-    def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
-        conn = pglib.connect(
-            host=host,
-            port=port,
-            database=dbname,
-            user=user,
-            password=password,
-        )
-        return conn
-
 
 def normalize_error(error):
     if isinstance(error, bytes):
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -97,7 +97,7 @@ def get_bin_path(filename):
     # check if it's already absolute
     if os.path.isabs(filename):
         return filename
-    if tconf.os_ops.remote:
+    if tconf.os_ops.conn_params.remote:
         pg_config = os.environ.get("PG_CONFIG_REMOTE") or os.environ.get("PG_CONFIG")
     else:
         # try PG_CONFIG - get from local machine
@@ -154,7 +154,7 @@ def cache_pg_config_data(cmd):
         return _pg_config_data
 
     # try specified pg_config path or PG_CONFIG
-    if tconf.os_ops.remote:
+    if tconf.os_ops.conn_params.remote:
         pg_config = pg_config_path or os.environ.get("PG_CONFIG_REMOTE") or os.environ.get("PG_CONFIG")
     else:
         # try PG_CONFIG - get from local machine
diff --git a/tests/test_remote.py b/tests/test_remote.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from testgres import ExecUtilException
+from testgres import ExecUtilException, get_remote_node, testgres_config
 from testgres import RemoteOperations
 from testgres import ConnectionParams
 
@@ -34,7 +34,7 @@ def test_exec_command_failure(self):
             exit_status, result, error = self.operations.exec_command(cmd, verbose=True, wait_exit=True)
         except ExecUtilException as e:
             error = e.message
-        assert error == b'Utility exited with non-zero code. Error: bash: line 1: nonexistent_command: command not found\n'
+        assert error == 'Utility exited with non-zero code. Error: bash: line 1: nonexistent_command: command not found\n'
 
     def test_is_executable_true(self):
         """
@@ -87,7 +87,7 @@ def test_makedirs_and_rmdirs_failure(self):
             exit_status, result, error = self.operations.rmdirs(path, verbose=True)
         except ExecUtilException as e:
             error = e.message
-        assert error == b"Utility exited with non-zero code. Error: rm: cannot remove '/root/test_dir': Permission denied\n"
+        assert error == "Utility exited with non-zero code. Error: rm: cannot remove '/root/test_dir': Permission denied\n"
 
     def test_listdir(self):
         """
@@ -192,3 +192,25 @@ def test_isfile_false(self):
         response = self.operations.isfile(filename)
 
         assert response is False
+
+    def test_skip_ssl(self):
+        conn_params = ConnectionParams(host=os.getenv('RDBMS_TESTPOOL1_HOST') or '127.0.0.1',
+                                       username=os.getenv('USER'),
+                                       remote=True,
+                                       skip_ssl=True)
+        os_ops = RemoteOperations(conn_params)
+        testgres_config.set_os_ops(os_ops=os_ops)
+        with get_remote_node().init().start() as node:
+            with node.connect() as con:
+                con.begin()
+                con.execute('create table test(val int)')
+                con.execute('insert into test values (1)')
+                con.commit()
+
+                con.begin()
+                con.execute('insert into test values (2)')
+                res = con.execute('select * from test order by val asc')
+                if isinstance(res, list):
+                    res.sort()
+                assert res == [(1,), (2,)]
+
diff --git a/tests/test_simple_remote.py b/tests/test_simple_remote.py
