Location via proxy:   [ UP ]  
[Report a bug]   [Manage cookies]                
Skip to content

[New] OsOps::execute_command supports a transfer of environment variables (exec_env) #239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions testgres/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -1020,7 +1020,7 @@ def get_control_data(self):

return out_dict

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

self.start(exec_env=exec_env)

if replica:
query = 'SELECT pg_is_in_recovery()'
Expand Down Expand Up @@ -1065,7 +1067,7 @@ def _detect_port_conflict(self, log_files0, log_files1):
return True
return False

def start(self, params=[], wait=True):
def start(self, params=[], wait=True, exec_env=None):
"""
Starts the PostgreSQL node using pg_ctl if node has not been started.
By default, it waits for the operation to complete before returning.
Expand All @@ -1079,7 +1081,7 @@ def start(self, params=[], wait=True):
Returns:
This instance of :class:`.PostgresNode`.
"""

assert exec_env is None or type(exec_env) == dict # noqa: E721
assert __class__._C_MAX_START_ATEMPTS > 1

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

def LOCAL__start_node():
# 'error' will be None on Windows
_, _, error = execute_utility2(self.os_ops, _params, self.utils_log_file, verbose=True)
_, _, error = execute_utility2(self.os_ops, _params, self.utils_log_file, verbose=True, exec_env=exec_env)
assert error is None or type(error) == str # noqa: E721
if error and 'does not exist' in error:
raise Exception(error)
Expand Down
66 changes: 60 additions & 6 deletions testgres/operations/local_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import socket

import psutil
import typing

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

def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None):
assert exec_env is None or type(exec_env) == dict # noqa: E721

# TODO: why don't we use the data from input?

extParams: typing.Dict[str, str] = dict()

if exec_env is None:
pass
elif len(exec_env) == 0:
pass
else:
env = os.environ.copy()
assert type(env) == dict # noqa: E721
for v in exec_env.items():
assert type(v) == tuple # noqa: E721
assert len(v) == 2
assert type(v[0]) == str # noqa: E721
assert v[0] != ""

if v[1] is None:
env.pop(v[0], None)
else:
assert type(v[1]) == str # noqa: E721
env[v[0]] = v[1]

extParams["env"] = env

with tempfile.NamedTemporaryFile(mode='w+b', delete=False) as temp_file:
stdout = temp_file
stderr = subprocess.STDOUT
Expand All @@ -58,6 +84,7 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process
stdin=stdin or subprocess.PIPE if input is not None else None,
stdout=stdout,
stderr=stderr,
**extParams,
)
if get_process:
return process, None, None
Expand All @@ -69,19 +96,45 @@ def _run_command__nt(self, cmd, shell, input, stdin, stdout, stderr, get_process
output, error = self._process_output(encoding, temp_file_path)
return process, output, error

def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding):
def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=None):
assert exec_env is None or type(exec_env) == dict # noqa: E721

input_prepared = None
if not get_process:
input_prepared = Helpers.PrepareProcessInput(input, encoding) # throw

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

extParams: typing.Dict[str, str] = dict()

if exec_env is None:
pass
elif len(exec_env) == 0:
pass
else:
env = os.environ.copy()
assert type(env) == dict # noqa: E721
for v in exec_env.items():
assert type(v) == tuple # noqa: E721
assert len(v) == 2
assert type(v[0]) == str # noqa: E721
assert v[0] != ""

if v[1] is None:
env.pop(v[0], None)
else:
assert type(v[1]) == str # noqa: E721
env[v[0]] = v[1]

extParams["env"] = env

process = subprocess.Popen(
cmd,
shell=shell,
stdin=stdin or subprocess.PIPE if input is not None else None,
stdout=stdout or subprocess.PIPE,
stderr=stderr or subprocess.PIPE,
**extParams
)
assert not (process is None)
if get_process:
Expand All @@ -100,25 +153,26 @@ def _run_command__generic(self, cmd, shell, input, stdin, stdout, stderr, get_pr
error = error.decode(encoding)
return process, output, error

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

return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
return method(self, cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env)

def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=False,
text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=False, timeout=None,
ignore_errors=False):
ignore_errors=False, exec_env=None):
"""
Execute a command in a subprocess and handle the output based on the provided parameters.
"""
assert type(expect_error) == bool # noqa: E721
assert type(ignore_errors) == bool # noqa: E721
assert exec_env is None or type(exec_env) == dict # noqa: E721

process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding, exec_env=exec_env)
if get_process:
return process

Expand Down
44 changes: 35 additions & 9 deletions testgres/operations/remote_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,14 +64,16 @@ def __enter__(self):

def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None,
stderr=None, get_process=None, timeout=None, ignore_errors=False):
stderr=None, get_process=None, timeout=None, ignore_errors=False,
exec_env=None):
"""
Execute a command in the SSH session.
Args:
- cmd (str): The command to be executed.
"""
assert type(expect_error) == bool # noqa: E721
assert type(ignore_errors) == bool # noqa: E721
assert exec_env is None or type(exec_env) == dict # noqa: E721

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

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

cmd_items = __class__._make_exec_env_list()
cmd_items = __class__._make_exec_env_list(exec_env=exec_env)
cmd_items.append(cmd_s)

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

@staticmethod
def _make_exec_env_list() -> typing.List[str]:
result: typing.List[str] = list()
def _make_exec_env_list(exec_env: typing.Dict) -> typing.List[str]:
env: typing.Dict[str, str] = dict()

# ---------------------------------- SYSTEM ENV
for envvar in os.environ.items():
if not __class__._does_put_envvar_into_exec_cmd(envvar[0]):
continue
qvalue = __class__._quote_envvar(envvar[1])
assert type(qvalue) == str # noqa: E721
result.append(envvar[0] + "=" + qvalue)
if __class__._does_put_envvar_into_exec_cmd(envvar[0]):
env[envvar[0]] = envvar[1]

# ---------------------------------- EXEC (LOCAL) ENV
if exec_env is None:
pass
else:
for envvar in exec_env.items():
assert type(envvar) == tuple # noqa: E721
assert len(envvar) == 2
assert type(envvar[0]) == str # noqa: E721
env[envvar[0]] = envvar[1]

# ---------------------------------- FINAL BUILD
result: typing.List[str] = list()
for envvar in env.items():
assert type(envvar) == tuple # noqa: E721
assert len(envvar) == 2
assert type(envvar[0]) == str # noqa: E721

if envvar[1] is None:
result.append("unset " + envvar[0])
else:
assert type(envvar[1]) == str # noqa: E721
qvalue = __class__._quote_envvar(envvar[1])
assert type(qvalue) == str # noqa: E721
result.append(envvar[0] + "=" + qvalue)
continue

return result
Expand Down
13 changes: 11 additions & 2 deletions testgres/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,17 +96,26 @@ def execute_utility(args, logfile=None, verbose=False):
return execute_utility2(tconf.os_ops, args, logfile, verbose)


def execute_utility2(os_ops: OsOperations, args, logfile=None, verbose=False, ignore_errors=False):
def execute_utility2(
os_ops: OsOperations,
args,
logfile=None,
verbose=False,
ignore_errors=False,
exec_env=None,
):
assert os_ops is not None
assert isinstance(os_ops, OsOperations)
assert type(verbose) == bool # noqa: E721
assert type(ignore_errors) == bool # noqa: E721
assert exec_env is None or type(exec_env) == dict # noqa: E721

exit_status, out, error = os_ops.exec_command(
args,
verbose=True,
ignore_errors=ignore_errors,
encoding=OsHelpers.GetDefaultEncoding())
encoding=OsHelpers.GetDefaultEncoding(),
exec_env=exec_env)

out = '' if not out else out

Expand Down
64 changes: 64 additions & 0 deletions tests/test_os_ops_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,70 @@ def test_exec_command_failure__expect_error(self, os_ops: OsOperations):
assert b"nonexistent_command" in error
assert b"not found" in error

def test_exec_command_with_exec_env(self, os_ops: OsOperations):
assert isinstance(os_ops, OsOperations)

RunConditions.skip_if_windows()

C_ENV_NAME = "TESTGRES_TEST__EXEC_ENV_20250414"

cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]

exec_env = {C_ENV_NAME: "Hello!"}

response = os_ops.exec_command(cmd, exec_env=exec_env)
assert response is not None
assert type(response) == bytes # noqa: E721
assert response == b'Hello!\n'

response = os_ops.exec_command(cmd)
assert response is not None
assert type(response) == bytes # noqa: E721
assert response == b'\n'

def test_exec_command__test_unset(self, os_ops: OsOperations):
assert isinstance(os_ops, OsOperations)

RunConditions.skip_if_windows()

C_ENV_NAME = "LANG"

cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]

response1 = os_ops.exec_command(cmd)
assert response1 is not None
assert type(response1) == bytes # noqa: E721

if response1 == b'\n':
logging.warning("Environment variable {} is not defined.".format(C_ENV_NAME))
return

exec_env = {C_ENV_NAME: None}
response2 = os_ops.exec_command(cmd, exec_env=exec_env)
assert response2 is not None
assert type(response2) == bytes # noqa: E721
assert response2 == b'\n'

response3 = os_ops.exec_command(cmd)
assert response3 is not None
assert type(response3) == bytes # noqa: E721
assert response3 == response1

def test_exec_command__test_unset_dummy_var(self, os_ops: OsOperations):
assert isinstance(os_ops, OsOperations)

RunConditions.skip_if_windows()

C_ENV_NAME = "TESTGRES_TEST__DUMMY_VAR_20250414"

cmd = ["sh", "-c", "echo ${}".format(C_ENV_NAME)]

exec_env = {C_ENV_NAME: None}
response2 = os_ops.exec_command(cmd, exec_env=exec_env)
assert response2 is not None
assert type(response2) == bytes # noqa: E721
assert response2 == b'\n'

def test_is_executable_true(self, os_ops: OsOperations):
"""
Test is_executable for an existing executable.
Expand Down