Skip to content

Commit 71f54f5

Browse files
author
v.shepard
committed
PBCKP-152 fix failed tests
1 parent ac77ef7 commit 71f54f5

File tree

7 files changed

+1086
-65
lines changed

7 files changed

+1086
-65
lines changed

‎testgres/cache.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ def call_initdb(initdb_dir, log=None):
3838
call_initdb(data_dir, logfile)
3939
else:
4040
# Fetch cached initdb dir
41-
cached_data_dir = testgres_config.cached_initdb_dir()
41+
cached_data_dir = testgres_config.cached_initdb_dir
4242

4343
# Initialize cached initdb
4444

‎testgres/connection.py

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -102,23 +102,15 @@ def rollback(self):
102102
return self
103103

104104
def execute(self, query, *args):
105+
self.cursor.execute(query, args)
105106
try:
106-
with self.connection.cursor() as cursor:
107-
cursor.execute(query, args)
108-
try:
109-
res = cursor.fetchall()
110-
111-
# pg8000 might return tuples
112-
if isinstance(res, tuple):
113-
res = [tuple(t) for t in res]
114-
115-
return res
116-
except (pglib.ProgrammingError, pglib.InternalError) as e:
117-
# An error occurred while trying to fetch results (e.g., no results to fetch)
118-
print(f"Error fetching results: {e}")
119-
return None
120-
except (pglib.Error, Exception) as e:
121-
# Handle other database errors
107+
res = self.cursor.fetchall()
108+
# pg8000 might return tuples
109+
if isinstance(res, tuple):
110+
res = [tuple(t) for t in res]
111+
112+
return res
113+
except Exception as e:
122114
print(f"Error executing query: {e}")
123115
return None
124116

‎testgres/node.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,7 @@ def pid(self):
201201

202202
if self.status():
203203
pid_file = os.path.join(self.data_dir, PG_PID_FILE)
204-
lines = self.os_ops.readlines(pid_file, num_lines=1)
204+
lines = self.os_ops.readlines(pid_file)
205205
pid = int(lines[0]) if lines else None
206206
return pid
207207

@@ -433,7 +433,8 @@ def _collect_special_files(self):
433433
if not self.os_ops.path_exists(f):
434434
continue
435435

436-
lines = b''.join(self.os_ops.readlines(f, num_lines, encoding='utf-8'))
436+
file_lines = self.os_ops.readlines(f, num_lines, binary=True, encoding=None)
437+
lines = b''.join(file_lines)
437438

438439
# fill list
439440
result.append((f, lines))
@@ -498,7 +499,7 @@ def default_conf(self,
498499
]
499500

500501
# write filtered lines
501-
self.os_ops.write(hba_conf_file, lines, truncate=True)
502+
self.os_ops.write(hba_conf, lines, truncate=True)
502503

503504
# replication-related settings
504505
if allow_streaming:
@@ -708,7 +709,8 @@ def start(self, params=[], wait=True):
708709
] + params # yapf: disable
709710

710711
try:
711-
execute_utility(_params, self.utils_log_file, os_ops=self.os_ops)
712+
res = execute_utility(_params, self.utils_log_file, os_ops=self.os_ops)
713+
print()
712714
except ExecUtilException as e:
713715
msg = 'Cannot start node'
714716
files = self._collect_special_files()
@@ -960,11 +962,9 @@ def psql(self,
960962
psql_params.append(dbname)
961963

962964
# start psql process
963-
process = self.os_ops.exec_command(psql_params)
965+
status_code, out, err = self.os_ops.exec_command(psql_params, shell=False, verbose=True, input=input)
964966

965-
# wait until it finishes and get stdout and stderr
966-
out, err = process.communicate(input=input)
967-
return process.returncode, out, err
967+
return status_code, out, err
968968

969969
@method_decorator(positional_args_hack(['dbname', 'query']))
970970
def safe_psql(self, query=None, expect_error=False, **kwargs):
@@ -1348,7 +1348,7 @@ def pgbench(self,
13481348
# should be the last one
13491349
_params.append(dbname)
13501350

1351-
proc = self.os_ops.exec_command(_params, wait_exit=True)
1351+
proc = self.os_ops.exec_command(_params, stdout=stdout, stderr=stderr, wait_exit=True, shell=False, proc=True)
13521352

13531353
return proc
13541354

‎testgres/os_ops/local_ops.py

Lines changed: 26 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,25 @@ def __init__(self, username=None):
1919
self.username = username or self.get_user()
2020

2121
# Command execution
22-
def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False):
23-
if isinstance(cmd, list):
24-
cmd = " ".join(cmd)
22+
def exec_command(self, cmd, wait_exit=False, verbose=False,
23+
expect_error=False, encoding=None, shell=True, text=False,
24+
input=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE, proc=None):
2525
log.debug(f"os_ops.exec_command: `{cmd}`; remote={self.remote}")
2626
# Source global profile file + execute command
2727
try:
28+
if proc:
29+
return subprocess.Popen(cmd,
30+
shell=shell,
31+
stdin=input or subprocess.PIPE,
32+
stdout=stdout,
33+
stderr=stderr)
2834
process = subprocess.run(
2935
cmd,
30-
shell=True,
31-
text=True,
32-
stdout=subprocess.PIPE,
33-
stderr=subprocess.PIPE,
36+
input=input,
37+
shell=shell,
38+
text=text,
39+
stdout=stdout,
40+
stderr=stderr,
3441
timeout=CMD_TIMEOUT_SEC,
3542
)
3643
exit_status = process.returncode
@@ -39,11 +46,11 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False):
3946

4047
if expect_error:
4148
raise Exception(result, error)
42-
if exit_status != 0 or "error" in error.lower():
49+
if exit_status != 0 or "error" in error.lower().decode(encoding or 'utf-8'): # Decode error for comparison
4350
log.error(
44-
f"Problem in executing command: `{cmd}`\nerror: {error}\nexit_code: {exit_status}"
51+
f"Problem in executing command: `{cmd}`\nerror: {error.decode(encoding or 'utf-8')}\nexit_code: {exit_status}"
52+
# Decode for logging
4553
)
46-
exit(1)
4754

4855
if verbose:
4956
return exit_status, result, error
@@ -152,9 +159,9 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
152159
"""
153160
mode = "wb" if binary else "w"
154161
if not truncate:
155-
mode = "a" + mode
162+
mode = "ab" if binary else "a"
156163
if read_and_write:
157-
mode = "r+" + mode
164+
mode = "r+b" if binary else "r+"
158165

159166
with open(filename, mode) as file:
160167
if isinstance(data, list):
@@ -174,26 +181,26 @@ def touch(self, filename):
174181
with open(filename, "a"):
175182
os.utime(filename, None)
176183

177-
def read(self, filename):
178-
with open(filename, "r") as file:
184+
def read(self, filename, encoding=None):
185+
with open(filename, "r", encoding=encoding) as file:
179186
return file.read()
180187

181-
def readlines(self, filename, num_lines=0, encoding=None):
188+
def readlines(self, filename, num_lines=0, binary=False, encoding=None):
182189
"""
183190
Read lines from a local file.
184191
If num_lines is greater than 0, only the last num_lines lines will be read.
185192
"""
186193
assert num_lines >= 0
187-
194+
mode = 'rb' if binary else 'r'
188195
if num_lines == 0:
189-
with open(filename, "r", encoding=encoding) as file:
196+
with open(filename, mode, encoding=encoding) as file: # open in binary mode
190197
return file.readlines()
191198

192199
else:
193200
bufsize = 8192
194201
buffers = 1
195202

196-
with open(filename, "r", encoding=encoding) as file:
203+
with open(filename, mode, encoding=encoding) as file: # open in binary mode
197204
file.seek(0, os.SEEK_END)
198205
end_pos = file.tell()
199206

@@ -205,7 +212,7 @@ def readlines(self, filename, num_lines=0, encoding=None):
205212
cur_lines = len(lines)
206213

207214
if cur_lines >= num_lines or pos == 0:
208-
return lines[-num_lines:]
215+
return lines[-num_lines:] # get last num_lines from lines
209216

210217
buffers = int(
211218
buffers * max(2, int(num_lines / max(cur_lines, 1)))

‎testgres/os_ops/remote_ops.py

Lines changed: 39 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import io
12
import os
23
import tempfile
34
from contextlib import contextmanager
@@ -65,30 +66,41 @@ def connect(self):
6566
return ssh
6667

6768
# Command execution
68-
def exec_command(
69-
self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding="utf-8"
70-
):
69+
def exec_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):
7172
if isinstance(cmd, list):
7273
cmd = " ".join(cmd)
7374
log.debug(f"os_ops.exec_command: `{cmd}`; remote={self.remote}")
7475
# Source global profile file + execute command
7576
try:
7677
cmd = f"source /etc/profile.d/custom.sh; {cmd}"
7778
with self.ssh_connect() as ssh:
78-
stdin, stdout, stderr = ssh.exec_command(cmd)
79+
if input:
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)
7986
exit_status = 0
8087
if wait_exit:
8188
exit_status = stdout.channel.recv_exit_status()
82-
result = stdout.read().decode(encoding)
83-
error = stderr.read().decode(encoding)
89+
if encoding:
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()
8497

8598
if expect_error:
8699
raise Exception(result, error)
87-
if exit_status != 0 or "error" in error.lower():
100+
if exit_status != 0 or 'error' in error_str:
88101
log.error(
89102
f"Problem in executing command: `{cmd}`\nerror: {error}\nexit_code: {exit_status}"
90103
)
91-
exit(1)
92104

93105
if verbose:
94106
return exit_status, result, error
@@ -203,9 +215,9 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
203215
"""
204216
mode = "wb" if binary else "w"
205217
if not truncate:
206-
mode = "a" + mode
218+
mode = "ab" if binary else "a"
207219
if read_and_write:
208-
mode = "r+" + mode
220+
mode = "r+b" if binary else "r+"
209221

210222
with tempfile.NamedTemporaryFile(mode=mode) as tmp_file:
211223
if isinstance(data, list):
@@ -229,17 +241,28 @@ def touch(self, filename):
229241
"""
230242
self.exec_command(f"touch {filename}")
231243

232-
def read(self, filename, encoding="utf-8"):
244+
def read(self, filename, binary=False, encoding=None):
233245
cmd = f"cat {filename}"
234-
return self.exec_command(cmd, encoding=encoding)
246+
result = self.exec_command(cmd, encoding=encoding)
247+
248+
if not binary and result:
249+
result = result.decode(encoding or 'utf-8')
250+
251+
return result
235252

236-
def readlines(self, filename, num_lines=0, encoding=None):
237-
encoding = encoding or "utf-8"
253+
def readlines(self, filename, num_lines=0, binary=False, encoding=None):
238254
if num_lines > 0:
239255
cmd = f"tail -n {num_lines} {filename}"
240-
lines = self.exec_command(cmd, encoding)
241256
else:
242-
lines = self.read(filename, encoding=encoding).splitlines()
257+
cmd = f"cat {filename}"
258+
259+
result = self.exec_command(cmd, encoding=encoding)
260+
261+
if not binary and result:
262+
lines = result.decode(encoding or 'utf-8').splitlines()
263+
else:
264+
lines = result.splitlines()
265+
243266
return lines
244267

245268
def isfile(self, remote_file):

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy