🌐 AI搜索 & 代理 主页
Skip to content

Commit 4e9e543

Browse files
author
v.shepard
committed
PBCKP-588 fix failed tests - psql, set_auto_conf
1 parent 190d084 commit 4e9e543

File tree

3 files changed

+54
-32
lines changed

3 files changed

+54
-32
lines changed

‎testgres/node.py‎

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import os
44
import random
55
import signal
6+
import subprocess
67
import threading
78
from queue import Queue
89

@@ -958,19 +959,31 @@ def psql(self,
958959

959960
# select query source
960961
if query:
961-
psql_params.extend(("-c", '"{}"'.format(query)))
962+
if self.os_ops.remote:
963+
psql_params.extend(("-c", '"{}"'.format(query)))
964+
else:
965+
psql_params.extend(("-c", query))
962966
elif filename:
963967
psql_params.extend(("-f", filename))
964968
else:
965969
raise QueryException('Query or filename must be provided')
966970

967971
# should be the last one
968972
psql_params.append(dbname)
973+
if not self.os_ops.remote:
974+
# start psql process
975+
process = subprocess.Popen(psql_params,
976+
stdin=subprocess.PIPE,
977+
stdout=subprocess.PIPE,
978+
stderr=subprocess.PIPE)
979+
980+
# wait until it finishes and get stdout and stderr
981+
out, err = process.communicate(input=input)
982+
return process.returncode, out, err
983+
else:
984+
status_code, out, err = self.os_ops.exec_command(psql_params, verbose=True, input=input)
969985

970-
# start psql process
971-
status_code, out, err = self.os_ops.exec_command(psql_params, verbose=True, input=input)
972-
973-
return status_code, out, err
986+
return status_code, out, err
974987

975988
@method_decorator(positional_args_hack(['dbname', 'query']))
976989
def safe_psql(self, query=None, expect_error=False, **kwargs):
@@ -1002,9 +1015,9 @@ def safe_psql(self, query=None, expect_error=False, **kwargs):
10021015
err = e.message
10031016
if ret:
10041017
if expect_error:
1005-
out = err or b''
1018+
out = (err or b'').decode('utf-8')
10061019
else:
1007-
raise QueryException(err or b'', query)
1020+
raise QueryException((err or b'').decode('utf-8'), query)
10081021
elif expect_error:
10091022
assert False, "Exception was expected, but query finished successfully: `{}` ".format(query)
10101023

@@ -1529,18 +1542,18 @@ def set_auto_conf(self, options, config='postgresql.auto.conf', rm_options={}):
15291542
Defaults to an empty set.
15301543
"""
15311544
# parse postgresql.auto.conf
1532-
auto_conf_file = os.path.join(self.data_dir, config)
1533-
raw_content = self.os_ops.read(auto_conf_file)
1545+
path = os.path.join(self.data_dir, config)
15341546

1547+
lines = self.os_ops.readlines(path)
15351548
current_options = {}
15361549
current_directives = []
1537-
for line in raw_content.splitlines():
1550+
for line in lines:
15381551

15391552
# ignore comments
15401553
if line.startswith('#'):
15411554
continue
15421555

1543-
if line == '':
1556+
if line.strip() == '':
15441557
continue
15451558

15461559
if line.startswith('include'):
@@ -1570,7 +1583,7 @@ def set_auto_conf(self, options, config='postgresql.auto.conf', rm_options={}):
15701583
for directive in current_directives:
15711584
auto_conf += directive + "\n"
15721585

1573-
self.os_ops.write(auto_conf_file, auto_conf)
1586+
self.os_ops.write(path, auto_conf, truncate=True)
15741587

15751588

15761589
class NodeApp:

‎testgres/operations/local_ops.py‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
import getpass
22
import os
33
import shutil
4+
import stat
45
import subprocess
56
import tempfile
67
from shutil import rmtree
78

89
import psutil
910

1011
from ..exceptions import ExecUtilException
11-
12-
from .os_ops import OsOperations, ConnectionParams
12+
from .os_ops import ConnectionParams, OsOperations
1313
from .os_ops import pglib
1414

1515
try:
@@ -18,20 +18,24 @@
1818
from distutils.spawn import find_executable
1919

2020
CMD_TIMEOUT_SEC = 60
21+
error_markers = [b'error', b'Permission denied', b'fatal']
2122

2223

2324
class LocalOperations(OsOperations):
24-
def __init__(self, conn_params: ConnectionParams = ConnectionParams()):
25-
super().__init__(conn_params.username)
25+
def __init__(self, conn_params=None):
26+
if conn_params is None:
27+
conn_params = ConnectionParams()
28+
super(LocalOperations, self).__init__(conn_params.username)
2629
self.conn_params = conn_params
2730
self.host = conn_params.host
2831
self.ssh_key = None
32+
self.remote = False
2933
self.username = conn_params.username or self.get_user()
3034

3135
# Command execution
3236
def exec_command(self, cmd, wait_exit=False, verbose=False,
33-
expect_error=False, encoding=None, shell=True, text=False,
34-
input=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, proc=None):
37+
expect_error=False, encoding=None, shell=False, text=False,
38+
input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, proc=None):
3539
"""
3640
Execute a command in a subprocess.
3741
@@ -49,9 +53,6 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
4953
- proc: The process to use for subprocess creation.
5054
:return: The output of the subprocess.
5155
"""
52-
if isinstance(cmd, list):
53-
cmd = ' '.join(item.decode('utf-8') if isinstance(item, bytes) else item for item in cmd)
54-
5556
if os.name == 'nt':
5657
with tempfile.NamedTemporaryFile() as buf:
5758
process = subprocess.Popen(cmd, stdout=buf, stderr=subprocess.STDOUT)
@@ -71,7 +72,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
7172
result, error = process.communicate(input)
7273
exit_status = process.returncode
7374

74-
found_error = "error" in error.decode(encoding or 'utf-8').lower()
75+
error_found = exit_status != 0 or any(marker in error for marker in error_markers)
7576

7677
if encoding:
7778
result = result.decode(encoding)
@@ -80,7 +81,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
8081
if expect_error:
8182
raise Exception(result, error)
8283

83-
if exit_status != 0 or found_error:
84+
if exit_status != 0 or error_found:
8485
if exit_status == 0:
8586
exit_status = 1
8687
raise ExecUtilException(message='Utility exited with non-zero code. Error `{}`'.format(error),
@@ -101,7 +102,7 @@ def find_executable(self, executable):
101102

102103
def is_executable(self, file):
103104
# Check if the file is executable
104-
return os.access(file, os.X_OK)
105+
return os.stat(file).st_mode & stat.S_IXUSR
105106

106107
def set_env(self, var_name, var_val):
107108
# Check if the directory is already in PATH
@@ -116,9 +117,12 @@ def get_name(self):
116117

117118
# Work with dirs
118119
def makedirs(self, path, remove_existing=False):
119-
if remove_existing and os.path.exists(path):
120-
shutil.rmtree(path)
121-
os.makedirs(path, exist_ok=True)
120+
if remove_existing:
121+
shutil.rmtree(path, ignore_errors=True)
122+
try:
123+
os.makedirs(path)
124+
except FileExistsError:
125+
pass
122126

123127
def rmdirs(self, path, ignore_errors=True):
124128
return rmtree(path, ignore_errors=ignore_errors)
@@ -141,7 +145,7 @@ def pathsep(self):
141145
return pathsep
142146

143147
def mkdtemp(self, prefix=None):
144-
return tempfile.mkdtemp(prefix=prefix)
148+
return tempfile.mkdtemp(prefix='{}'.format(prefix))
145149

146150
def mkstemp(self, prefix=None):
147151
fd, filename = tempfile.mkstemp(prefix=prefix)

‎testgres/operations/remote_ops.py‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
sshtunnel.TUNNEL_TIMEOUT = 5.0
1818

1919

20-
error_markers = [b'error', b'Permission denied']
20+
error_markers = [b'error', b'Permission denied', b'fatal']
2121

2222

2323
class PsUtilProcessProxy:
@@ -43,6 +43,7 @@ def __init__(self, conn_params: ConnectionParams):
4343
self.host = conn_params.host
4444
self.ssh_key = conn_params.ssh_key
4545
self.ssh = self.ssh_connect()
46+
self.remote = True
4647
self.username = conn_params.username or self.get_user()
4748
self.tunnel = None
4849

@@ -89,7 +90,7 @@ def _read_ssh_key(self):
8990
ExecUtilException(message="An error occurred while reading the ssh key: {}".format(e))
9091

9192
def exec_command(self, cmd: str, wait_exit=False, verbose=False, expect_error=False,
92-
encoding=None, shell=True, text=False, input=None, stdout=None,
93+
encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None,
9394
stderr=None, proc=None):
9495
"""
9596
Execute a command in the SSH session.
@@ -131,7 +132,11 @@ def exec_command(self, cmd: str, wait_exit=False, verbose=False, expect_error=Fa
131132
if error_found:
132133
if exit_status == 0:
133134
exit_status = 1
134-
raise ExecUtilException(message="Utility exited with non-zero code. Error: {}".format(error.decode(encoding or 'utf-8')),
135+
if encoding:
136+
message = "Utility exited with non-zero code. Error: {}".format(error.decode(encoding))
137+
else:
138+
message = b"Utility exited with non-zero code. Error: " + error
139+
raise ExecUtilException(message=message,
135140
command=cmd,
136141
exit_code=exit_status,
137142
out=result)
@@ -429,7 +434,7 @@ def db_connect(self, dbname, user, password=None, host="127.0.0.1", port=5432, s
429434
conn = pglib.connect(
430435
host=host, # change to 'localhost' because we're connecting through a local ssh tunnel
431436
port=self.tunnel.local_bind_port, # use the local bind port set up by the tunnel
432-
dbname=dbname,
437+
database=dbname,
433438
user=user or self.username,
434439
password=password
435440
)

0 commit comments

Comments
 (0)