From cf16ea6ecd6f5537cb7f2a14cb7c9a50be58248e Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 12:07:58 -0400
Subject: [PATCH 01/27] checkpoint

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 73 ++++++++++++++++++++++++++++++++++
 1 file changed, 73 insertions(+)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 70990647f1..56e2a82eac 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -237,3 +237,76 @@ def execute(self, **kwargs) -> typing.Any:
 
     def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
         return self._config_task_instance.post_execute(user_params, rval)
+
+
+class PortableShellTask(ShellTask):
+
+    def __init__(
+        self,
+        name: str,
+        script_file: str,
+        debug: bool = False,
+        task_config: T = None,
+        inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
+        output_locs: typing.Optional[typing.List[OutputLocation]] = None,
+        **kwargs,
+    ):
+        self._script = """
+        #!/bin/bash
+        
+        set -uexo pipefail
+        
+        {inputs.export_env}
+        
+        bash {inputs.script_file} {inputs.args}
+        """
+
+        def execute(self, **kwargs) -> typing.Any:
+            """
+            Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem
+            """
+            logger.info(f"Running shell script as type {self.task_type}")
+            if self.script_file:
+                with open(self.script_file) as f:
+                    self._script = f.read()
+
+            outputs: typing.Dict[str, str] = {}
+            if self._output_locs:
+                for v in self._output_locs:
+                    outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs)
+
+            if os.name == "nt":
+                self._script = self._script.lstrip().rstrip().replace("\n", "&&")
+
+            gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
+            if self._debug:
+                print("\n==============================================\n")
+                print(gen_script)
+                print("\n==============================================\n")
+
+            try:
+                subprocess.check_call(gen_script, shell=True)
+            except subprocess.CalledProcessError as e:
+                files = os.listdir(".")
+                fstr = "\n-".join(files)
+                logger.error(
+                    f"Failed to Execute Script, return-code {e.returncode} \n"
+                    f"StdErr: {e.stderr}\n"
+                    f"StdOut: {e.stdout}\n"
+                    f" Current directory contents: .\n-{fstr}"
+                )
+                raise
+
+            final_outputs = []
+            for v in self._output_locs:
+                if issubclass(v.var_type, FlyteFile):
+                    final_outputs.append(FlyteFile(outputs[v.var]))
+                if issubclass(v.var_type, FlyteDirectory):
+                    final_outputs.append(FlyteDirectory(outputs[v.var]))
+            if len(final_outputs) == 1:
+                return final_outputs[0]
+            if len(final_outputs) > 1:
+                return tuple(final_outputs)
+            return None
+
+

From eff0a7dc0e4dbd59540257e876c1256d44325bf1 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 13:05:01 -0400
Subject: [PATCH 02/27] Experimental implementation works well. Instead of
 messing with the class, we utilize the interpolation and context to
 dynamically generate a wrapper around the desired script. In the wrapper, we
 cd to the ctx.working_directory, export the env variables (handled by an
 additional method to convert dict to str), and then pass the arguments to the
 script

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 95 ++++++++++------------------------
 1 file changed, 28 insertions(+), 67 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 56e2a82eac..014e7a2576 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -184,6 +184,12 @@ def script(self) -> typing.Optional[str]:
     def script_file(self) -> typing.Optional[os.PathLike]:
         return self._script_file
 
+    def make_export_string_from_env_dict(self, d):
+        items = []
+        for k, v in d.items():
+            items.append(f"export {k}={v}")
+        return "\n".join(items)
+
     def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
         return self._config_task_instance.pre_execute(user_params)
 
@@ -204,6 +210,9 @@ def execute(self, **kwargs) -> typing.Any:
         if os.name == "nt":
             self._script = self._script.lstrip().rstrip().replace("\n", "&&")
 
+        if "env" in kwargs:
+            kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"])
+
         gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
         if self._debug:
             print("\n==============================================\n")
@@ -239,74 +248,26 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty
         return self._config_task_instance.post_execute(user_params, rval)
 
 
-class PortableShellTask(ShellTask):
-
-    def __init__(
-        self,
-        name: str,
-        script_file: str,
-        debug: bool = False,
-        task_config: T = None,
-        inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
-        output_locs: typing.Optional[typing.List[OutputLocation]] = None,
-        **kwargs,
-    ):
-        self._script = """
-        #!/bin/bash
-        
-        set -uexo pipefail
-        
-        {inputs.export_env}
-        
-        bash {inputs.script_file} {inputs.args}
-        """
+portable_shell_task = ShellTask(
+    name="portable_shell_task_instance",
+    debug=True,
+    inputs=kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
+    output_locs=[
+        OutputLocation(
+            var="k",
+            var_type=FlyteFile,
+            location="{ctx.working_directory}/test_output.txt",
+        )
+    ],
+    script="""
+    #!/bin/bash
 
-        def execute(self, **kwargs) -> typing.Any:
-            """
-            Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem
-            """
-            logger.info(f"Running shell script as type {self.task_type}")
-            if self.script_file:
-                with open(self.script_file) as f:
-                    self._script = f.read()
-
-            outputs: typing.Dict[str, str] = {}
-            if self._output_locs:
-                for v in self._output_locs:
-                    outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs)
-
-            if os.name == "nt":
-                self._script = self._script.lstrip().rstrip().replace("\n", "&&")
-
-            gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
-            if self._debug:
-                print("\n==============================================\n")
-                print(gen_script)
-                print("\n==============================================\n")
-
-            try:
-                subprocess.check_call(gen_script, shell=True)
-            except subprocess.CalledProcessError as e:
-                files = os.listdir(".")
-                fstr = "\n-".join(files)
-                logger.error(
-                    f"Failed to Execute Script, return-code {e.returncode} \n"
-                    f"StdErr: {e.stderr}\n"
-                    f"StdOut: {e.stdout}\n"
-                    f" Current directory contents: .\n-{fstr}"
-                )
-                raise
+    set -uexo pipefail
 
-            final_outputs = []
-            for v in self._output_locs:
-                if issubclass(v.var_type, FlyteFile):
-                    final_outputs.append(FlyteFile(outputs[v.var]))
-                if issubclass(v.var_type, FlyteDirectory):
-                    final_outputs.append(FlyteDirectory(outputs[v.var]))
-            if len(final_outputs) == 1:
-                return final_outputs[0]
-            if len(final_outputs) > 1:
-                return tuple(final_outputs)
-            return None
+    cd {ctx.working_directory}
 
+    {inputs.export_env}
 
+    bash {inputs.script_file} {inputs.script_args}
+    """
+)
\ No newline at end of file

From afa9c5f482faec37ef8228071dc65c22079767ff Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 13:08:26 -0400
Subject: [PATCH 03/27] fix spacing

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 014e7a2576..d424004315 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -212,7 +212,7 @@ def execute(self, **kwargs) -> typing.Any:
 
         if "env" in kwargs:
             kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"])
-
+        breakpoint()
         gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
         if self._debug:
             print("\n==============================================\n")
@@ -251,7 +251,7 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty
 portable_shell_task = ShellTask(
     name="portable_shell_task_instance",
     debug=True,
-    inputs=kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
+    inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
     output_locs=[
         OutputLocation(
             var="k",
@@ -260,14 +260,14 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty
         )
     ],
     script="""
-    #!/bin/bash
+#!/bin/bash
 
-    set -uexo pipefail
+set -uexo pipefail
 
-    cd {ctx.working_directory}
+cd {ctx.working_directory}
 
-    {inputs.export_env}
+{inputs.export_env}
 
-    bash {inputs.script_file} {inputs.script_args}
+bash {inputs.script_file} {inputs.script_args}
     """
 )
\ No newline at end of file

From 369eca0ed05ba1323a17aaed6edcfe55eef3c338 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 13:10:51 -0400
Subject: [PATCH 04/27] remove breakpoint

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index d424004315..79d3927374 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -212,7 +212,7 @@ def execute(self, **kwargs) -> typing.Any:
 
         if "env" in kwargs:
             kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"])
-        breakpoint()
+
         gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
         if self._debug:
             print("\n==============================================\n")

From a29367e6b1c746c23f6cd4bcde3a0c7c57ee2ea9 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 13:23:10 -0400
Subject: [PATCH 05/27] Output ctx.working_directory as single output

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 79d3927374..a98d2b2b1a 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -255,8 +255,8 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty
     output_locs=[
         OutputLocation(
             var="k",
-            var_type=FlyteFile,
-            location="{ctx.working_directory}/test_output.txt",
+            var_type=FlyteDirectory,
+            location="{ctx.working_directory}",
         )
     ],
     script="""

From a64774b5f35305f0a5f9d6bab4d282a46843c631 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 13:39:41 -0400
Subject: [PATCH 06/27] more doc strings

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index a98d2b2b1a..a18b44f882 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -184,7 +184,15 @@ def script(self) -> typing.Optional[str]:
     def script_file(self) -> typing.Optional[os.PathLike]:
         return self._script_file
 
-    def make_export_string_from_env_dict(self, d):
+    def make_export_string_from_env_dict(self, d) -> str:
+        """
+        Utility function to convert a dictionary of desired environment variable key: value pairs into a string of
+        `
+        export k1=v1
+        export k2=v2
+        ...
+        `
+        """
         items = []
         for k, v in d.items():
             items.append(f"export {k}={v}")
@@ -210,7 +218,7 @@ def execute(self, **kwargs) -> typing.Any:
         if os.name == "nt":
             self._script = self._script.lstrip().rstrip().replace("\n", "&&")
 
-        if "env" in kwargs:
+        if "env" in kwargs and isinstance(kwargs["env"], dict):
             kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"])
 
         gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)

From 78292811334c241245d7d803fcde2d476467745e Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Thu, 31 Mar 2022 13:42:57 -0400
Subject: [PATCH 07/27] more comments

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index a18b44f882..452c9af10c 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -255,7 +255,9 @@ def execute(self, **kwargs) -> typing.Any:
     def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
         return self._config_task_instance.post_execute(user_params, rval)
 
-
+# The portable_shell_task is an instance of ShellTask which wraps a 'pure' shell script
+# This utility function allows for the specification of env variables, arguments, and the actual script within the
+# workflow definition rather than at `ShellTask` instantiation
 portable_shell_task = ShellTask(
     name="portable_shell_task_instance",
     debug=True,

From 15453aa93202e54039d5c30336ea70a5536be6c4 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Fri, 1 Apr 2022 11:01:17 -0400
Subject: [PATCH 08/27] Added comments

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 452c9af10c..7edb5dd6c7 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -219,6 +219,8 @@ def execute(self, **kwargs) -> typing.Any:
             self._script = self._script.lstrip().rstrip().replace("\n", "&&")
 
         if "env" in kwargs and isinstance(kwargs["env"], dict):
+            # This supports the portable_shell_task by adding an additional key:value pair to kwargs/input
+            # This will cause collisions if a user tries to use `env` AND `export_env` in their inputs
             kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"])
 
         gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)

From 0abcfdc3565759e99e369c2c9a5de709a0e81ba5 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Fri, 1 Apr 2022 15:40:17 -0400
Subject: [PATCH 09/27] Add tests and test files from other branch

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 .../flytekit/unit/extras/tasks/test_shell.py  | 36 ++++++++++++++++++-
 .../extras/tasks/testdata/script_args_env.sh  | 22 ++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index ad1d75e272..8ded46048b 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -9,7 +9,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask, portable_shell_task
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -20,6 +20,7 @@
     script_sh = os.path.join(testdata, "script.exe")
 else:
     script_sh = os.path.join(testdata, "script.sh")
+    script_sh_2 = os.path.join(testdata, "script_args_env.sh")
 
 
 def test_shell_task_no_io():
@@ -249,3 +250,36 @@ def test_shell_script():
 
     assert t.script_file == script_sh
     t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+
+
+def test_portable_shell_task_with_args(capfd):
+    portable_shell_task(
+        script_file=script_sh_2,
+        script_args="first_arg second_arg",
+        env={}
+    )
+    cap = capfd.readouterr()
+    assert "first_arg" in cap.out
+    assert "second_arg" in cap.out
+
+
+def test_shell_task_with_env(capfd):
+    portable_shell_task(
+        script_file=script_sh_2,
+        env={"A": "AAAA", "B": "BBBB"},
+        script_args=""
+    )
+    cap = capfd.readouterr()
+    assert "AAAA" in cap.out
+    assert "BBBB" in cap.out
+
+
+def test_shell_task_properly_restores_env_after_execution():
+    env_as_dict = os.environ.copy()
+    portable_shell_task(
+        script_file=script_sh_2,
+        env={"A": "AAAA", "B": "BBBB"},
+        script_args=""
+    )
+    env_as_dict_after = os.environ.copy()
+    assert env_as_dict == env_as_dict_after
\ No newline at end of file
diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
new file mode 100644
index 0000000000..b63d4c0a29
--- /dev/null
+++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+set -exo pipefail
+
+echo "A: $A"
+echo "B: $B"
+
+if [ ! -z $1 ]; then
+    echo $1
+else
+    echo "Unset first positional argument"
+fi
+
+if [ ! -z $2 ]; then
+    echo $2
+else
+    echo "Unset second positional argument"
+fi
+
+SOME_VAR="This var is set"
+
+echo "Reading SOME_VAR: ${SOME_VAR}"
\ No newline at end of file

From e0efff65e093bc9e6c5ac35f44cadd7127adc5b3 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Wed, 6 Apr 2022 15:28:33 -0400
Subject: [PATCH 10/27] minor fix, have function return the instance rather
 than create. It seems to get registered when flyte packages your project

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 7edb5dd6c7..653fbdb25e 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -257,10 +257,12 @@ def execute(self, **kwargs) -> typing.Any:
     def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
         return self._config_task_instance.post_execute(user_params, rval)
 
+
 # The portable_shell_task is an instance of ShellTask which wraps a 'pure' shell script
 # This utility function allows for the specification of env variables, arguments, and the actual script within the
 # workflow definition rather than at `ShellTask` instantiation
-portable_shell_task = ShellTask(
+def get_portable_shell_task() -> ShellTask:
+    return ShellTask(
     name="portable_shell_task_instance",
     debug=True,
     inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),

From 4abd44707a34c15845513e59f256c71f8beb272e Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Mon, 11 Apr 2022 13:00:44 -0400
Subject: [PATCH 11/27] fix tests

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 tests/flytekit/unit/extras/tasks/test_shell.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 8ded46048b..1f443bf47b 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -9,7 +9,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask, portable_shell_task
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_portable_shell_task
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -253,7 +253,8 @@ def test_shell_script():
 
 
 def test_portable_shell_task_with_args(capfd):
-    portable_shell_task(
+    pst = get_portable_shell_task()
+    pst(
         script_file=script_sh_2,
         script_args="first_arg second_arg",
         env={}
@@ -264,7 +265,8 @@ def test_portable_shell_task_with_args(capfd):
 
 
 def test_shell_task_with_env(capfd):
-    portable_shell_task(
+    pst = get_portable_shell_task()
+    pst(
         script_file=script_sh_2,
         env={"A": "AAAA", "B": "BBBB"},
         script_args=""
@@ -276,7 +278,8 @@ def test_shell_task_with_env(capfd):
 
 def test_shell_task_properly_restores_env_after_execution():
     env_as_dict = os.environ.copy()
-    portable_shell_task(
+    pst = get_portable_shell_task()
+    pst(
         script_file=script_sh_2,
         env={"A": "AAAA", "B": "BBBB"},
         script_args=""

From 6e2964b5075b2e1684daf6f222a5f950a1d3ffae Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Tue, 12 Apr 2022 09:12:32 -0400
Subject: [PATCH 12/27] remove set flags not supported by sh

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 653fbdb25e..1701ff5ae2 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -276,7 +276,7 @@ def get_portable_shell_task() -> ShellTask:
     script="""
 #!/bin/bash
 
-set -uexo pipefail
+set -uex
 
 cd {ctx.working_directory}
 

From d16b54ed20d7701051ea8da0b16bb196b76d6d88 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Tue, 12 Apr 2022 18:57:01 -0400
Subject: [PATCH 13/27] fix spellcheck lint errors

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 .../unit/extras/tasks/testdata/script_args_env.sh         | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
index b63d4c0a29..dd77131e22 100644
--- a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
+++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
@@ -5,14 +5,14 @@ set -exo pipefail
 echo "A: $A"
 echo "B: $B"
 
-if [ ! -z $1 ]; then
-    echo $1
+if [ ! -n "$1" ]; then
+    echo "$1"
 else
     echo "Unset first positional argument"
 fi
 
-if [ ! -z $2 ]; then
-    echo $2
+if [ ! -n "$2" ]; then
+    echo "$2"
 else
     echo "Unset second positional argument"
 fi

From 9a989de47d051a27d4a9aa358af5b7528a6b02fd Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Fri, 15 Apr 2022 12:42:49 -0400
Subject: [PATCH 14/27] don't have a windows equivalent test script so
 bypassing those tests for now

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 tests/flytekit/unit/extras/tasks/test_shell.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 1f443bf47b..2eafe90468 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -18,6 +18,7 @@
 test_csv = os.path.join(testdata, "test.csv")
 if os.name == "nt":
     script_sh = os.path.join(testdata, "script.exe")
+    script_sh_2 = None
 else:
     script_sh = os.path.join(testdata, "script.sh")
     script_sh_2 = os.path.join(testdata, "script_args_env.sh")
@@ -253,6 +254,8 @@ def test_shell_script():
 
 
 def test_portable_shell_task_with_args(capfd):
+    if script_sh_2 is None:
+        return
     pst = get_portable_shell_task()
     pst(
         script_file=script_sh_2,
@@ -265,6 +268,8 @@ def test_portable_shell_task_with_args(capfd):
 
 
 def test_shell_task_with_env(capfd):
+    if script_sh_2 is None:
+        return
     pst = get_portable_shell_task()
     pst(
         script_file=script_sh_2,
@@ -277,6 +282,8 @@ def test_shell_task_with_env(capfd):
 
 
 def test_shell_task_properly_restores_env_after_execution():
+    if script_sh_2 is None:
+        return
     env_as_dict = os.environ.copy()
     pst = get_portable_shell_task()
     pst(

From e217806820b9ad64307b197b3babb1c214ccbd10 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Mon, 18 Apr 2022 09:20:39 -0400
Subject: [PATCH 15/27] Add typing to make_export_string_from_env_dict params

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 1701ff5ae2..d22dc06691 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -184,7 +184,7 @@ def script(self) -> typing.Optional[str]:
     def script_file(self) -> typing.Optional[os.PathLike]:
         return self._script_file
 
-    def make_export_string_from_env_dict(self, d) -> str:
+    def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str:
         """
         Utility function to convert a dictionary of desired environment variable key: value pairs into a string of
         `

From 2bce64e95fe18e7951e497787efde9b83faf1129 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Mon, 18 Apr 2022 09:52:49 -0400
Subject: [PATCH 16/27] fixup doc string

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index d22dc06691..ec49ffeb32 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -187,11 +187,11 @@ def script_file(self) -> typing.Optional[os.PathLike]:
     def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str:
         """
         Utility function to convert a dictionary of desired environment variable key: value pairs into a string of
-        `
+        ```
         export k1=v1
         export k2=v2
         ...
-        `
+        ```
         """
         items = []
         for k, v in d.items():

From 3d26075dfaf023711e36a3960d0778d583d95eed Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Wed, 27 Apr 2022 15:43:35 -0400
Subject: [PATCH 17/27] Refactored the new behavior out into a new class. Did
 not change implementation details at all

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py                | 144 +++++++++++++++---
 .../flytekit/unit/extras/tasks/test_shell.py  |   8 +-
 2 files changed, 123 insertions(+), 29 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index ec49ffeb32..71b33a5bc8 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -184,6 +184,106 @@ def script(self) -> typing.Optional[str]:
     def script_file(self) -> typing.Optional[os.PathLike]:
         return self._script_file
 
+    def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
+        return self._config_task_instance.pre_execute(user_params)
+
+    def execute(self, **kwargs) -> typing.Any:
+        """
+        Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem
+        """
+        logger.info(f"Running shell script as type {self.task_type}")
+        if self.script_file:
+            with open(self.script_file) as f:
+                self._script = f.read()
+
+        outputs: typing.Dict[str, str] = {}
+        if self._output_locs:
+            for v in self._output_locs:
+                outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs)
+
+        if os.name == "nt":
+            self._script = self._script.lstrip().rstrip().replace("\n", "&&")
+
+        gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
+        if self._debug:
+            print("\n==============================================\n")
+            print(gen_script)
+            print("\n==============================================\n")
+
+        try:
+            subprocess.check_call(gen_script, shell=True)
+        except subprocess.CalledProcessError as e:
+            files = os.listdir(".")
+            fstr = "\n-".join(files)
+            logger.error(
+                f"Failed to Execute Script, return-code {e.returncode} \n"
+                f"StdErr: {e.stderr}\n"
+                f"StdOut: {e.stdout}\n"
+                f" Current directory contents: .\n-{fstr}"
+            )
+            raise
+
+        final_outputs = []
+        for v in self._output_locs:
+            if issubclass(v.var_type, FlyteFile):
+                final_outputs.append(FlyteFile(outputs[v.var]))
+            if issubclass(v.var_type, FlyteDirectory):
+                final_outputs.append(FlyteDirectory(outputs[v.var]))
+        if len(final_outputs) == 1:
+            return final_outputs[0]
+        if len(final_outputs) > 1:
+            return tuple(final_outputs)
+        return None
+
+    def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
+        return self._config_task_instance.post_execute(user_params, rval)
+
+
+class _RawShellTask(ShellTask):
+    """ """
+
+    def __init__(
+        self,
+        name: str,
+        debug: bool = False,
+        script: typing.Optional[str] = None,
+        script_file: typing.Optional[str] = None,
+        task_config: T = None,
+        inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
+        output_locs: typing.Optional[typing.List[OutputLocation]] = None,
+        **kwargs,
+    ):
+        """
+        The `_RawShellTask` is a minimal extension of the existing `ShellTask`. It's purpose is to support wrapping a
+        "raw" or "pure" shell script which needs to be executed with some environment variables set, and some arguments,
+        which may not be known until execution time.
+
+        This class is not meant to be instantiated into tasks by users, but used with the factory function
+        `get_raw_shell_task()`. An instance of this class will be returned with either user-specified or default
+        template. The template itself will export the desired environment variables, and subsequently execute the
+        desired "raw" script with the specified arguments.
+
+        .. note::
+            This means that within your workflow, you can dynamically control the env variables, arguments, and even the
+            actual script you want to run.
+
+        .. note::
+            The downside is that a dynamic workflow will be required. The "raw" script passed in at execution time must
+            be at the specified location.
+
+        These args are forwarded directly to the parent `ShellTask` constructor as behavior does not diverge
+        """
+        super().__init__(
+            name=name,
+            debug=debug,
+            script=script,
+            script_file=script_file,
+            task_config=task_config,
+            inputs=inputs,
+            output_locs=output_locs,
+            **kwargs,
+        )
+
     def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str:
         """
         Utility function to convert a dictionary of desired environment variable key: value pairs into a string of
@@ -198,9 +298,6 @@ def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str:
             items.append(f"export {k}={v}")
         return "\n".join(items)
 
-    def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
-        return self._config_task_instance.pre_execute(user_params)
-
     def execute(self, **kwargs) -> typing.Any:
         """
         Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem
@@ -219,8 +316,6 @@ def execute(self, **kwargs) -> typing.Any:
             self._script = self._script.lstrip().rstrip().replace("\n", "&&")
 
         if "env" in kwargs and isinstance(kwargs["env"], dict):
-            # This supports the portable_shell_task by adding an additional key:value pair to kwargs/input
-            # This will cause collisions if a user tries to use `env` AND `export_env` in their inputs
             kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"])
 
         gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs)
@@ -254,26 +349,25 @@ def execute(self, **kwargs) -> typing.Any:
             return tuple(final_outputs)
         return None
 
-    def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
-        return self._config_task_instance.post_execute(user_params, rval)
-
 
-# The portable_shell_task is an instance of ShellTask which wraps a 'pure' shell script
+# The raw_shell_task is an instance of _RawShellTask and wraps a 'pure' shell script
 # This utility function allows for the specification of env variables, arguments, and the actual script within the
-# workflow definition rather than at `ShellTask` instantiation
-def get_portable_shell_task() -> ShellTask:
-    return ShellTask(
-    name="portable_shell_task_instance",
-    debug=True,
-    inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
-    output_locs=[
-        OutputLocation(
-            var="k",
-            var_type=FlyteDirectory,
-            location="{ctx.working_directory}",
-        )
-    ],
-    script="""
+# workflow definition rather than at `_RawShellTask` instantiation
+def get_raw_shell_task(name: str = None) -> _RawShellTask:
+    _name = name if name else "raw_shell_task_instance"
+
+    return _RawShellTask(
+        name=_name,
+        debug=True,
+        inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
+        output_locs=[
+            OutputLocation(
+                var="out",
+                var_type=FlyteDirectory,
+                location="{ctx.working_directory}",
+            )
+        ],
+        script="""
 #!/bin/bash
 
 set -uex
@@ -283,5 +377,5 @@ def get_portable_shell_task() -> ShellTask:
 {inputs.export_env}
 
 bash {inputs.script_file} {inputs.script_args}
-    """
-)
\ No newline at end of file
+"""
+    )
\ No newline at end of file
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 2eafe90468..197296eaea 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -9,7 +9,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_portable_shell_task
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -256,7 +256,7 @@ def test_shell_script():
 def test_portable_shell_task_with_args(capfd):
     if script_sh_2 is None:
         return
-    pst = get_portable_shell_task()
+    pst = get_raw_shell_task()
     pst(
         script_file=script_sh_2,
         script_args="first_arg second_arg",
@@ -270,7 +270,7 @@ def test_portable_shell_task_with_args(capfd):
 def test_shell_task_with_env(capfd):
     if script_sh_2 is None:
         return
-    pst = get_portable_shell_task()
+    pst = get_raw_shell_task()
     pst(
         script_file=script_sh_2,
         env={"A": "AAAA", "B": "BBBB"},
@@ -285,7 +285,7 @@ def test_shell_task_properly_restores_env_after_execution():
     if script_sh_2 is None:
         return
     env_as_dict = os.environ.copy()
-    pst = get_portable_shell_task()
+    pst = get_raw_shell_task()
     pst(
         script_file=script_sh_2,
         env={"A": "AAAA", "B": "BBBB"},

From ce4603fe59cf4a7ff2b2d3e1d4f86470357898b3 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Mon, 2 May 2022 09:47:06 -0400
Subject: [PATCH 18/27] Address linter issues and up test cov

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py                |  2 +-
 .../flytekit/unit/extras/tasks/test_shell.py  | 45 ++++++++++++++++---
 .../extras/tasks/testdata/script_args_env.sh  |  4 +-
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 71b33a5bc8..1e249d0cb1 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -378,4 +378,4 @@ def get_raw_shell_task(name: str = None) -> _RawShellTask:
 
 bash {inputs.script_file} {inputs.script_args}
 """
-    )
\ No newline at end of file
+    )
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 197296eaea..cef54e22c9 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -1,6 +1,7 @@
 import datetime
 import os
 import tempfile
+import typing
 from dataclasses import dataclass
 from subprocess import CalledProcessError
 
@@ -9,7 +10,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task, _RawShellTask
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -253,7 +254,7 @@ def test_shell_script():
     t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
 
 
-def test_portable_shell_task_with_args(capfd):
+def test_raw_shell_task_with_args(capfd):
     if script_sh_2 is None:
         return
     pst = get_raw_shell_task()
@@ -267,7 +268,7 @@ def test_portable_shell_task_with_args(capfd):
     assert "second_arg" in cap.out
 
 
-def test_shell_task_with_env(capfd):
+def test_raw_shell_task_with_env(capfd):
     if script_sh_2 is None:
         return
     pst = get_raw_shell_task()
@@ -281,7 +282,7 @@ def test_shell_task_with_env(capfd):
     assert "BBBB" in cap.out
 
 
-def test_shell_task_properly_restores_env_after_execution():
+def test_raw_shell_task_properly_restores_env_after_execution():
     if script_sh_2 is None:
         return
     env_as_dict = os.environ.copy()
@@ -292,4 +293,38 @@ def test_shell_task_properly_restores_env_after_execution():
         script_args=""
     )
     env_as_dict_after = os.environ.copy()
-    assert env_as_dict == env_as_dict_after
\ No newline at end of file
+    assert env_as_dict == env_as_dict_after
+
+
+def test_raw_shell_task_instantiation(capfd):
+    pst = _RawShellTask(
+        name="test",
+        debug=True,
+        inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
+        output_locs=[
+            OutputLocation(
+                var="out",
+                var_type=FlyteDirectory,
+                location="{ctx.working_directory}",
+            )
+        ],
+        script="""
+#!/bin/bash
+
+set -uex
+
+cd {ctx.working_directory}
+
+{inputs.export_env}
+
+bash {inputs.script_file} {inputs.script_args}
+"""
+    )
+    pst(
+        script_file=script_sh_2,
+        script_args="first_arg second_arg",
+        env={}
+    )
+    cap = capfd.readouterr()
+    assert "first_arg" in cap.out
+    assert "second_arg" in cap.out
diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
index dd77131e22..ceb4bdf975 100644
--- a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
+++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
@@ -5,13 +5,13 @@ set -exo pipefail
 echo "A: $A"
 echo "B: $B"
 
-if [ ! -n "$1" ]; then
+if [ ! -z "$1" ]; then
     echo "$1"
 else
     echo "Unset first positional argument"
 fi
 
-if [ ! -n "$2" ]; then
+if [ ! -z "$2" ]; then
     echo "$2"
 else
     echo "Unset second positional argument"

From ead64955b19f23101621d5f4a0a43ccbc703472c Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Mon, 2 May 2022 11:02:42 -0400
Subject: [PATCH 19/27] Skip test on windows, no equivalent script

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 tests/flytekit/unit/extras/tasks/test_shell.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index cef54e22c9..d6ab80eb94 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -297,6 +297,8 @@ def test_raw_shell_task_properly_restores_env_after_execution():
 
 
 def test_raw_shell_task_instantiation(capfd):
+    if script_sh_2 is None:
+        return
     pst = _RawShellTask(
         name="test",
         debug=True,

From 2ebeecd59d381d6270431c928c3d5cd9b140c569 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Tue, 3 May 2022 10:41:04 -0400
Subject: [PATCH 20/27] Run black and isort, address SC2236

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py                |  2 +-
 .../flytekit/unit/extras/tasks/test_shell.py  | 28 ++++---------------
 .../extras/tasks/testdata/script_args_env.sh  |  6 ++--
 3 files changed, 10 insertions(+), 26 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 1e249d0cb1..e3de1756fb 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -377,5 +377,5 @@ def get_raw_shell_task(name: str = None) -> _RawShellTask:
 {inputs.export_env}
 
 bash {inputs.script_file} {inputs.script_args}
-"""
+""",
     )
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index d6ab80eb94..c369787761 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -10,7 +10,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task, _RawShellTask
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask, _RawShellTask, get_raw_shell_task
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -258,11 +258,7 @@ def test_raw_shell_task_with_args(capfd):
     if script_sh_2 is None:
         return
     pst = get_raw_shell_task()
-    pst(
-        script_file=script_sh_2,
-        script_args="first_arg second_arg",
-        env={}
-    )
+    pst(script_file=script_sh_2, script_args="first_arg second_arg", env={})
     cap = capfd.readouterr()
     assert "first_arg" in cap.out
     assert "second_arg" in cap.out
@@ -272,11 +268,7 @@ def test_raw_shell_task_with_env(capfd):
     if script_sh_2 is None:
         return
     pst = get_raw_shell_task()
-    pst(
-        script_file=script_sh_2,
-        env={"A": "AAAA", "B": "BBBB"},
-        script_args=""
-    )
+    pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="")
     cap = capfd.readouterr()
     assert "AAAA" in cap.out
     assert "BBBB" in cap.out
@@ -287,11 +279,7 @@ def test_raw_shell_task_properly_restores_env_after_execution():
         return
     env_as_dict = os.environ.copy()
     pst = get_raw_shell_task()
-    pst(
-        script_file=script_sh_2,
-        env={"A": "AAAA", "B": "BBBB"},
-        script_args=""
-    )
+    pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="")
     env_as_dict_after = os.environ.copy()
     assert env_as_dict == env_as_dict_after
 
@@ -320,13 +308,9 @@ def test_raw_shell_task_instantiation(capfd):
 {inputs.export_env}
 
 bash {inputs.script_file} {inputs.script_args}
-"""
-    )
-    pst(
-        script_file=script_sh_2,
-        script_args="first_arg second_arg",
-        env={}
+""",
     )
+    pst(script_file=script_sh_2, script_args="first_arg second_arg", env={})
     cap = capfd.readouterr()
     assert "first_arg" in cap.out
     assert "second_arg" in cap.out
diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
index ceb4bdf975..0927512a05 100644
--- a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
+++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh
@@ -5,13 +5,13 @@ set -exo pipefail
 echo "A: $A"
 echo "B: $B"
 
-if [ ! -z "$1" ]; then
+if [ -n "$1" ]; then
     echo "$1"
 else
     echo "Unset first positional argument"
 fi
 
-if [ ! -z "$2" ]; then
+if [ -n "$2" ]; then
     echo "$2"
 else
     echo "Unset second positional argument"
@@ -19,4 +19,4 @@ fi
 
 SOME_VAR="This var is set"
 
-echo "Reading SOME_VAR: ${SOME_VAR}"
\ No newline at end of file
+echo "Reading SOME_VAR: ${SOME_VAR}"

From e4c121dc69437fa091bd4ff541792641519831f5 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Tue, 3 May 2022 10:45:00 -0400
Subject: [PATCH 21/27] Refactored class name to remove _, make utility
 function require name so multiple uses in a workflow don't break

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py                 | 12 ++++++------
 tests/flytekit/unit/extras/tasks/test_shell.py |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index e3de1756fb..8cd2e39422 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -239,7 +239,7 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty
         return self._config_task_instance.post_execute(user_params, rval)
 
 
-class _RawShellTask(ShellTask):
+class RawShellTask(ShellTask):
     """ """
 
     def __init__(
@@ -254,7 +254,7 @@ def __init__(
         **kwargs,
     ):
         """
-        The `_RawShellTask` is a minimal extension of the existing `ShellTask`. It's purpose is to support wrapping a
+        The `RawShellTask` is a minimal extension of the existing `ShellTask`. It's purpose is to support wrapping a
         "raw" or "pure" shell script which needs to be executed with some environment variables set, and some arguments,
         which may not be known until execution time.
 
@@ -350,13 +350,13 @@ def execute(self, **kwargs) -> typing.Any:
         return None
 
 
-# The raw_shell_task is an instance of _RawShellTask and wraps a 'pure' shell script
+# The raw_shell_task is an instance of RawShellTask and wraps a 'pure' shell script
 # This utility function allows for the specification of env variables, arguments, and the actual script within the
-# workflow definition rather than at `_RawShellTask` instantiation
-def get_raw_shell_task(name: str = None) -> _RawShellTask:
+# workflow definition rather than at `RawShellTask` instantiation
+def get_raw_shell_task(name) -> RawShellTask:
     _name = name if name else "raw_shell_task_instance"
 
-    return _RawShellTask(
+    return RawShellTask(
         name=_name,
         debug=True,
         inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index c369787761..45f4801a80 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -10,7 +10,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask, _RawShellTask, get_raw_shell_task
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask, RawShellTask, get_raw_shell_task
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -287,7 +287,7 @@ def test_raw_shell_task_properly_restores_env_after_execution():
 def test_raw_shell_task_instantiation(capfd):
     if script_sh_2 is None:
         return
-    pst = _RawShellTask(
+    pst = RawShellTask(
         name="test",
         debug=True,
         inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),

From 215f48c51f6f39e900dd0a08d583de4e2b62fe58 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Tue, 3 May 2022 11:15:54 -0400
Subject: [PATCH 22/27] fix utility function

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 flytekit/extras/tasks/shell.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index 8cd2e39422..812c0a3749 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -353,11 +353,10 @@ def execute(self, **kwargs) -> typing.Any:
 # The raw_shell_task is an instance of RawShellTask and wraps a 'pure' shell script
 # This utility function allows for the specification of env variables, arguments, and the actual script within the
 # workflow definition rather than at `RawShellTask` instantiation
-def get_raw_shell_task(name) -> RawShellTask:
-    _name = name if name else "raw_shell_task_instance"
+def get_raw_shell_task(name: str) -> RawShellTask:
 
     return RawShellTask(
-        name=_name,
+        name=name,
         debug=True,
         inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str),
         output_locs=[

From 5b756919dbad6f974d4cb0c4f16bcc978d90572a Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Tue, 3 May 2022 11:16:26 -0400
Subject: [PATCH 23/27] fix utility function call in tests

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 tests/flytekit/unit/extras/tasks/test_shell.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 45f4801a80..11e9fa3d6f 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -257,7 +257,7 @@ def test_shell_script():
 def test_raw_shell_task_with_args(capfd):
     if script_sh_2 is None:
         return
-    pst = get_raw_shell_task()
+    pst = get_raw_shell_task(name="test")
     pst(script_file=script_sh_2, script_args="first_arg second_arg", env={})
     cap = capfd.readouterr()
     assert "first_arg" in cap.out
@@ -267,7 +267,7 @@ def test_raw_shell_task_with_args(capfd):
 def test_raw_shell_task_with_env(capfd):
     if script_sh_2 is None:
         return
-    pst = get_raw_shell_task()
+    pst = get_raw_shell_task(name="test")
     pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="")
     cap = capfd.readouterr()
     assert "AAAA" in cap.out
@@ -278,7 +278,7 @@ def test_raw_shell_task_properly_restores_env_after_execution():
     if script_sh_2 is None:
         return
     env_as_dict = os.environ.copy()
-    pst = get_raw_shell_task()
+    pst = get_raw_shell_task(name="test")
     pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="")
     env_as_dict_after = os.environ.copy()
     assert env_as_dict == env_as_dict_after

From 53e88f60dbd36d6b257c3ad997655eedde9b2838 Mon Sep 17 00:00:00 2001
From: Mike Zhong <mzhong@embarkvet.com>
Date: Wed, 4 May 2022 09:40:34 -0400
Subject: [PATCH 24/27] Fix isort on test_shell.py

Signed-off-by: Mike Zhong <mzhong@embarkvet.com>
---
 tests/flytekit/unit/extras/tasks/test_shell.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 11e9fa3d6f..7081d8d65f 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -10,7 +10,7 @@
 
 import flytekit
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import OutputLocation, ShellTask, RawShellTask, get_raw_shell_task
+from flytekit.extras.tasks.shell import OutputLocation, RawShellTask, ShellTask, get_raw_shell_task
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 

From 6d5e67af7be4352d1d50bc9baa19d4dd231f5afd Mon Sep 17 00:00:00 2001
From: Calvin Leather <cleather@embarkvet.com>
Date: Tue, 9 Aug 2022 21:29:24 -0400
Subject: [PATCH 25/27] adding logging settings to papermill plugin

Signed-off-by: Calvin Leather <cleather@embarkvet.com>
---
 .../flytekitplugins/papermill/task.py                | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
index 304932a828..fcceb7d70b 100644
--- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
+++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
@@ -1,5 +1,7 @@
 import json
+import logging
 import os
+import sys
 import typing
 from typing import Any
 
@@ -112,6 +114,7 @@ def __init__(
         name: str,
         notebook_path: str,
         render_deck: bool = False,
+        stream_logs: bool = False,
         task_config: T = None,
         inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
         outputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
@@ -132,6 +135,13 @@ def __init__(
         self._notebook_path = os.path.abspath(notebook_path)
 
         self._render_deck = render_deck
+        self._stream_logs = stream_logs
+
+        # Send the papermill logger to stdout so that it appears in pod logs. Note that papermill doesn't allow
+        # injecting a logger, so we cannot redirect logs to the flyte child loggers (e.g., the userspace logger)
+        # and we instead send to stdout directly
+        if self._stream_logs:
+            logging.getLogger("papermill").addHandler(logging.StreamHandler(sys.stdout))
 
         if not os.path.exists(self._notebook_path):
             raise ValueError(f"Illegal notebook path passed in {self._notebook_path}")
@@ -207,7 +217,7 @@ def execute(self, **kwargs) -> Any:
         """
         logger.info(f"Hijacking the call for task-type {self.task_type}, to call notebook.")
         # Execute Notebook via Papermill.
-        pm.execute_notebook(self._notebook_path, self.output_notebook_path, parameters=kwargs)  # type: ignore
+        pm.execute_notebook(self._notebook_path, self.output_notebook_path, parameters=kwargs, log_output=self._stream_logs)  # type: ignore
 
         outputs = self.extract_outputs(self.output_notebook_path)
         self.render_nb_html(self.output_notebook_path, self.rendered_output_path)

From 8b85e13c18e7860d30027abfb54db2f2e21f2778 Mon Sep 17 00:00:00 2001
From: Calvin Leather <cleather@embarkvet.com>
Date: Wed, 10 Aug 2022 06:29:27 -0400
Subject: [PATCH 26/27] set level to info logging

---
 plugins/flytekit-papermill/flytekitplugins/papermill/task.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
index fcceb7d70b..00ec942fcd 100644
--- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
+++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
@@ -141,7 +141,9 @@ def __init__(
         # injecting a logger, so we cannot redirect logs to the flyte child loggers (e.g., the userspace logger)
         # and we instead send to stdout directly
         if self._stream_logs:
-            logging.getLogger("papermill").addHandler(logging.StreamHandler(sys.stdout))
+            papermill_logger = logging.getLogger("papermill")
+            papermill_logger.addHandler(logging.StreamHandler(sys.stdout))
+            papermill_logger.setLevel(logging.INFO)
 
         if not os.path.exists(self._notebook_path):
             raise ValueError(f"Illegal notebook path passed in {self._notebook_path}")

From c854024ff9f49e2da39732111550606531b2a7c3 Mon Sep 17 00:00:00 2001
From: Calvin Leather <cleather@embarkvet.com>
Date: Wed, 10 Aug 2022 06:40:00 -0400
Subject: [PATCH 27/27] improved comments/docs

---
 .../flytekitplugins/papermill/task.py                  | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
index 00ec942fcd..9e9535dec2 100644
--- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
+++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py
@@ -86,6 +86,13 @@ class NotebookTask(PythonInstanceTask[T]):
 
         Users can access these notebooks after execution of the task locally or from remote servers.
 
+    .. note:
+
+        By default, print statements in your notebook won't be transmitted to the pod logs/stdout. If you would
+        like to have logs forwarded as the notebook executes, pass the stream_logs argument. Note that notebook
+        logs can be quite verbose, so ensure you are prepared for any downstream log ingestion costs
+        (e.g., cloudwatch)
+
     .. todo:
 
         Implicit extraction of SparkConfiguration from the notebook is not supported.
@@ -139,10 +146,11 @@ def __init__(
 
         # Send the papermill logger to stdout so that it appears in pod logs. Note that papermill doesn't allow
         # injecting a logger, so we cannot redirect logs to the flyte child loggers (e.g., the userspace logger)
-        # and we instead send to stdout directly
+        # and inherit their settings, but we instead must send logs to stdout directly
         if self._stream_logs:
             papermill_logger = logging.getLogger("papermill")
             papermill_logger.addHandler(logging.StreamHandler(sys.stdout))
+            # Papermill leaves the default level of DEBUG. We increase it here.
             papermill_logger.setLevel(logging.INFO)
 
         if not os.path.exists(self._notebook_path):
