Repository: botify-labs/simpleflow
Issue: Improve execute.python

* don't hide the output
* harden exception catching
* modular logging (simpleflow.execute may not have a logger)

For this,

- [x] don't capture stdout or stderr
- [x] use a specific temp file or fd to communicate the exception, if any (using a file backend prevents pipe limits)
- [x] same for the result, in fact
- [x] accept a "logger_name" argument

Changed files:
simpleflow/swf/process/worker/base.py, simpleflow/exceptions.py, tests/test_simpleflow/swf/process/test_worker.py, simpleflow/execute.py, tests/test_simpleflow/test_execute.py, tests/integration/workflow.py

diff --git a/simpleflow/exceptions.py b/simpleflow/exceptions.py
index 50ff1be1..0f6185a9 100644
--- a/simpleflow/exceptions.py
+++ b/simpleflow/exceptions.py
@@ -124,3 +124,7 @@ class AggregateException(Exception):
 
     def __eq__(self, other):
         return self.exceptions == other.exceptions
+
+
+class ExecutionError(Exception):
+    pass
diff --git a/simpleflow/execute.py b/simpleflow/execute.py
index 555988ad..13e7c76f 100644
--- a/simpleflow/execute.py
+++ b/simpleflow/execute.py
@@ -1,29 +1,29 @@
 from __future__ import absolute_import, print_function
-import re
+
+import os
 import sys
-import subprocess
+
+try:
+    import subprocess32 as subprocess
+except ImportError:
+    import subprocess
 import functools
+import logging
+import tempfile
+import traceback
 
+# noinspection PyCompatibility
 from builtins import map
 
 from future.utils import iteritems
 
-from simpleflow import compat
 from swf import format
-
-try:
-    import cPickle as pickle
-except ImportError:
-    import pickle
-import base64
-import logging
-
+from simpleflow import compat
+from simpleflow.exceptions import ExecutionError
 from simpleflow.utils import json_dumps
 
 __all__ = ['program', 'python']
 
-logger = logging.getLogger(__name__)
-
 
 class RequiredArgument(object):
     pass
@@ -61,7 +61,7 @@ def format_arguments(*args, **kwargs):
             return '-' + str(key)  # short option -c
         return '--' + str(key)  # long option --val
 
-    return ['{}="{}"'.format(arg(key), value) for key, value in
+    return ['{}="{}"'.format(arg(k), v) for k, v in
             iteritems(kwargs)] + list(map(str, args))
 
 
@@ -139,7 +139,7 @@ def get_name(func):
     return '.'.join([prefix, name])
 
 
-def python(interpreter='python'):
+def python(interpreter='python', logger_name=__name__):
     """
     Execute a callable as an external Python program.
 
@@ -153,56 +153,58 @@ def python(interpreter='python'):
     def wrap_callable(func):
         @functools.wraps(func)
         def execute(*args, **kwargs):
+            logger = logging.getLogger(logger_name)
             command = 'simpleflow.execute'  # name of a module.
-            full_command = [
-                interpreter, '-m', command,  # execute module a script.
-                get_name(func), format_arguments_json(*args, **kwargs),
-            ]
-            try:
-                output = subprocess.check_output(
+            sys.stdout.flush()
+            sys.stderr.flush()
+            result_str = None  # useless
+            with tempfile.TemporaryFile() as result_fd, tempfile.TemporaryFile() as error_fd:
+                dup_result_fd = os.dup(result_fd.fileno())  # remove FD_CLOEXEC
+                dup_error_fd = os.dup(error_fd.fileno())  # remove FD_CLOEXEC
+                # print('error_fd: {}'.format(dup_error_fd))
+                full_command = [
+                    interpreter, '-m', command,  # execute module a script.
+                    get_name(func), format_arguments_json(*args, **kwargs),
+                    '--logger-name={}'.format(logger_name),
+                    '--result-fd={}'.format(dup_result_fd),
+                    '--error-fd={}'.format(dup_error_fd),
+                ]
+                if compat.PY2:  # close_fds doesn't work with python2 (using its C _posixsubprocess helper)
+                    close_fds = False
+                    pass_fds = ()
+                else:
+                    close_fds = True
+                    pass_fds = (dup_result_fd, dup_error_fd)
+                process = subprocess.Popen(
                     full_command,
-                    # Redirect stderr to stdout to get traceback on error.
-                    stderr=subprocess.STDOUT,
-                )
-            except subprocess.CalledProcessError as err:
-                err_output = err.output
-                if not compat.PY2:
-                    err_output = err_output.decode('utf-8', errors='replace')
-                logger.info(
-                    "Got a subprocess.CalledProcessError on command: {}\\n"
-                    "Original error output: '{}'".format(
-                        full_command, err_output, type(err_output)
-                    )
+                    bufsize=-1,
+                    close_fds=close_fds,
+                    pass_fds=pass_fds,
                 )
-                exclines = err_output.rstrip().rsplit('\\n', 2)
-                excline = exclines[-1]
-
-                try:
-                    exception = pickle.loads(
-                        base64.b64decode(excline.rstrip()))
-                except (TypeError, pickle.UnpicklingError):
-                    exception = Exception(excline)
-                    if ':' in excline:
-                        cls, msg = excline.split(':', 1)
-                        if re.match(r'\\s*[\\w.]+\\s*', cls):
-                            try:
-                                exception = eval('{}("{}")'.format(
-                                    cls.strip(),
-                                    msg.strip(),
-                                ))
-                            except BaseException as ex:
-                                logger.warning('{}'.format(ex))
-
-                raise exception
+                rc = process.wait()
+                os.close(dup_result_fd)
+                os.close(dup_error_fd)
+                if rc:
+                    error_fd.seek(0)
+                    err_output = error_fd.read()
+                    if err_output:
+                        if not compat.PY2:
+                            err_output = err_output.decode('utf-8', errors='replace')
+                    raise ExecutionError(err_output)
+
+                result_fd.seek(0)
+                result_str = result_fd.read()
+
+            if not result_str:
+                return None
             try:
                 if not compat.PY2:
-                    output = output.decode('utf-8', errors='replace')
-                last_line = output.rstrip().rsplit('\\n', 1)[-1]
-                d = format.decode(last_line)
-                return d
+                    result_str = result_str.decode('utf-8', errors='replace')
+                result = format.decode(result_str)
+                return result
             except BaseException as ex:
-                logger.warning('Exception in python.execute: {}'.format(ex))
-                logger.warning(repr(output))
+                logger.exception('Exception in python.execute: {} {}'.format(ex.__class__.__name__, ex))
+                logger.warning('%r', result_str)
 
         # Not automatically assigned in python < 3.2.
         execute.__wrapped__ = func
@@ -306,7 +308,7 @@ def make_callable(funcname):
     return callable_
 
 
-if __name__ == '__main__':
+def main():
     """
     When executed as a script, this module expects the name of a callable as
     its first argument and the arguments of the callable encoded in a JSON
@@ -345,7 +347,6 @@ if __name__ == '__main__':
 
     """
     import argparse
-
     parser = argparse.ArgumentParser()
     parser.add_argument(
         'funcname',
@@ -355,20 +356,38 @@ if __name__ == '__main__':
         'funcargs',
         help='callable arguments in JSON',
     )
-
+    parser.add_argument(
+        '--logger-name',
+        help='logger name',
+    )
+    parser.add_argument(
+        '--result-fd',
+        type=int,
+        default=1,
+        metavar='N',
+        help='result file descriptor',
+    )
+    parser.add_argument(
+        '--error-fd',
+        type=int,
+        default=2,
+        metavar='N',
+        help='error file descriptor',
+    )
     cmd_arguments = parser.parse_args()
-
     funcname = cmd_arguments.funcname
     try:
         arguments = format.decode(cmd_arguments.funcargs)
     except:
         raise ValueError('cannot load arguments from {}'.format(
             cmd_arguments.funcargs))
-
+    if cmd_arguments.logger_name:
+        logger = logging.getLogger(cmd_arguments.logger_name)
+    else:
+        logger = logging.getLogger(__name__)
     callable_ = make_callable(funcname)
     if hasattr(callable_, '__wrapped__'):
         callable_ = callable_.__wrapped__
-
     args = arguments.get('args', ())
     kwargs = arguments.get('kwargs', {})
     try:
@@ -378,13 +397,31 @@ if __name__ == '__main__':
             result = callable_(*args, **kwargs)
     except Exception as err:
         logger.error('Exception: {}'.format(err))
-        # Use base64 encoding to avoid carriage returns and special characters.
-        # FIXME change this: brittle, missing traceback
-        encoded_err = base64.b64encode(pickle.dumps(err))
+        exc_type, exc_value, exc_traceback = sys.exc_info()
+        tb = traceback.format_tb(exc_traceback)
+        details = json_dumps(
+            {
+                'error': exc_type.__name__,
+                'message': str(exc_value),
+                'traceback': tb,
+            },
+            default=repr,
+        )
+        if cmd_arguments.error_fd == 2:
+            sys.stderr.flush()
         if not compat.PY2:
-            # Convert bytes to string
-            encoded_err = encoded_err.decode('utf-8', errors='replace')
-        print(encoded_err)
+            details = details.encode('utf-8')
+        os.write(cmd_arguments.error_fd, details)
         sys.exit(1)
-    else:
-        print(json_dumps(result))
+
+    if cmd_arguments.result_fd == 1:  # stdout (legacy)
+        sys.stdout.flush()  # may have print's in flight
+        os.write(cmd_arguments.result_fd, b'\\n')
+    result = json_dumps(result)
+    if not compat.PY2:
+        result = result.encode('utf-8')
+    os.write(cmd_arguments.result_fd, result)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/simpleflow/swf/process/worker/base.py b/simpleflow/swf/process/worker/base.py
index d81d4643..8a68a299 100644
--- a/simpleflow/swf/process/worker/base.py
+++ b/simpleflow/swf/process/worker/base.py
@@ -3,17 +3,21 @@ import logging
 import multiprocessing
 import os
 import signal
+import sys
 import traceback
 
 import psutil
+
+from simpleflow.exceptions import ExecutionError
 from swf import format
 import swf.actors
 import swf.exceptions
 from simpleflow.process import Supervisor, with_state
 from simpleflow.swf.process import Poller
+
 from simpleflow.swf.task import ActivityTask
 from simpleflow.swf.utils import sanitize_activity_context
-from simpleflow.utils import format_exc
+from simpleflow.utils import format_exc, json_dumps
 
 from .dispatch import dynamic_dispatcher
 
@@ -141,17 +145,37 @@ class ActivityWorker(object):
             context = sanitize_activity_context(task.context)
             context['domain_name'] = poller.domain.name
             result = ActivityTask(activity, *args, context=context, **kwargs).execute()
-        except Exception as err:
-            logger.exception("process error: {}".format(str(err)))
-            tb = traceback.format_exc()
-            return poller.fail_with_retry(token, task, reason=format_exc(err), details=tb)
+        except Exception:
+            exc_type, exc_value, exc_traceback = sys.exc_info()
+            logger.exception("process error: {}".format(str(exc_value)))
+            if isinstance(exc_value, ExecutionError) and len(exc_value.args):
+                details = exc_value.args[0]
+                reason = format_exc(exc_value)  # FIXME json.loads and rebuild?
+            else:
+                tb = traceback.format_tb(exc_traceback)
+                reason = format_exc(exc_value)
+                details = json_dumps(
+                    {
+                        'error': exc_type.__name__,
+                        'message': str(exc_value),
+                        'traceback': tb,
+                    },
+                    default=repr
+                )
+            return poller.fail_with_retry(
+                token,
+                task,
+                reason=reason,
+                details=details
+            )
 
         try:
             poller.complete_with_retry(token, result)
         except Exception as err:
             logger.exception("complete error")
-            reason = 'cannot complete task {}: {}'.format(
+            reason = 'cannot complete task {}: {} {}'.format(
                 task.activity_id,
+                err.__class__.__name__,
                 err,
             )
             poller.fail_with_retry(token, task, reason)
diff --git a/tests/integration/workflow.py b/tests/integration/workflow.py
index 8b449759..3bc59f96 100644
--- a/tests/integration/workflow.py
+++ b/tests/integration/workflow.py
@@ -190,7 +190,6 @@ class SignaledWorkflow(Workflow):
         return future.result
 
 
-
 class WorkflowToCancel(Workflow):
     name = 'example'
     version = 'example'
diff --git a/tests/test_simpleflow/swf/process/test_worker.py b/tests/test_simpleflow/swf/process/test_worker.py
index a4beb58f..381951a2 100644
--- a/tests/test_simpleflow/swf/process/test_worker.py
+++ b/tests/test_simpleflow/swf/process/test_worker.py
@@ -1,5 +1,5 @@
 from collections import namedtuple
-from mock import patch, ANY
+from mock import patch
 import unittest
 
 from moto import mock_swf
@@ -30,3 +30,7 @@ class TestActivityWorker(unittest.TestCase):
         self.assertEquals(1, mock.call_count)
         self.assertEquals(mock.call_args[0], ("token", task))
         self.assertIn("No module named ", mock.call_args[1]["reason"])
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_simpleflow/test_execute.py b/tests/test_simpleflow/test_execute.py
index f834c0a7..c1eff4f4 100644
--- a/tests/test_simpleflow/test_execute.py
+++ b/tests/test_simpleflow/test_execute.py
@@ -1,5 +1,7 @@
 # coding: utf-8
 from __future__ import print_function
+
+import json
 import tempfile
 import os.path
 import platform
@@ -7,6 +9,7 @@ import platform
 import pytest
 
 from simpleflow import execute
+from simpleflow.exceptions import ExecutionError
 
 
 @execute.program(path='ls')
@@ -139,15 +142,17 @@ def test_function_as_program_with_kwargs():
 
 
 def test_function_as_program_raises_builtin_exception():
-    with pytest.raises(TypeError):
+    with pytest.raises(ExecutionError) as excinfo:
         add('1')
-    with pytest.raises(TypeError):
+    assert '"error":"TypeError"' in str(excinfo.value)
+    with pytest.raises(ExecutionError) as excinfo:
         Add('1')
+    assert '"error":"TypeError"' in str(excinfo.value)
 
 
 @execute.python()
 def print_string(s, retval):
-    print(s)
+    print(s, end='')
     return retval
 
 
@@ -163,8 +168,10 @@ class PrintString(object):
 
 
 def test_function_with_print():
-    assert print_string("This isn't part of the return value", None) is None
-    assert PrintString("This isn't part of the return value", None) is None
+    actual = print_string("This isn't part of the return value", None)
+    assert actual is None, actual
+    actual = PrintString("This isn't part of the return value", None)
+    assert actual is None, actual
 
 
 def test_function_with_print_and_return():
@@ -197,10 +204,12 @@ class RaiseDummyException(object):
 
 
 def test_function_as_program_raises_custom_exception():
-    with pytest.raises(DummyException):
+    with pytest.raises(ExecutionError) as excinfo:
         raise_dummy_exception()
-    with pytest.raises(DummyException):
+    assert '"error":"DummyException"' in str(excinfo.value)
+    with pytest.raises(ExecutionError) as excinfo:
         RaiseDummyException()
+    assert '"error":"DummyException"' in str(excinfo.value)
 
 
 @execute.python()
@@ -210,12 +219,9 @@ def raise_timeout_error():
 
 
 def test_function_as_program_raises_module_exception():
-    from simpleflow.exceptions import TimeoutError
-
-    with pytest.raises(TimeoutError) as err:
+    with pytest.raises(ExecutionError) as excinfo:
         raise_timeout_error()
-        assert err.timeout_type == 'timeout'
-        assert err.timeout_value == 1
+    assert '"error":"TimeoutError"' in str(excinfo.value)
 
 
 @execute.python()
@@ -246,6 +252,8 @@ def raise_dummy_exception_with_unicode():
 
 
 def test_exception_with_unicode():
-    with pytest.raises(DummyException) as err:
+    with pytest.raises(ExecutionError) as excinfo:
         raise_dummy_exception_with_unicode()
-        assert err.args[0] == u'ʘ‿ʘ'
+    assert '"error":"DummyException"' in str(excinfo.value)
+    error = json.loads(excinfo.value.args[0])
+    assert error['message'] == u'ʘ‿ʘ'