Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit263ff9c

Browse files
author
ВашÐViktoria Shepard
committed
Remove sshtunnel
1 parent46eb92a commit263ff9c

File tree

6 files changed

+68
-59
lines changed

6 files changed

+68
-59
lines changed

‎setup.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,7 @@
1111
"port-for>=0.4",
1212
"six>=1.9.0",
1313
"psutil",
14-
"packaging",
15-
"sshtunnel"
14+
"packaging"
1615
]
1716

1817
# Add compatibility enum class

‎testgres/exceptions.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,16 @@ def __str__(self):
3232
ifself.out:
3333
msg.append(u'----\n{}'.format(self.out))
3434

35-
returnsix.text_type('\n').join(msg)
35+
returnself.convert_and_join(msg)
36+
37+
@staticmethod
38+
defconvert_and_join(msg_list):
39+
# Convert each byte element in the list to str
40+
str_list= [six.text_type(item,'utf-8')ifisinstance(item,bytes)elsesix.text_type(item)foritemin
41+
msg_list]
42+
43+
# Join the list into a single string with the specified delimiter
44+
returnsix.text_type('\n').join(str_list)
3645

3746

3847
@six.python_2_unicode_compatible

‎testgres/node.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1371,7 +1371,7 @@ def pgbench(self,
13711371
# should be the last one
13721372
_params.append(dbname)
13731373

1374-
proc=self.os_ops.exec_command(_params,stdout=stdout,stderr=stderr,wait_exit=True,proc=True)
1374+
proc=self.os_ops.exec_command(_params,stdout=stdout,stderr=stderr,wait_exit=True,get_process=True)
13751375

13761376
returnproc
13771377

‎testgres/operations/local_ops.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
fromdistutils.spawnimportfind_executable
1919
fromdistutilsimportrmtree
2020

21-
2221
CMD_TIMEOUT_SEC=60
2322
error_markers= [b'error',b'Permission denied',b'fatal']
2423

@@ -37,7 +36,8 @@ def __init__(self, conn_params=None):
3736
# Command execution
3837
defexec_command(self,cmd,wait_exit=False,verbose=False,
3938
expect_error=False,encoding=None,shell=False,text=False,
40-
input=None,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE,proc=None):
39+
input=None,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE,
40+
get_process=None,timeout=None):
4141
"""
4242
Execute a command in a subprocess.
4343
@@ -69,9 +69,14 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
6969
stdout=stdout,
7070
stderr=stderr,
7171
)
72-
ifproc:
72+
ifget_process:
7373
returnprocess
74-
result,error=process.communicate(input)
74+
75+
try:
76+
result,error=process.communicate(input,timeout=timeout)
77+
exceptsubprocess.TimeoutExpired:
78+
process.kill()
79+
raiseExecUtilException("Command timed out after {} seconds.".format(timeout))
7580
exit_status=process.returncode
7681

7782
error_found=exit_status!=0orany(markerinerrorformarkerinerror_markers)

‎testgres/operations/remote_ops.py

Lines changed: 43 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,19 @@
33
importos
44
importsubprocess
55
importtempfile
6-
importtime
76

8-
importsshtunnel
7+
# we support both pg8000 and psycopg2
8+
try:
9+
importpsycopg2aspglib
10+
exceptImportError:
11+
try:
12+
importpg8000aspglib
13+
exceptImportError:
14+
raiseImportError("You must have psycopg2 or pg8000 modules installed")
915

1016
from ..exceptionsimportExecUtilException
1117

1218
from .os_opsimportOsOperations,ConnectionParams
13-
from .os_opsimportpglib
14-
15-
sshtunnel.SSH_TIMEOUT=5.0
16-
sshtunnel.TUNNEL_TIMEOUT=5.0
1719

1820
ConsoleEncoding=locale.getdefaultlocale()[1]
1921
ifnotConsoleEncoding:
@@ -50,21 +52,28 @@ def __init__(self, conn_params: ConnectionParams):
5052
self.remote=True
5153
self.username=conn_params.usernameorself.get_user()
5254
self.add_known_host(self.host)
55+
self.tunnel_process=None
5356

5457
def__enter__(self):
5558
returnself
5659

5760
def__exit__(self,exc_type,exc_val,exc_tb):
58-
self.close_tunnel()
61+
self.close_ssh_tunnel()
5962

60-
defclose_tunnel(self):
61-
ifgetattr(self,'tunnel',None):
62-
self.tunnel.stop(force=True)
63-
start_time=time.time()
64-
whileself.tunnel.is_active:
65-
iftime.time()-start_time>sshtunnel.TUNNEL_TIMEOUT:
66-
break
67-
time.sleep(0.5)
63+
defestablish_ssh_tunnel(self,local_port,remote_port):
64+
"""
65+
Establish an SSH tunnel from a local port to a remote PostgreSQL port.
66+
"""
67+
ssh_cmd= ['-N','-L',f"{local_port}:localhost:{remote_port}"]
68+
self.tunnel_process=self.exec_command(ssh_cmd,get_process=True,timeout=300)
69+
70+
defclose_ssh_tunnel(self):
71+
ifhasattr(self,'tunnel_process'):
72+
self.tunnel_process.terminate()
73+
self.tunnel_process.wait()
74+
delself.tunnel_process
75+
else:
76+
print("No active tunnel to close.")
6877

6978
defadd_known_host(self,host):
7079
cmd='ssh-keyscan -H %s >> /home/%s/.ssh/known_hosts'% (host,os.getlogin())
@@ -78,21 +87,29 @@ def add_known_host(self, host):
7887
raiseExecUtilException(message="Failed to add %s to known_hosts. Error: %s"% (host,str(e)),command=cmd,
7988
exit_code=e.returncode,out=e.stderr)
8089

81-
defexec_command(self,cmd:str,wait_exit=False,verbose=False,expect_error=False,
90+
defexec_command(self,cmd,wait_exit=False,verbose=False,expect_error=False,
8291
encoding=None,shell=True,text=False,input=None,stdin=None,stdout=None,
83-
stderr=None,proc=None):
92+
stderr=None,get_process=None,timeout=None):
8493
"""
8594
Execute a command in the SSH session.
8695
Args:
8796
- cmd (str): The command to be executed.
8897
"""
98+
ssh_cmd= []
8999
ifisinstance(cmd,str):
90100
ssh_cmd= ['ssh',f"{self.username}@{self.host}",'-i',self.ssh_key,cmd]
91101
elifisinstance(cmd,list):
92102
ssh_cmd= ['ssh',f"{self.username}@{self.host}",'-i',self.ssh_key]+cmd
93103
process=subprocess.Popen(ssh_cmd,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
104+
ifget_process:
105+
returnprocess
106+
107+
try:
108+
result,error=process.communicate(input,timeout=timeout)
109+
exceptsubprocess.TimeoutExpired:
110+
process.kill()
111+
raiseExecUtilException("Command timed out after {} seconds.".format(timeout))
94112

95-
result,error=process.communicate(input)
96113
exit_status=process.returncode
97114

98115
ifencoding:
@@ -372,41 +389,19 @@ def get_process_children(self, pid):
372389
raiseExecUtilException(f"Error in getting process children. Error:{result.stderr}")
373390

374391
# Database control
375-
defdb_connect(self,dbname,user,password=None,host="127.0.0.1",port=5432,ssh_key=None):
392+
defdb_connect(self,dbname,user,password=None,host="localhost",port=5432):
376393
"""
377-
Connects to a PostgreSQL database on the remote system.
378-
Args:
379-
- dbname (str): The name of the database to connect to.
380-
- user (str): The username for the database connection.
381-
- password (str, optional): The password for the database connection. Defaults to None.
382-
- host (str, optional): The IP address of the remote system. Defaults to "localhost".
383-
- port (int, optional): The port number of the PostgreSQL service. Defaults to 5432.
384-
385-
This function establishes a connection to a PostgreSQL database on the remote system using the specified
386-
parameters. It returns a connection object that can be used to interact with the database.
394+
Established SSH tunnel and Connects to a PostgreSQL
387395
"""
388-
self.close_tunnel()
389-
self.tunnel=sshtunnel.open_tunnel(
390-
(self.host,22),# Remote server IP and SSH port
391-
ssh_username=self.username,
392-
ssh_pkey=self.ssh_key,
393-
remote_bind_address=(self.host,port),# PostgreSQL server IP and PostgreSQL port
394-
local_bind_address=('localhost',0)
395-
# Local machine IP and available port (0 means it will pick any available port)
396-
)
397-
self.tunnel.start()
398-
396+
self.establish_ssh_tunnel(local_port=port,remote_port=5432)
399397
try:
400-
# Use localhost and self.tunnel.local_bind_port to connect
401398
conn=pglib.connect(
402-
host='localhost',# Connect to localhost
403-
port=self.tunnel.local_bind_port,# use the local bindport set up by the tunnel
399+
host=host,
400+
port=port,
404401
database=dbname,
405-
user=userorself.username,
406-
password=password
402+
user=user,
403+
password=password,
407404
)
408-
409405
returnconn
410406
exceptExceptionase:
411-
self.tunnel.stop()
412-
raiseExecUtilException("Could not create db tunnel. {}".format(e))
407+
raiseException(f"Could not connect to the database. Error:{e}")

‎tests/test_simple_remote.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -735,9 +735,10 @@ def test_pgbench(self):
735735
options=['-q']).pgbench_run(time=2)
736736

737737
# run TPC-B benchmark
738-
out=node.pgbench(stdout=subprocess.PIPE,
739-
stderr=subprocess.STDOUT,
740-
options=['-T3'])
738+
proc=node.pgbench(stdout=subprocess.PIPE,
739+
stderr=subprocess.STDOUT,
740+
options=['-T3'])
741+
out=proc.communicate()[0]
741742
self.assertTrue(b'tps = 'inout)
742743

743744
deftest_pg_config(self):

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp