|
5 | 5 | importplatform |
6 | 6 | importtime |
7 | 7 |
|
8 | | -from ..utilsimportreserve_port |
9 | | - |
10 | 8 | # we support both pg8000 and psycopg2 |
11 | 9 | try: |
12 | 10 | importpsycopg2aspglib |
|
17 | 15 | raiseImportError("You must have psycopg2 or pg8000 modules installed") |
18 | 16 |
|
19 | 17 | from ..exceptionsimportExecUtilException |
20 | | -from ..utilsimportreserve_port |
21 | 18 | from .os_opsimportOsOperations,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 | exceptsocket.error: |
77 | 74 | returnFalse |
78 | 75 |
|
79 | | -defestablish_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 | | -ifhost!='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 | | -whiletime.time()-start_time<timeout: |
91 | | -ifself.is_port_open('localhost',local_port): |
92 | | -print("SSH tunnel established.") |
93 | | -return |
94 | | -time.sleep(0.5) |
95 | | -raiseException("Failed to establish SSH tunnel within the timeout period.") |
96 | | - |
97 | 76 | defclose_ssh_tunnel(self): |
98 | 77 | ifself.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 | defdb_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 | | -returnconn |
430 | | -exceptExceptionase: |
431 | | -print(f"Error connecting to the database:{str(e)}") |
432 | | -ifself.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 | +returnconn |