PBCKP-152 multihost

v.shepard · v.shepard · commit 4f38bd505ee9 · 2023-05-03T11:50:35.000+02:00
diff --git a/testgres/cache.py b/testgres/cache.py
@@ -39,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
@@ -55,12 +55,6 @@ def cached_initdb_dir(self):
     def cached_initdb_dir(self, value):
         self._cached_initdb_dir = value
 
-        # NOTE: assign initial cached dir for initdb
-        if self.os_ops:
-            testgres_config.cached_initdb_dir = self.os_ops.mkdtemp(prefix=TMP_CACHE)
-        else:
-            testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
-
         if value:
             cached_initdb_dirs.add(value)
         return testgres_config.cached_initdb_dir
@@ -208,3 +202,7 @@ def configure_testgres(**options):
     """
 
     testgres_config.update(options)
+
+
+# NOTE: assign initial cached dir for initdb
+testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
diff --git a/testgres/connection.py b/testgres/connection.py
@@ -42,8 +42,7 @@ def __init__(self,
 
         self._node = node
 
-        self.os_ops = OsOperations(node.host, node.hostname, node.ssh_key, node.username)
-        self._connection = self.os_ops.db_connect(dbname=dbname,
+        self._connection = node.os_ops.db_connect(dbname=dbname,
                                                   user=username,
                                                   password=password,
                                                   host=node.host,
diff --git a/testgres/os_ops.py b/testgres/os_ops.py
@@ -262,32 +262,25 @@ def get_user(self):
             user = getpass.getuser()
         return user
 
-    @contextmanager
     def db_connect(self, dbname, user, password=None, host='localhost', port=5432):
         if self.remote:
-            with self.ssh_connect() as ssh:
-                # Set up a local port forwarding on a random port
-                local_port = ssh.forward_remote_port(host, port)
-                conn = pglib.connect(
-                    host=host,
-                    port=local_port,
-                    dbname=dbname,
-                    user=user,
-                    password=password,
-                )
-                try:
-                    yield conn
-                finally:
-                    conn.close()
-                    ssh.close_forwarded_tcp(local_port)
+            local_port = self.ssh.forward_remote_port(host, port)
+            conn = pglib.connect(
+                host=host,
+                port=local_port,
+                dbname=dbname,
+                user=user,
+                password=password,
+            )
         else:
-            with pglib.connect(
+            conn = pglib.connect(
                     host=host,
                     port=port,
                     dbname=dbname,
                     user=user,
                     password=password,
-            ) as conn:
-                yield conn
+            )
+        return conn
+
 
 
