follow aws logic of handler importing on auto instrument. Use the runtime's logic if possible (#215)

Author: saartochner-lumigo

src/lumigo_tracer/auto_instrument_handler.py

@@ -1,56 +1,34 @@
-import warnings
 import os
-import importlib
 
-from lumigo_tracer import lumigo_tracer
+try:
+    # Try to import AWS's current _get_handler logic
+    from bootstrap import _get_handler as aws_get_handler
+except Exception:
+    # Import a snapshot of _get_handler from python38 runtime
+    from lumigo_tracer.libs.bootstrap import _get_handler as aws_get_handler
 
-with warnings.catch_warnings():
-    warnings.filterwarnings("ignore", category=DeprecationWarning)
-    import imp
+from lumigo_tracer import lumigo_tracer
 
 ORIGINAL_HANDLER_KEY = "LUMIGO_ORIGINAL_HANDLER"
 
 
-def parse_handler():
+def get_original_handler():
     try:
-        module_name, unit_name = os.environ[ORIGINAL_HANDLER_KEY].rsplit(".", 1)
-        file_handle, pathname, desc = imp.find_module(module_name)
+        return aws_get_handler(os.environ[ORIGINAL_HANDLER_KEY])
     except KeyError:
         raise Exception(
             "Could not find the original handler. Please contact Lumigo for more information."
         ) from None
-    except ValueError as e:
-        raise ValueError(
-            f"Runtime.MalformedHandlerName: Bad handler '{os.environ[ORIGINAL_HANDLER_KEY]}': {str(e)}"
-        ) from None
-    return module_name, unit_name, file_handle, pathname, desc
 
 
 @lumigo_tracer()
 def _handler(*args, **kwargs):
-    handler_module = ""
-    try:
-        handler_module, unit_name, file_handle, pathname, desc = parse_handler()
-        original_module = imp.load_module(handler_module, file_handle, pathname, desc)
-    except ImportError as e:
-        raise ImportError(
-            f"Runtime.ImportModuleError: Unable to import module '{handler_module}': {str(e)}"
-        ) from None
-    except SyntaxError as e:
-        raise SyntaxError(
-            f"Runtime.UserCodeSyntaxError: Syntax error in module '{handler_module}': {str(e)}"
-        ) from None
-    try:
-        original_handler = getattr(original_module, unit_name)
-    except AttributeError:
-        raise Exception(
-            f"Runtime.HandlerNotFound: Handler '{unit_name}' missing on module '{handler_module}'"
-        ) from None
+    original_handler = get_original_handler()
     return original_handler(*args, **kwargs)
 
 
 try:
     # import handler during runtime initialization, as usual.
-    importlib.import_module(parse_handler()[0])
+    get_original_handler()
 except Exception:
     pass

src/lumigo_tracer/libs/bootstrap.py

@@ -0,0 +1,80 @@
+"""
+Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+"""
+
+import warnings
+
+from lumigo_tracer.libs.lambda_runtime_exception import FaultException
+
+
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=DeprecationWarning)
+    import imp
+
+ERROR_LOG_LINE_TERMINATE = "\r"
+ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
+
+
+def _get_handler(handler):
+    try:
+        (modname, fname) = handler.rsplit(".", 1)
+    except ValueError as e:
+        fault = FaultException(
+            FaultException.MALFORMED_HANDLER_NAME, "Bad handler '{}': {}".format(handler, str(e))
+        )
+        return make_fault_handler(fault)
+
+    file_handle, pathname, desc = None, None, None
+    try:
+        # Recursively loading handler in nested directories
+        for segment in modname.split("."):
+            if pathname is not None:
+                pathname = [pathname]
+            file_handle, pathname, desc = imp.find_module(segment, pathname)
+        if file_handle is None:
+            module_type = desc[2]
+            if module_type == imp.C_BUILTIN:
+                fault = FaultException(
+                    FaultException.BUILT_IN_MODULE_CONFLICT,
+                    "Cannot use built-in module {} as a handler module".format(modname),
+                )
+                request_handler = make_fault_handler(fault)
+                return request_handler
+        m = imp.load_module(modname, file_handle, pathname, desc)
+    except ImportError as e:
+        fault = FaultException(
+            FaultException.IMPORT_MODULE_ERROR,
+            "Unable to import module '{}': {}".format(modname, str(e)),
+        )
+        request_handler = make_fault_handler(fault)
+        return request_handler
+    except SyntaxError as e:
+        trace = ['  File "%s" Line %s\n    %s' % (e.filename, e.lineno, e.text)]
+        fault = FaultException(
+            FaultException.USER_CODE_SYNTAX_ERROR,
+            "Syntax error in module '{}': {}".format(modname, str(e)),
+            trace,
+        )
+        request_handler = make_fault_handler(fault)
+        return request_handler
+    finally:
+        if file_handle is not None:
+            file_handle.close()
+
+    try:
+        request_handler = getattr(m, fname)
+    except AttributeError:
+        fault = FaultException(
+            FaultException.HANDLER_NOT_FOUND,
+            "Handler '{}' missing on module '{}'".format(fname, modname),
+            None,
+        )
+        request_handler = make_fault_handler(fault)
+    return request_handler
+
+
+def make_fault_handler(fault):
+    def result(*args):
+        raise fault
+
+    return result

src/lumigo_tracer/libs/lambda_runtime_exception.py

@@ -0,0 +1,19 @@
+"""
+Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+"""
+
+
+class FaultException(Exception):
+    MARSHAL_ERROR = "Runtime.MarshalError"
+    UNMARSHAL_ERROR = "Runtime.UnmarshalError"
+    USER_CODE_SYNTAX_ERROR = "Runtime.UserCodeSyntaxError"
+    HANDLER_NOT_FOUND = "Runtime.HandlerNotFound"
+    IMPORT_MODULE_ERROR = "Runtime.ImportModuleError"
+    BUILT_IN_MODULE_CONFLICT = "Runtime.BuiltInModuleConflict"
+    MALFORMED_HANDLER_NAME = "Runtime.MalformedHandlerName"
+    LAMBDA_CONTEXT_UNMARSHAL_ERROR = "Runtime.LambdaContextUnmarshalError"
+
+    def __init__(self, exception_type, msg, trace=None):
+        self.msg = msg
+        self.exception_type = exception_type
+        self.trace = trace

src/test/unit/test_auto_instrument_handler.py

@@ -6,46 +6,52 @@
 
 import pytest
 from lumigo_tracer.auto_instrument_handler import _handler, ORIGINAL_HANDLER_KEY
+from lumigo_tracer.libs.lambda_runtime_exception import FaultException
 
 
 def abc(*args, **kwargs):
     return {"hello": "world"}
 
 
-def test_happy_flow(monkeypatch):
+def test_happy_flow(monkeypatch, context):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "test_auto_instrument_handler.abc")
-    assert _handler({}, {}) == {"hello": "world"}
+    assert _handler({}, context) == {"hello": "world"}
 
 
-def test_hierarchy_happy_flow(monkeypatch):
+def test_hierarchy_happy_flow(monkeypatch, context):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "lumigo_tracer/test_module/test.handler")
-    assert _handler({}, {}) == {"hello": "world"}
+    assert _handler({}, context) == {"hello": "world"}
 
 
-def test_import_error(monkeypatch):
+def test_hierarchy_happy_flow_with_dots(monkeypatch, context):
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "lumigo_tracer.test_module.test.handler")
+    assert _handler({}, context) == {"hello": "world"}
+
+
+def test_import_error(monkeypatch, context):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "blabla.not.exists")
 
     try:
-        _handler({}, {})
-    except ImportError as e:
+        _handler({}, context)
+    except FaultException as e:
         # Note: We're not using pytest.raises in order to get the exception context
         assert "Runtime.ImportModuleError" in str(e)
         assert "another exception occurred" not in traceback.format_exc()
     else:
         assert False
 
 
-def test_no_env_handler_error(monkeypatch):
+def test_no_env_handler_error(monkeypatch, context):
     monkeypatch.delenv(ORIGINAL_HANDLER_KEY, None)
 
     with pytest.raises(Exception) as e:
-        _handler({}, {})
+        _handler({}, context)
     assert "Could not find the original handler" in str(e.value)
 
 
 def test_error_in_original_handler_no_extra_exception_log(monkeypatch, context):
     monkeypatch.setattr(imp, "load_module", mock.Mock(side_effect=ZeroDivisionError))
-    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "lumigo_tracer.test_module.test.handler")
 
     try:
         _handler({}, context)
@@ -58,37 +64,37 @@ def test_error_in_original_handler_no_extra_exception_log(monkeypatch, context):
 
 def test_error_in_original_handler_syntax_error(monkeypatch, context):
     monkeypatch.setattr(imp, "load_module", mock.Mock(side_effect=SyntaxError))
-    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.exit")
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "lumigo_tracer.test_module.test.handler")
 
     try:
         _handler({}, context)
-    except SyntaxError as e:
+    except FaultException as e:
         # Note: We're not using pytest.raises in order to get the exception context
         assert "Runtime.UserCodeSyntaxError" in str(e)
         assert "another exception occurred" not in traceback.format_exc()
     else:
         assert False
 
 
-def test_handler_bad_format(monkeypatch):
+def test_handler_bad_format(monkeypatch, context):
     monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "no_method")
 
     try:
-        _handler({}, {})
-    except ValueError as e:
+        _handler({}, context)
+    except FaultException as e:
         # Note: We're not using pytest.raises in order to get the exception context
         assert "Runtime.MalformedHandlerName" in str(e)
         assert "another exception occurred" not in traceback.format_exc()
     else:
         assert False
 
 
-def test_handler_not_found(monkeypatch):
-    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "sys.not_found")
+def test_handler_not_found(monkeypatch, context):
+    monkeypatch.setenv(ORIGINAL_HANDLER_KEY, "lumigo_tracer.not_found")
 
     try:
-        _handler({}, {})
-    except Exception as e:
+        _handler({}, context)
+    except FaultException as e:
         # Note: We're not using pytest.raises in order to get the exception context
         assert "Runtime.HandlerNotFound" in str(e)
         assert "another exception occurred" not in traceback.format_exc()