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

dmitry-lipetsk · web-flow · commit ac782bb3354c · 2025-04-15T14:50:35.000+03:00
* [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.
diff --git a/testgres/node.py b/testgres/node.py
@@ -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
@@ -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()'
@@ -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.
@@ -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:
@@ -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)
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -9,6 +9,7 @@
 import socket
 
 import psutil
+import typing
 
 from ..exceptions import ExecUtilException
 from ..exceptions import InvalidOperationException
@@ -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
@@ -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
@@ -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:
@@ -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
 
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -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:
@@ -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)
@@ -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
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -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
 
diff --git a/tests/test_os_ops_common.py b/tests/test_os_ops_common.py
@@ -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.
