Skip to content

Commit edb5708

Browse files
author
vshepard
committed
Fix initdb error on Windows
1 parent 846c05f commit edb5708

File tree

6 files changed

+184
-51
lines changed

6 files changed

+184
-51
lines changed

‎setup.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,16 @@
2727
readme = f.read()
2828

2929
setup(
30-
version='1.9.2',
30+
version='1.9.3',
3131
name='testgres',
3232
packages=['testgres', 'testgres.operations'],
3333
description='Testing utility for PostgreSQL and its extensions',
3434
url='https://github.com/postgrespro/testgres',
3535
long_description=readme,
3636
long_description_content_type='text/markdown',
3737
license='PostgreSQL',
38-
author='Ildar Musin',
39-
author_email='zildermann@gmail.com',
38+
author='Postgres Professional',
39+
author_email='testgres@postgrespro.ru',
4040
keywords=['test', 'testing', 'postgresql'],
4141
install_requires=install_requires,
4242
classifiers=[],

‎testgres/operations/local_ops.py

Lines changed: 85 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,7 @@
88
import psutil
99

1010
from ..exceptions import ExecUtilException
11-
from .os_ops import ConnectionParams, OsOperations
12-
from .os_ops import pglib
11+
from .os_ops import ConnectionParams, OsOperations, pglib, get_default_encoding
1312

1413
try:
1514
from shutil import which as find_executable
@@ -22,6 +21,12 @@
2221
error_markers = [b'error', b'Permission denied', b'fatal']
2322

2423

24+
def has_errors(output):
25+
if isinstance(output, str):
26+
output = output.encode(get_default_encoding())
27+
return any(marker in output for marker in error_markers)
28+
29+
2530
class LocalOperations(OsOperations):
2631
def __init__(self, conn_params=None):
2732
if conn_params is None:
@@ -33,7 +38,38 @@ def __init__(self, conn_params=None):
3338
self.remote = False
3439
self.username = conn_params.username or self.get_user()
3540

36-
# Command execution
41+
@staticmethod
42+
def _run_command(cmd, shell, input, timeout, encoding, temp_file=None):
43+
"""Execute a command and return the process."""
44+
if temp_file is not None:
45+
stdout = temp_file
46+
stderr = subprocess.STDOUT
47+
else:
48+
stdout = subprocess.PIPE
49+
stderr = subprocess.PIPE
50+
51+
process = subprocess.Popen(
52+
cmd,
53+
shell=shell,
54+
stdin=subprocess.PIPE if input is not None else None,
55+
stdout=stdout,
56+
stderr=stderr,
57+
)
58+
59+
try:
60+
return process.communicate(input=input.encode(encoding) if input else None, timeout=timeout), process
61+
except subprocess.TimeoutExpired:
62+
process.kill()
63+
raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
64+
65+
@staticmethod
66+
def _raise_exec_exception(message, command, exit_code, output):
67+
"""Raise an ExecUtilException."""
68+
raise ExecUtilException(message=message.format(output),
69+
command=command,
70+
exit_code=exit_code,
71+
out=output)
72+
3773
def exec_command(self, cmd, wait_exit=False, verbose=False,
3874
expect_error=False, encoding=None, shell=False, text=False,
3975
input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
@@ -56,16 +92,15 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
5692
:return: The output of the subprocess.
5793
"""
5894
if os.name == 'nt':
59-
with tempfile.NamedTemporaryFile() as buf:
60-
process = subprocess.Popen(cmd, stdout=buf, stderr=subprocess.STDOUT)
61-
process.communicate()
62-
buf.seek(0)
63-
result = buf.read().decode(encoding)
64-
return result
95+
return self._exec_command_windows(cmd, wait_exit=wait_exit, verbose=verbose,
96+
expect_error=expect_error, encoding=encoding, shell=shell, text=text,
97+
input=input, stdin=stdin, stdout=stdout, stderr=stderr,
98+
get_process=get_process, timeout=timeout)
6599
else:
66100
process = subprocess.Popen(
67101
cmd,
68102
shell=shell,
103+
stdin=stdin,
69104
stdout=stdout,
70105
stderr=stderr,
71106
)
@@ -79,7 +114,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
79114
raise ExecUtilException("Command timed out after {} seconds.".format(timeout))
80115
exit_status = process.returncode
81116

82-
error_found = exit_status != 0 or any(marker in error for marker in error_markers)
117+
error_found = exit_status != 0 or has_errors(error)
83118

84119
if encoding:
85120
result = result.decode(encoding)
@@ -91,15 +126,50 @@ def exec_command(self, cmd, wait_exit=False, verbose=False,
91126
if exit_status != 0 or error_found:
92127
if exit_status == 0:
93128
exit_status = 1
94-
raise ExecUtilException(message='Utility exited with non-zero code. Error `{}`'.format(error),
95-
command=cmd,
96-
exit_code=exit_status,
97-
out=result)
129+
self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, exit_status, result)
98130
if verbose:
99131
return exit_status, result, error
100132
else:
101133
return result
102134

135+
@staticmethod
136+
def _process_output(process, encoding, temp_file=None):
137+
"""Process the output of a command."""
138+
if temp_file is not None:
139+
temp_file.seek(0)
140+
output = temp_file.read()
141+
else:
142+
output = process.stdout.read()
143+
144+
if encoding:
145+
output = output.decode(encoding)
146+
147+
return output
148+
149+
def _exec_command_windows(self, cmd, wait_exit=False, verbose=False,
150+
expect_error=False, encoding=None, shell=False, text=False,
151+
input=None, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
152+
get_process=None, timeout=None):
153+
with tempfile.NamedTemporaryFile(mode='w+b') as temp_file:
154+
_, process = self._run_command(cmd, shell, input, timeout, encoding, temp_file)
155+
if get_process:
156+
return process
157+
output = self._process_output(process, encoding, temp_file)
158+
159+
if process.returncode != 0 or has_errors(output):
160+
if process.returncode == 0:
161+
process.returncode = 1
162+
if expect_error:
163+
if verbose:
164+
return process.returncode, output, output
165+
else:
166+
return output
167+
else:
168+
self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode,
169+
output)
170+
171+
return (process.returncode, output, output) if verbose else output
172+
103173
# Environment setup
104174
def environ(self, var_name):
105175
return os.environ.get(var_name)
@@ -210,7 +280,7 @@ def read(self, filename, encoding=None, binary=False):
210280
if binary:
211281
return content
212282
if isinstance(content, bytes):
213-
return content.decode(encoding or 'utf-8')
283+
return content.decode(encoding or get_default_encoding())
214284
return content
215285

216286
def readlines(self, filename, num_lines=0, binary=False, encoding=None):

‎testgres/operations/os_ops.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import locale
2+
13
try:
24
import psycopg2 as pglib # noqa: F401
35
except ImportError:
@@ -14,6 +16,10 @@ def __init__(self, host='127.0.0.1', ssh_key=None, username=None):
1416
self.username = username
1517

1618

19+
def get_default_encoding():
20+
return locale.getdefaultlocale()[1] or 'UTF-8'
21+
22+
1723
class OsOperations:
1824
def __init__(self, username=None):
1925
self.ssh_key = None

‎testgres/operations/remote_ops.py

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import locale
21
import logging
32
import os
43
import subprocess
@@ -15,12 +14,7 @@
1514
raise ImportError("You must have psycopg2 or pg8000 modules installed")
1615

1716
from ..exceptions import ExecUtilException
18-
19-
from .os_ops import OsOperations, ConnectionParams
20-
21-
ConsoleEncoding = locale.getdefaultlocale()[1]
22-
if not ConsoleEncoding:
23-
ConsoleEncoding = 'UTF-8'
17+
from .os_ops import OsOperations, ConnectionParams, get_default_encoding
2418

2519
error_markers = [b'error', b'Permission denied', b'fatal', b'No such file or directory']
2620

@@ -36,7 +30,7 @@ def kill(self):
3630

3731
def cmdline(self):
3832
command = "ps -p {} -o cmd --no-headers".format(self.pid)
39-
stdin, stdout, stderr = self.ssh.exec_command(command, verbose=True, encoding=ConsoleEncoding)
33+
stdin, stdout, stderr = self.ssh.exec_command(command, verbose=True, encoding=get_default_encoding())
4034
cmdline = stdout.strip()
4135
return cmdline.split()
4236

@@ -145,7 +139,7 @@ def environ(self, var_name: str) -> str:
145139
- var_name (str): The name of the environment variable.
146140
"""
147141
cmd = "echo ${}".format(var_name)
148-
return self.exec_command(cmd, encoding=ConsoleEncoding).strip()
142+
return self.exec_command(cmd, encoding=get_default_encoding()).strip()
149143

150144
def find_executable(self, executable):
151145
search_paths = self.environ("PATH")
@@ -176,11 +170,11 @@ def set_env(self, var_name: str, var_val: str):
176170

177171
# Get environment variables
178172
def get_user(self):
179-
return self.exec_command("echo $USER", encoding=ConsoleEncoding).strip()
173+
return self.exec_command("echo $USER", encoding=get_default_encoding()).strip()
180174

181175
def get_name(self):
182176
cmd = 'python3 -c "import os; print(os.name)"'
183-
return self.exec_command(cmd, encoding=ConsoleEncoding).strip()
177+
return self.exec_command(cmd, encoding=get_default_encoding()).strip()
184178

185179
# Work with dirs
186180
def makedirs(self, path, remove_existing=False):
@@ -227,7 +221,7 @@ def listdir(self, path):
227221
return result.splitlines()
228222

229223
def path_exists(self, path):
230-
result = self.exec_command("test -e {}; echo $?".format(path), encoding=ConsoleEncoding)
224+
result = self.exec_command("test -e {}; echo $?".format(path), encoding=get_default_encoding())
231225
return int(result.strip()) == 0
232226

233227
@property
@@ -264,9 +258,9 @@ def mkdtemp(self, prefix=None):
264258

265259
def mkstemp(self, prefix=None):
266260
if prefix:
267-
temp_dir = self.exec_command("mktemp {}XXXXX".format(prefix), encoding=ConsoleEncoding)
261+
temp_dir = self.exec_command("mktemp {}XXXXX".format(prefix), encoding=get_default_encoding())
268262
else:
269-
temp_dir = self.exec_command("mktemp", encoding=ConsoleEncoding)
263+
temp_dir = self.exec_command("mktemp", encoding=get_default_encoding())
270264

271265
if temp_dir:
272266
if not os.path.isabs(temp_dir):
@@ -283,7 +277,9 @@ def copytree(self, src, dst):
283277
return self.exec_command("cp -r {} {}".format(src, dst))
284278

285279
# Work with files
286-
def write(self, filename, data, truncate=False, binary=False, read_and_write=False, encoding=ConsoleEncoding):
280+
def write(self, filename, data, truncate=False, binary=False, read_and_write=False, encoding=None):
281+
if not encoding:
282+
encoding = get_default_encoding()
287283
mode = "wb" if binary else "w"
288284
if not truncate:
289285
mode = "ab" if binary else "a"
@@ -302,7 +298,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
302298
data = data.encode(encoding)
303299

304300
if isinstance(data, list):
305-
data = [(s if isinstance(s, str) else s.decode(ConsoleEncoding)).rstrip('\n') + '\n' for s in data]
301+
data = [(s if isinstance(s, str) else s.decode(get_default_encoding())).rstrip('\n') + '\n' for s in data]
306302
tmp_file.writelines(data)
307303
else:
308304
tmp_file.write(data)
@@ -334,7 +330,7 @@ def read(self, filename, binary=False, encoding=None):
334330
result = self.exec_command(cmd, encoding=encoding)
335331

336332
if not binary and result:
337-
result = result.decode(encoding or ConsoleEncoding)
333+
result = result.decode(encoding or get_default_encoding())
338334

339335
return result
340336

@@ -347,7 +343,7 @@ def readlines(self, filename, num_lines=0, binary=False, encoding=None):
347343
result = self.exec_command(cmd, encoding=encoding)
348344

349345
if not binary and result:
350-
lines = result.decode(encoding or ConsoleEncoding).splitlines()
346+
lines = result.decode(encoding or get_default_encoding()).splitlines()
351347
else:
352348
lines = result.splitlines()
353349

@@ -375,7 +371,7 @@ def kill(self, pid, signal):
375371

376372
def get_pid(self):
377373
# Get current process id
378-
return int(self.exec_command("echo $$", encoding=ConsoleEncoding))
374+
return int(self.exec_command("echo $$", encoding=get_default_encoding()))
379375

380376
def get_process_children(self, pid):
381377
command = ["ssh", "-i", self.ssh_key, f"{self.username}@{self.host}", f"pgrep -P {pid}"]

‎testgres/utils.py

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,16 @@
44
from __future__ import print_function
55

66
import os
7-
import port_for
7+
import random
8+
import socket
9+
810
import sys
911

1012
from contextlib import contextmanager
1113
from packaging.version import Version, InvalidVersion
1214
import re
1315

16+
from port_for import PortForException
1417
from six import iteritems
1518

1619
from .exceptions import ExecUtilException
@@ -37,13 +40,49 @@ def reserve_port():
3740
"""
3841
Generate a new port and add it to 'bound_ports'.
3942
"""
40-
41-
port = port_for.select_random(exclude_ports=bound_ports)
43+
port = select_random(exclude_ports=bound_ports)
4244
bound_ports.add(port)
4345

4446
return port
4547

4648

49+
def select_random(
50+
ports=None,
51+
exclude_ports=None,
52+
) -> int:
53+
"""
54+
Return random unused port number.
55+
Standard function from port_for does not work on Windows because of error
56+
'port_for.exceptions.PortForException: Can't select a port'
57+
We should update it.
58+
"""
59+
if ports is None:
60+
ports = set(range(1024, 65535))
61+
62+
if exclude_ports is None:
63+
exclude_ports = set()
64+
65+
ports.difference_update(set(exclude_ports))
66+
67+
sampled_ports = random.sample(tuple(ports), min(len(ports), 100))
68+
69+
for port in sampled_ports:
70+
if is_port_free(port):
71+
return port
72+
73+
raise PortForException("Can't select a port")
74+
75+
76+
def is_port_free(port: int) -> bool:
77+
"""Check if a port is free to use."""
78+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
79+
try:
80+
s.bind(("", port))
81+
return True
82+
except OSError:
83+
return False
84+
85+
4786
def release_port(port):
4887
"""
4988
Free port provided by reserve_port().
@@ -80,7 +119,8 @@ def execute_utility(args, logfile=None, verbose=False):
80119
lines = [u'\n'] + ['# ' + line for line in out.splitlines()] + [u'\n']
81120
tconf.os_ops.write(filename=logfile, data=lines)
82121
except IOError:
83-
raise ExecUtilException("Problem with writing to logfile `{}` during run command `{}`".format(logfile, args))
122+
raise ExecUtilException(
123+
"Problem with writing to logfile `{}` during run command `{}`".format(logfile, args))
84124
if verbose:
85125
return exit_status, out, error
86126
else:

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