Merge pull request #75 from fsteggink/execwithenvvars

Introduced env_args and env_separator config settings to outputs.execoutput.ExecOutput
geopython · Jul 3, 2018 · c53f21a · c53f21a
2 parents 93188f4 + dbf3e28
commit c53f21a
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 15 deletions.
diff --git a/stetl/filters/execfilter.py b/stetl/filters/execfilter.py
@@ -7,6 +7,7 @@
 #
 import subprocess
 import os
+from stetl.component import Config
 from stetl.filter import Filter
 from stetl.util import Util
 from stetl.packet import FORMAT
@@ -19,21 +20,40 @@ class ExecFilter(Filter):
     Executes any command (abstract base class).
     """
 
+    @Config(ptype=str, default='', required=False)
+    def env_args(self):
+        """
+        Provides of list of environment variables which will be used when executing the given command.
+
+        Example: env_args = pgpassword=postgres othersetting=value~with~spaces
+        """
+        pass
+
+    @Config(ptype=str, default='=', required=False)
+    def env_separator(self):
+        """
+        Provides the separator to split the environment variable names from their values.
+        """
+        pass
+
     def __init__(self, configdict, section, consumes, produces):
         Filter.__init__(self, configdict, section, consumes, produces)
 
     def invoke(self, packet):
         return packet
 
     def execute_cmd(self, cmd):
-        use_shell = True
-        if os.name == 'nt':
-            use_shell = False
-
-        log.info("executing cmd=%s" % cmd)
-        result = subprocess.check_output(cmd, shell=use_shell)
-        log.info("execute done")
-        return result
+        env_vars = Util.string_to_dict(self.env_args, self.env_separator)
+        old_environ = os.environ.copy()
+
+        try:
+            os.environ.update(env_vars)
+            log.info("executing cmd=%s" % cmd)
+            result = subprocess.check_output(cmd, shell=True)
+            log.info("execute done")
+            return result
+        finally:
+            os.environ = old_environ
 
 
 class CommandExecFilter(ExecFilter):

diff --git a/stetl/outputs/execoutput.py b/stetl/outputs/execoutput.py
@@ -20,20 +20,39 @@ class ExecOutput(Output):
     Executes any command (abstract base class).
     """
 
+    @Config(ptype=str, default='', required=False)
+    def env_args(self):
+        """
+        Provides of list of environment variables which will be used when executing the given command.
+
+        Example: env_args = pgpassword=postgres othersetting=value~with~spaces
+        """
+        pass
+
+    @Config(ptype=str, default='=', required=False)
+    def env_separator(self):
+        """
+        Provides the separator to split the environment variable names from their values.
+        """
+        pass
+
     def __init__(self, configdict, section, consumes):
         Output.__init__(self, configdict, section, consumes)
 
     def write(self, packet):
         return packet
 
     def execute_cmd(self, cmd):
-        use_shell = True
-        if os.name == 'nt':
-            use_shell = False
-
-        log.info("executing cmd=%s" % cmd)
-        subprocess.call(cmd, shell=use_shell)
-        log.info("execute done")
+        env_vars = Util.string_to_dict(self.env_args, self.env_separator)
+        old_environ = os.environ.copy()
+
+        try:
+            os.environ.update(env_vars)
+            log.info("executing cmd=%s" % cmd)
+            subprocess.call(cmd, shell=True)
+            log.info("execute done")
+        finally:
+            os.environ = old_environ
 
 
 class CommandExecOutput(ExecOutput):

diff --git a/tests/filters/configs/commandexecfilter.cfg b/tests/filters/configs/commandexecfilter.cfg
@@ -9,6 +9,8 @@ file_path = tests/data/commandexecfilter.txt
 
 [command_executor]
 class = filters.execfilter.CommandExecFilter
+env_args = pgpassword!postgres
+env_separator = !
 
 [packet_buffer]
 class = filters.packetbuffer.PacketBuffer

diff --git a/tests/outputs/configs/commandexecoutput.cfg b/tests/outputs/configs/commandexecoutput.cfg
@@ -10,3 +10,5 @@ file_path = tests/data/command.txt
 [output_exec]
 class = outputs.execoutput.CommandExecOutput
 input_format = string
+env_args = pgpassword!postgres
+env_separator = !