|
5 | 5 | import platform |
6 | 6 | import time |
7 | 7 |
|
8 | | -from ..utils import reserve_port |
9 | | - |
10 | 8 | # we support both pg8000 and psycopg2 |
11 | 9 | try: |
12 | 10 | import psycopg2 as pglib |
|
17 | 15 | raise ImportError("You must have psycopg2 or pg8000 modules installed") |
18 | 16 |
|
19 | 17 | from ..exceptions import ExecUtilException |
20 | | -from ..utils import reserve_port |
21 | 18 | from .os_ops import OsOperations, ConnectionParams, get_default_encoding |
22 | 19 |
|
23 | 20 | error_markers = [b'error', b'Permission denied', b'fatal', b'No such file or directory'] |
@@ -76,24 +73,6 @@ def is_port_open(host, port): |
76 | 73 | except socket.error: |
77 | 74 | return False |
78 | 75 |
|
79 | | - def establish_ssh_tunnel(self, local_port, remote_port, host): |
80 | | - """ |
81 | | - Establish an SSH tunnel from a local port to a remote PostgreSQL port. |
82 | | - """ |
83 | | - if host != 'localhost': |
84 | | - ssh_cmd = ['-N', '-L', f"localhost:{local_port}:{host}:{remote_port}"] |
85 | | - else: |
86 | | - ssh_cmd = ['-N', '-L', f"{local_port}:{host}:{remote_port}"] |
87 | | - self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300) |
88 | | - timeout = 10 |
89 | | - start_time = time.time() |
90 | | - while time.time() - start_time < timeout: |
91 | | - if self.is_port_open('localhost', local_port): |
92 | | - print("SSH tunnel established.") |
93 | | - return |
94 | | - time.sleep(0.5) |
95 | | - raise Exception("Failed to establish SSH tunnel within the timeout period.") |
96 | | - |
97 | 76 | def close_ssh_tunnel(self): |
98 | 77 | if self.tunnel_process: |
99 | 78 | self.tunnel_process.terminate() |
@@ -410,26 +389,11 @@ def get_process_children(self, pid): |
410 | 389 |
|
411 | 390 | # Database control |
412 | 391 | def db_connect(self, dbname, user, password=None, host="localhost", port=5432): |
413 | | - """ |
414 | | - Establish SSH tunnel and connect to a PostgreSQL database. |
415 | | - """ |
416 | | - local_port = reserve_port() |
417 | | - self.tunnel_port = local_port |
418 | | - self.establish_ssh_tunnel(local_port=local_port, remote_port=port, host=host) |
419 | | - try: |
420 | | - conn = pglib.connect( |
421 | | - host='localhost', |
422 | | - port=local_port, |
423 | | - database=dbname, |
424 | | - user=user, |
425 | | - password=password, |
426 | | - timeout=10 |
427 | | - ) |
428 | | - print("Database connection established successfully.") |
429 | | - return conn |
430 | | - except Exception as e: |
431 | | - print(f"Error connecting to the database: {str(e)}") |
432 | | - if self.tunnel_process: |
433 | | - self.tunnel_process.terminate() |
434 | | - print("SSH tunnel closed due to connection failure.") |
435 | | - raise |
| 392 | + conn = pglib.connect( |
| 393 | + host=host, |
| 394 | + port=port, |
| 395 | + database=dbname, |
| 396 | + user=user, |
| 397 | + password=password, |
| 398 | + ) |
| 399 | + return conn |
0 commit comments