Skip to content

Commit ac782bb

Browse files
dmitry-lipetskdmitry-lipetsk
authored
[New] OsOps::execute_command supports a transfer of environment variables (exec_env) (#239)
* [New] OsOps::execute_command supports a transfer of environment variables (exec_env) New feature allows to pass environment variables to an executed program. If variable in exec_env has None value, then this variable will be unset. PostgresNode::start and PostgresNode::slow_start supports exec_env.
1 parent 2401474 commit ac782bb

File tree

5 files changed

+177
-22
lines changed

5 files changed

+177
-22
lines changed

‎testgres/node.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1020,7 +1020,7 @@ def get_control_data(self):
10201020

10211021
return out_dict
10221022

1023-
def slow_start(self, replica=False, dbname='template1', username=None, max_attempts=0):
1023+
def slow_start(self, replica=False, dbname='template1', username=None, max_attempts=0, exec_env=None):
10241024
"""
10251025
Starts the PostgreSQL instance and then polls the instance
10261026
until it reaches the expected state (primary or replica). The state is checked
@@ -1033,7 +1033,9 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem
10331033
If False, waits for the instance to be in primary mode. Default is False.
10341034
max_attempts:
10351035
"""
1036-
self.start()
1036+
assert exec_env is None or type(exec_env) == dict # noqa: E721
1037+
1038+
self.start(exec_env=exec_env)
10371039

10381040
if replica:
10391041
query = 'SELECT pg_is_in_recovery()'
@@ -1065,7 +1067,7 @@ def _detect_port_conflict(self, log_files0, log_files1):
10651067
return True
10661068
return False
10671069

1068-
def start(self, params=[], wait=True):
1070+
def start(self, params=[], wait=True, exec_env=None):
10691071
"""
10701072
Starts the PostgreSQL node using pg_ctl if node has not been started.
10711073
By default, it waits for the operation to complete before returning.
@@ -1079,7 +1081,7 @@ def start(self, params=[], wait=True):
10791081
Returns:
10801082
This instance of :class:`.PostgresNode`.
10811083
"""
1082-
1084+
assert exec_env is None or type(exec_env) == dict # noqa: E721
10831085
assert __class__._C_MAX_START_ATEMPTS > 1
10841086

10851087
if self.is_started:
@@ -1098,7 +1100,7 @@ def start(self, params=[], wait=True):
10981100

10991101
def LOCAL__start_node():
11001102
# 'error' will be None on Windows
1101-
_, _, error = execute_utility2(self.os_ops, _params, self.utils_log_file, verbose=True)
1103+
_, _, error = execute_utility2(self.os_ops, _params, self.utils_log_file, verbose=True, exec_env=exec_env)
11021104
assert error is None or type(error) == str # noqa: E721
11031105
if error and 'does not exist' in error:
11041106
raise Exception(error)

‎testgres/operations/local_ops.py

Lines changed: 60 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import socket
1010

1111
import psutil
12+
import typing
1213

1314
from ..exceptions import ExecUtilException
1415
from ..exceptions import InvalidOperationException
@@ -46,9 +47,34 @@ def _process_output(encoding, temp_file_path):
4647
output = output.decode(encoding)
4748
return output, None # In Windows stderr writing in stdout
4849

49-
def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
50+
def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None):
51+
assert exec_env is None or type(exec_env) == dict # noqa: E721
52+
5053
# TODO: why don't we use the data from input?
5154

55+
extParams: typing.Dict[str, str] = dict()
56+
57+
if exec_env is None:
58+
pass
59+
elif len(exec_env) == 0:
60+
pass
61+
else:
62+
env = os.environ.copy()
63+
assert type(env) == dict # noqa: E721
64+
for v in exec_env.items():
65+
assert type(v) == tuple # noqa: E721
66+
assert len(v) == 2
67+
assert type(v[0]) == str # noqa: E721
68+
assert v[0] != ""
69+
70+
if v[1] is None:
71+
env.pop(v[0], None)
72+
else:
73+
assert type(v[1]) == str # noqa: E721
74+
env[v[0]] = v[1]
75+
76+
extParams["env"] = env
77+
5278
with tempfile.NamedTemporaryFile(mode='w+b', delete=False) as temp_file:
5379
stdout = temp_file
5480
stderr = subprocess.STDOUT
@@ -58,6 +84,7 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process
5884
stdin=stdin or subprocess.PIPE if input is not None else None,
5985
stdout=stdout,
6086
stderr=stderr,
87+
**extParams,
6188
)
6289
if get_process:
6390
return process, None, None
@@ -69,19 +96,45 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process
6996
output, error = self._process_output(encoding, temp_file_path)
7097
return process, output, error
7198

72-
def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
99+
def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None):
100+
assert exec_env is None or type(exec_env) == dict # noqa: E721
101+
73102
input_prepared = None
74103
if not get_process:
75104
input_prepared = Helpers.PrepareProcessInput(input, encoding) # throw
76105

77106
assert input_prepared is None or (type(input_prepared) == bytes) # noqa: E721
78107

108+
extParams: typing.Dict[str, str] = dict()
109+
110+
if exec_env is None:
111+
pass
112+
elif len(exec_env) == 0:
113+
pass
114+
else:
115+
env = os.environ.copy()
116+
assert type(env) == dict # noqa: E721
117+
for v in exec_env.items():
118+
assert type(v) == tuple # noqa: E721
119+
assert len(v) == 2
120+
assert type(v[0]) == str # noqa: E721
121+
assert v[0] != ""
122+
123+
if v[1] is None:
124+
env.pop(v[0], None)
125+
else:
126+
assert type(v[1]) == str # noqa: E721
127+
env[v[0]] = v[1]
128+
129+
extParams["env"] = env
130+
79131
process = subprocess.Popen(
80132
cmd,
81133
shell=shell,
82134
stdin=stdin or subprocess.PIPE if input is not None else None,
83135
stdout=stdout or subprocess.PIPE,
84136
stderr=stderr or subprocess.PIPE,
137+
**extParams
85138
)
86139
assert not (process is None)
87140
if get_process:
@@ -100,25 +153,26 @@ def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_pr
100153
error = error.decode(encoding)
101154
return process, output, error
102155

103-
def _run_command(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
156+
def _run_command(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None):
104157
"""Execute a command and return the process and its output."""
105158
if os.name == 'nt' and stdout is None: # Windows
106159
method = __class__._run_command__nt
107160
else: # Other OS
108161
method = __class__._run_command__generic
109162

110-
return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
163+
return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env)
111164

112165
def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=False,
113166
text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None,
114-
ignore_errors=False):
167+
ignore_errors=False, exec_env=None):
115168
"""
116169
Execute a command in a subprocess and handle the output based on the provided parameters.
117170
"""
118171
assert type(expect_error) == bool # noqa: E721
119172
assert type(ignore_errors) == bool # noqa: E721
173+
assert exec_env is None or type(exec_env) == dict # noqa: E721
120174

121-
process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
175+
process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env)
122176
if get_process:
123177
return process
124178

‎testgres/operations/remote_ops.py

Lines changed: 35 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -64,14 +64,16 @@ def __enter__(self):
6464

6565
def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
6666
encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None,
67-
stderr=None, get_process=None, timeout=None, ignore_errors=False):
67+
stderr=None, get_process=None, timeout=None, ignore_errors=False,
68+
exec_env=None):
6869
"""
6970
Execute a command in the SSH session.
7071
Args:
7172
- cmd (str): The command to be executed.
7273
"""
7374
assert type(expect_error) == bool # noqa: E721
7475
assert type(ignore_errors) == bool # noqa: E721
76+
assert exec_env is None or type(exec_env) == dict # noqa: E721
7577

7678
input_prepared = None
7779
if not get_process:
@@ -88,7 +90,7 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
8890

8991
assert type(cmd_s) == str # noqa: E721
9092

91-
cmd_items = __class__._make_exec_env_list()
93+
cmd_items = __class__._make_exec_env_list(exec_env=exec_env)
9294
cmd_items.append(cmd_s)
9395

9496
env_cmd_s = ';'.join(cmd_items)
@@ -670,14 +672,38 @@ def _is_port_free__process_1(error: str) -> bool:
670672
return True
671673

672674
@staticmethod
673-
def _make_exec_env_list() -> typing.List[str]:
674-
result: typing.List[str] = list()
675+
def _make_exec_env_list(exec_env: typing.Dict) -> typing.List[str]:
676+
env: typing.Dict[str, str] = dict()
677+
678+
# ---------------------------------- SYSTEM ENV
675679
for envvar in os.environ.items():
676-
if not __class__._does_put_envvar_into_exec_cmd(envvar[0]):
677-
continue
678-
qvalue = __class__._quote_envvar(envvar[1])
679-
assert type(qvalue) == str # noqa: E721
680-
result.append(envvar[0] + "=" + qvalue)
680+
if __class__._does_put_envvar_into_exec_cmd(envvar[0]):
681+
env[envvar[0]] = envvar[1]
682+
683+
# ---------------------------------- EXEC (LOCAL) ENV
684+
if exec_env is None:
685+
pass
686+
else:
687+
for envvar in exec_env.items():
688+
assert type(envvar) == tuple # noqa: E721
689+
assert len(envvar) == 2
690+
assert type(envvar[0]) == str # noqa: E721
691+
env[envvar[0]] = envvar[1]
692+
693+
# ---------------------------------- FINAL BUILD
694+
result: typing.List[str] = list()
695+
for envvar in env.items():
696+
assert type(envvar) == tuple # noqa: E721
697+
assert len(envvar) == 2
698+
assert type(envvar[0]) == str # noqa: E721
699+
700+
if envvar[1] is None:
701+
result.append("unset " + envvar[0])
702+
else:
703+
assert type(envvar[1]) == str # noqa: E721
704+
qvalue = __class__._quote_envvar(envvar[1])
705+
assert type(qvalue) == str # noqa: E721
706+
result.append(envvar[0] + "=" + qvalue)
681707
continue
682708

683709
return result

‎testgres/utils.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,17 +96,26 @@ def execute_utility(args, logfile=None, verbose=False):
9696
return execute_utility2(tconf.os_ops, args, logfile, verbose)
9797

9898

99-
def execute_utility2(os_ops: OsOperations, args, logfile=None, verbose=False, ignore_errors=False):
99+
def execute_utility2(
100+
os_ops: OsOperations,
101+
args,
102+
logfile=None,
103+
verbose=False,
104+
ignore_errors=False,
105+
exec_env=None,
106+
):
100107
assert os_ops is not None
101108
assert isinstance(os_ops, OsOperations)
102109
assert type(verbose) == bool # noqa: E721
103110
assert type(ignore_errors) == bool # noqa: E721
111+
assert exec_env is None or type(exec_env) == dict # noqa: E721
104112

105113
exit_status, out, error = os_ops.exec_command(
106114
args,
107115
verbose=True,
108116
ignore_errors=ignore_errors,
109-
encoding=OsHelpers.GetDefaultEncoding())
117+
encoding=OsHelpers.GetDefaultEncoding(),
118+
exec_env=exec_env)
110119

111120
out = '' if not out else out
112121

‎tests/test_os_ops_common.py

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,70 @@ def test_exec_command_failure__expect_error(self, os_ops: OsOperations):
9393
assert b"nonexistent_command" in error
9494
assert b"not found" in error
9595

96+
def test_exec_command_with_exec_env(self, os_ops: OsOperations):
97+
assert isinstance(os_ops, OsOperations)
98+
99+
RunConditions.skip_if_windows()
100+
101+
C_ENV_NAME = "TESTGRES_TEST__EXEC_ENV_20250414"
102+
103+
cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]
104+
105+
exec_env = {C_ENV_NAME: "Hello!"}
106+
107+
response = os_ops.exec_command(cmd, exec_env=exec_env)
108+
assert response is not None
109+
assert type(response) == bytes # noqa: E721
110+
assert response == b'Hello!\n'
111+
112+
response = os_ops.exec_command(cmd)
113+
assert response is not None
114+
assert type(response) == bytes # noqa: E721
115+
assert response == b'\n'
116+
117+
def test_exec_command__test_unset(self, os_ops: OsOperations):
118+
assert isinstance(os_ops, OsOperations)
119+
120+
RunConditions.skip_if_windows()
121+
122+
C_ENV_NAME = "LANG"
123+
124+
cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]
125+
126+
response1 = os_ops.exec_command(cmd)
127+
assert response1 is not None
128+
assert type(response1) == bytes # noqa: E721
129+
130+
if response1 == b'\n':
131+
logging.warning("Environment variable {} is not defined.".format(C_ENV_NAME))
132+
return
133+
134+
exec_env = {C_ENV_NAME: None}
135+
response2 = os_ops.exec_command(cmd, exec_env=exec_env)
136+
assert response2 is not None
137+
assert type(response2) == bytes # noqa: E721
138+
assert response2 == b'\n'
139+
140+
response3 = os_ops.exec_command(cmd)
141+
assert response3 is not None
142+
assert type(response3) == bytes # noqa: E721
143+
assert response3 == response1
144+
145+
def test_exec_command__test_unset_dummy_var(self, os_ops: OsOperations):
146+
assert isinstance(os_ops, OsOperations)
147+
148+
RunConditions.skip_if_windows()
149+
150+
C_ENV_NAME = "TESTGRES_TEST__DUMMY_VAR_20250414"
151+
152+
cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]
153+
154+
exec_env = {C_ENV_NAME: None}
155+
response2 = os_ops.exec_command(cmd, exec_env=exec_env)
156+
assert response2 is not None
157+
assert type(response2) == bytes # noqa: E721
158+
assert response2 == b'\n'
159+
96160
def test_is_executable_true(self, os_ops: OsOperations):
97161
"""
98162
Test is_executable for an existing executable.

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