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

Commit1c405ef

Browse files
author
v.shepard
committed
PBCKP-152 add tests for remote_ops.py
1 parente098b97 commit1c405ef

File tree

3 files changed

+313
-146
lines changed

3 files changed

+313
-146
lines changed

‎testgres/node.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -152,10 +152,10 @@ def __init__(self, name=None, port=None, base_dir=None,
152152
self.host=host
153153
self.hostname=hostname
154154
self.ssh_key=ssh_key
155-
ifhostname=='localhost'orhost=='127.0.0.1':
156-
self.os_ops=LocalOperations(username=username)
157-
else:
155+
ifhostname!='localhost'orhost!='127.0.0.1':
158156
self.os_ops=RemoteOperations(host,hostname,ssh_key)
157+
else:
158+
self.os_ops=LocalOperations(username=username)
159159

160160
testgres_config.os_ops=self.os_ops
161161
# defaults for __exit__()

‎testgres/operations/remote_ops.py

Lines changed: 151 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -1,106 +1,98 @@
1-
importio
21
importos
32
importtempfile
4-
fromcontextlibimportcontextmanager
3+
fromtypingimportOptional
54

6-
fromtestgres.loggerimportlog
5+
importparamiko
6+
fromparamikoimportSSHClient
77

8+
fromloggerimportlog
89
from .os_opsimportOsOperations
910
from .os_opsimportpglib
1011

11-
importparamiko
12+
error_markers= [b'error',b'Permission denied']
1213

1314

1415
classRemoteOperations(OsOperations):
15-
"""
16-
This class specifically supports work with Linux systems. It utilizes the SSH
17-
for making connections and performing various file and directory operations, command executions,
18-
environment setup and management, process control, and database connections.
19-
It uses the Paramiko library for SSH connections and operations.
20-
21-
Some methods are designed to work with specific Linux shell commands, and thus may not work as expected
22-
on other non-Linux systems.
23-
24-
Attributes:
25-
- hostname (str): The remote system's hostname. Default 'localhost'.
26-
- host (str): The remote system's IP address. Default '127.0.0.1'.
27-
- ssh_key (str): Path to the SSH private key for authentication.
28-
- username (str): Username for the remote system.
29-
- ssh (paramiko.SSHClient): SSH connection to the remote system.
30-
"""
31-
32-
def__init__(
33-
self,hostname="localhost",host="127.0.0.1",ssh_key=None,username=None
34-
):
16+
def__init__(self,hostname="localhost",host="127.0.0.1",ssh_key=None,username=None):
3517
super().__init__(username)
36-
self.hostname=hostname
3718
self.host=host
3819
self.ssh_key=ssh_key
3920
self.remote=True
40-
self.ssh=self.connect()
21+
self.ssh=self.ssh_connect()
4122
self.username=usernameorself.get_user()
4223

4324
def__del__(self):
4425
ifself.ssh:
4526
self.ssh.close()
4627

47-
@contextmanager
48-
defssh_connect(self):
28+
defssh_connect(self)->Optional[SSHClient]:
4929
ifnotself.remote:
50-
yieldNone
30+
returnNone
5131
else:
32+
key=self._read_ssh_key()
33+
ssh=paramiko.SSHClient()
34+
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
35+
ssh.connect(self.host,username=self.username,pkey=key)
36+
returnssh
37+
38+
def_read_ssh_key(self):
39+
try:
5240
withopen(self.ssh_key,"r")asf:
5341
key_data=f.read()
5442
if"BEGIN OPENSSH PRIVATE KEY"inkey_data:
5543
key=paramiko.Ed25519Key.from_private_key_file(self.ssh_key)
5644
else:
5745
key=paramiko.RSAKey.from_private_key_file(self.ssh_key)
46+
returnkey
47+
exceptFileNotFoundError:
48+
log.error(f"No such file or directory: '{self.ssh_key}'")
49+
exceptExceptionase:
50+
log.error(f"An error occurred while reading the ssh key:{e}")
5851

59-
withparamiko.SSHClient()asssh:
60-
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
61-
ssh.connect(self.host,username=self.username,pkey=key)
62-
yieldssh
63-
64-
defconnect(self):
65-
withself.ssh_connect()asssh:
66-
returnssh
67-
68-
# Command execution
69-
defexec_command(self,cmd,wait_exit=False,verbose=False,
70-
expect_error=False,encoding=None,shell=True,text=False,
71-
input=None,stdout=None,stderr=None,proc=None):
52+
defexec_command(self,cmd:str,wait_exit=False,verbose=False,expect_error=False,
53+
encoding=None,shell=True,text=False,input=None,stdout=None,
54+
stderr=None,proc=None):
55+
"""
56+
Execute a command in the SSH session.
57+
Args:
58+
- cmd (str): The command to be executed.
59+
"""
7260
ifisinstance(cmd,list):
7361
cmd=" ".join(cmd)
74-
log.debug(f"os_ops.exec_command: `{cmd}`; remote={self.remote}")
75-
# Source global profile file + execute command
7662
try:
77-
cmd=f"source /etc/profile.d/custom.sh;{cmd}"
78-
withself.ssh_connect()asssh:
79-
ifinput:
80-
# encode input and feed it to stdin
81-
stdin,stdout,stderr=ssh.exec_command(cmd)
82-
stdin.write(input)
83-
stdin.flush()
84-
else:
85-
stdin,stdout,stderr=ssh.exec_command(cmd)
86-
exit_status=0
87-
ifwait_exit:
88-
exit_status=stdout.channel.recv_exit_status()
89-
ifencoding:
90-
result=stdout.read().decode(encoding)
91-
error=stderr.read().decode(encoding)
92-
else:
93-
# Save as binary string
94-
result=io.BytesIO(stdout.read()).getvalue()
95-
error=io.BytesIO(stderr.read()).getvalue()
96-
error_str=stderr.read()
63+
ifinput:
64+
stdin,stdout,stderr=self.ssh.exec_command(cmd)
65+
stdin.write(input.encode("utf-8"))
66+
stdin.flush()
67+
else:
68+
stdin,stdout,stderr=self.ssh.exec_command(cmd)
69+
exit_status=0
70+
ifwait_exit:
71+
exit_status=stdout.channel.recv_exit_status()
72+
73+
ifencoding:
74+
result=stdout.read().decode(encoding)
75+
error=stderr.read().decode(encoding)
76+
else:
77+
result=stdout.read()
78+
error=stderr.read()
9779

9880
ifexpect_error:
9981
raiseException(result,error)
100-
ifexit_status!=0or'error'inerror_str:
82+
83+
ifencoding:
84+
error_found=exit_status!=0orany(
85+
marker.decode(encoding)inerrorformarkerinerror_markers)
86+
else:
87+
error_found=exit_status!=0orany(
88+
markerinerrorformarkerinerror_markers)
89+
90+
iferror_found:
10191
log.error(
10292
f"Problem in executing command: `{cmd}`\nerror:{error}\nexit_code:{exit_status}"
10393
)
94+
ifexit_status==0:
95+
exit_status=1
10496

10597
ifverbose:
10698
returnexit_status,result,error
@@ -112,7 +104,12 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
112104
returnNone
113105

114106
# Environment setup
115-
defenviron(self,var_name):
107+
defenviron(self,var_name:str)->str:
108+
"""
109+
Get the value of an environment variable.
110+
Args:
111+
- var_name (str): The name of the environment variable.
112+
"""
116113
cmd=f"echo ${var_name}"
117114
returnself.exec_command(cmd).strip()
118115

@@ -131,7 +128,8 @@ def find_executable(self, executable):
131128

132129
defis_executable(self,file):
133130
# Check if the file is executable
134-
returnself.exec_command(f"test -x{file} && echo OK")=="OK\n"
131+
is_exec=self.exec_command(f"test -x{file} && echo OK")
132+
returnis_exec==b"OK\n"
135133

136134
defadd_to_path(self,new_path):
137135
pathsep=self.pathsep
@@ -144,8 +142,13 @@ def add_to_path(self, new_path):
144142
os.environ["PATH"]=f"{new_path}{pathsep}{path}"
145143
returnpathsep
146144

147-
defset_env(self,var_name,var_val):
148-
# Check if the directory is already in PATH
145+
defset_env(self,var_name:str,var_val:str)->None:
146+
"""
147+
Set the value of an environment variable.
148+
Args:
149+
- var_name (str): The name of the environment variable.
150+
- var_val (str): The value to be set for the environment variable.
151+
"""
149152
returnself.exec_command(f"export{var_name}={var_val}")
150153

151154
# Get environment variables
@@ -158,22 +161,47 @@ def get_name(self):
158161

159162
# Work with dirs
160163
defmakedirs(self,path,remove_existing=False):
164+
"""
165+
Create a directory in the remote server.
166+
Args:
167+
- path (str): The path to the directory to be created.
168+
- remove_existing (bool): If True, the existing directory at the path will be removed.
169+
"""
161170
ifremove_existing:
162171
cmd=f"rm -rf{path} && mkdir -p{path}"
163172
else:
164173
cmd=f"mkdir -p{path}"
165-
returnself.exec_command(cmd)
174+
exit_status,result,error=self.exec_command(cmd,verbose=True)
175+
ifexit_status!=0:
176+
raiseException(f"Couldn't create dir{path} because of error{error}")
177+
returnresult
166178

167-
defrmdirs(self,path,ignore_errors=True):
179+
defrmdirs(self,path,verbose=False,ignore_errors=True):
180+
"""
181+
Remove a directory in the remote server.
182+
Args:
183+
- path (str): The path to the directory to be removed.
184+
- verbose (bool): If True, return exit status, result, and error.
185+
- ignore_errors (bool): If True, do not raise error if directory does not exist.
186+
"""
168187
cmd=f"rm -rf{path}"
169-
returnself.exec_command(cmd)
188+
exit_status,result,error=self.exec_command(cmd,verbose=True)
189+
ifverbose:
190+
returnexit_status,result,error
191+
else:
192+
returnresult
170193

171194
deflistdir(self,path):
195+
"""
196+
List all files and directories in a directory.
197+
Args:
198+
path (str): The path to the directory.
199+
"""
172200
result=self.exec_command(f"ls{path}")
173201
returnresult.splitlines()
174202

175203
defpath_exists(self,path):
176-
result=self.exec_command(f"test -e{path}; echo $?")
204+
result=self.exec_command(f"test -e{path}; echo $?",encoding='utf-8')
177205
returnint(result.strip())==0
178206

179207
@property
@@ -188,7 +216,12 @@ def pathsep(self):
188216
returnpathsep
189217

190218
defmkdtemp(self,prefix=None):
191-
temp_dir=self.exec_command(f"mkdtemp -d{prefix}")
219+
"""
220+
Creates a temporary directory in the remote server.
221+
Args:
222+
prefix (str): The prefix of the temporary directory name.
223+
"""
224+
temp_dir=self.exec_command(f"mkdtemp -d{prefix}",encoding='utf-8')
192225
returntemp_dir.strip()
193226

194227
defmkstemp(self,prefix=None):
@@ -200,18 +233,19 @@ def copytree(self, src, dst):
200233
returnself.exec_command(f"cp -r{src}{dst}")
201234

202235
# Work with files
203-
defwrite(self,filename,data,truncate=False,binary=False,read_and_write=False):
236+
defwrite(self,filename,data,truncate=False,binary=False,read_and_write=False,encoding='utf-8'):
204237
"""
205238
Write data to a file on a remote host
239+
206240
Args:
207-
filename: The file path where the data will be written.
208-
data: The data to be written to the file.
209-
truncate: If True, the file will be truncated before writing ('w' or 'wb' option);
210-
if False (default), data will be appended ('a' or 'ab' option).
211-
binary: If True, the data will be written in binary mode ('wb' or 'ab' option);
212-
if False (default), the data will be written in text mode ('w' or 'a' option).
213-
read_and_write: If True, the file will be opened with read and write permissions ('r+' option);
214-
if False (default), only write permission will be used ('w', 'a', 'wb', or 'ab' option)
241+
-filename (str): The file path where the data will be written.
242+
- data (bytes or str): The data to be written to the file.
243+
-truncate (bool): If True, the file will be truncated before writing ('w' or 'wb' option);
244+
if False (default), data will be appended ('a' or 'ab' option).
245+
-binary (bool): If True, the data will be written in binary mode ('wb' or 'ab' option);
246+
if False (default), the data will be written in text mode ('w' or 'a' option).
247+
-read_and_write (bool): If True, the file will be opened with read and write permissions ('r+' option);
248+
if False (default), only write permission will be used ('w', 'a', 'wb', or 'ab' option).
215249
"""
216250
mode="wb"ifbinaryelse"w"
217251
ifnottruncate:
@@ -220,15 +254,18 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
220254
mode="r+b"ifbinaryelse"r+"
221255

222256
withtempfile.NamedTemporaryFile(mode=mode)astmp_file:
223-
ifisinstance(data,list):
224-
tmp_file.writelines(data)
225-
else:
226-
tmp_file.write(data)
257+
ifisinstance(data,bytes)andnotbinary:
258+
data=data.decode(encoding)
259+
elifisinstance(data,str)andbinary:
260+
data=data.encode(encoding)
261+
262+
tmp_file.write(data)
227263
tmp_file.flush()
228264

229-
sftp=self.ssh.open_sftp()
230-
sftp.put(tmp_file.name,filename)
231-
sftp.close()
265+
withself.ssh_connect()asssh:
266+
sftp=ssh.open_sftp()
267+
sftp.put(tmp_file.name,filename)
268+
sftp.close()
232269

233270
deftouch(self,filename):
234271
"""
@@ -281,8 +318,29 @@ def get_pid(self):
281318
returnself.exec_command("echo $$")
282319

283320
# Database control
284-
defdb_connect(self,dbname,user,password=None,host="localhost",port=5432):
285-
local_port=self.ssh.forward_remote_port(host,port)
321+
defdb_connect(self,dbname,user,password=None,host="127.0.0.1",hostname="localhost",port=5432):
322+
"""
323+
Connects to a PostgreSQL database on the remote system.
324+
Args:
325+
- dbname (str): The name of the database to connect to.
326+
- user (str): The username for the database connection.
327+
- password (str, optional): The password for the database connection. Defaults to None.
328+
- host (str, optional): The IP address of the remote system. Defaults to "127.0.0.1".
329+
- hostname (str, optional): The hostname of the remote system. Defaults to "localhost".
330+
- port (int, optional): The port number of the PostgreSQL service. Defaults to 5432.
331+
332+
This function establishes a connection to a PostgreSQL database on the remote system using the specified
333+
parameters. It returns a connection object that can be used to interact with the database.
334+
"""
335+
transport=self.ssh.get_transport()
336+
local_port=9090# or any other available port
337+
338+
transport.open_channel(
339+
'direct-tcpip',
340+
(hostname,port),
341+
(host,local_port)
342+
)
343+
286344
conn=pglib.connect(
287345
host=host,
288346
port=local_port,
@@ -291,3 +349,4 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
291349
password=password,
292350
)
293351
returnconn
352+

0 commit comments

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