Update lambda handler function extraction to work with python 3.9

codeguru_profiler_agent/aws_lambda/lambda_handler.py

@@ -1,34 +1,24 @@
 import os
-import logging
+import importlib
 from codeguru_profiler_agent.aws_lambda.profiler_decorator import with_lambda_profiler
 from codeguru_profiler_agent.agent_metadata.aws_lambda import HANDLER_ENV_NAME_FOR_CODEGURU_KEY
 HANDLER_ENV_NAME = "_HANDLER"
-logger = logging.getLogger(__name__)
 
 
 def restore_handler_env(original_handler, env=os.environ):
     env[HANDLER_ENV_NAME] = original_handler
 
 
-def load_handler(bootstrap_module, env=os.environ, original_handler_env_key=HANDLER_ENV_NAME_FOR_CODEGURU_KEY):
+def load_handler(handler_extractor, env=os.environ, original_handler_env_key=HANDLER_ENV_NAME_FOR_CODEGURU_KEY):
     try:
         original_handler_name = env.get(original_handler_env_key)
         if not original_handler_name:
             raise ValueError("Could not find module and function name from " + HANDLER_ENV_NAME_FOR_CODEGURU_KEY
                              + " environment variable")
 
-        # Delegate to the lambda code to load the customer's module.
-        if hasattr(bootstrap_module, '_get_handler'):
-            customer_handler_function = bootstrap_module._get_handler(original_handler_name)
-        else:
-            # TODO FIXME Review if the support for python 3.6 bootstrap can be improved.
-            # This returns both a init_handler and the function, we apply the init right away as we are in init process
-            init_handler, customer_handler_function = bootstrap_module._get_handlers(
-                handler=original_handler_name,
-                mode='event',  # with 'event' it will return the function as is (handlerfn in the lambda code)
-                               # 'http' would return wsgi.handle_one(sockfd, ('localhost', 80), handlerfn) instead
-                invokeid='unknown_id')  # FIXME invokeid is used for error handling, need to see if we can get it
-            init_handler()
+        # Delegate to the lambda code to load the customer's module and function.
+        customer_handler_function = handler_extractor(original_handler_name)
+
         restore_handler_env(original_handler_name, env)
         return customer_handler_function
     except:
@@ -39,10 +29,60 @@ def load_handler(bootstrap_module, env=os.environ, original_handler_env_key=HAND
         raise
 
 
-# Load the customer's handler, this should be done at import time which means it is done when lambda frameworks
-# loads our module. We load the bootstrap module by string name so that IDE does not complain
-lambda_bootstrap_module = __import__("bootstrap")
-handler_function = load_handler(lambda_bootstrap_module)
+def _python36_extractor(bootstrap_module, original_handler_name):
+    """
+    The lambda bootstrap code for python 3.6 was different than for later versions, instead of the _get_handler
+    function there was a more complex _get_handlers function with more parameters
+    """
+    # TODO FIXME Review if the support for python 3.6 bootstrap can be improved.
+    # This returns both a init_handler and the function, we apply the init right away as we are in init process
+    init_handler, customer_handler_function = bootstrap_module._get_handlers(
+        handler=original_handler_name,
+        mode='event',  # with 'event' it will return the function as is (handlerfn in the lambda code)
+        # 'http' would return wsgi.handle_one(sockfd, ('localhost', 80), handlerfn) instead
+        invokeid='unknown_id')  # FIXME invokeid is used for error handling, need to see if we can get it
+    init_handler()
+    return customer_handler_function
+
+
+def get_lambda_handler_extractor():
+    """
+    This loads and returns a function from lambda or RIC source code that is able to load the customer's
+    handler function.
+    WARNING !! This is a bit dangerous since we are calling internal functions from other modules that we do not
+    officially depend on. The idea is that this code should run only in a lambda function environment where we can know
+    what is available. However if lambda developers decide to change their internal code it could impact this !
+    """
+    # First try to load the lambda RIC if it is available (i.e. python 3.9)
+    # See https://github.com/aws/aws-lambda-python-runtime-interface-client
+    ric_bootstrap_module = _try_to_load_module("awslambdaric.bootstrap")
+    if ric_bootstrap_module is not None and hasattr(ric_bootstrap_module, '_get_handler'):
+        return ric_bootstrap_module._get_handler
+
+    # If no RIC module is available there should be a bootstrap module available
+    # do not catch ModuleNotFoundError exceptions here as we cannot do anything if this fails.
+    bootstrap_module = importlib.import_module("bootstrap")
+    if hasattr(bootstrap_module, '_get_handler'):
+        return bootstrap_module._get_handler
+    else:
+        return lambda handler_name: _python36_extractor(bootstrap_module, handler_name)
+
+
+def _try_to_load_module(module_name):
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError:
+        return None
+
+
+# We need to load the customer's handler function since the lambda framework loaded our function instead.
+# We want to delegate this work to the lambda framework so we need to find the appropriate method that does it
+# (depends on python versions) so we can call it.
+# This should be done at import time which means it is done when lambda frameworks loads our module
+handler_extractor = get_lambda_handler_extractor()
+
+# Now load the actual customer's handler function.
+handler_function = load_handler(handler_extractor)
 
 
 # WARNING: Do not rename this file, this function or HANDLER_ENV_NAME_FOR_CODEGURU without changing the bootstrap script

test/unit/aws_lambda/test_lambda_handler.py

@@ -11,6 +11,7 @@ def handler_function(event, context):
 
 
 init_handler_has_been_called = False
+python36_extractor_has_been_called = False
 
 
 def init_handler():
@@ -27,10 +28,18 @@ def _get_handler(self, handler):
 class BootstrapPython36ModuleMock:
     # for python3.6 version of lambda runtime bootstrap
     def _get_handlers(self, handler, mode, invokeid):
+        global python36_extractor_has_been_called
+        python36_extractor_has_been_called = True
         if handler == "handler_module.handler_function" and mode == "event":
             return init_handler, handler_function
 
 
+class RicBootstrapModuleMock:
+    def _get_handler(self, handler):
+        if handler == "handler_module.handler_function":
+            return handler_function
+
+
 class TestLambdaHandler:
     class TestWhenLambdaHandlerModuleIsLoaded:
         @pytest.fixture(autouse=True)
@@ -53,48 +62,86 @@ def test_call_handler_calls_the_inner_handler(self):
             assert lambda_handler_module.call_handler(event="expected_event",
                                                       context="expected_context") == "expected result"
 
-    class TestLoadModuleFunction:
-        class TestWhenPython38LambdaBootstrapCalls:
-            class TestWhenHandlerEnvIsSetProperly:
-                @before
-                def before(self):
-                    self.bootstrap = BootstrapModuleMock()
-                    self.env = {"HANDLER_ENV_NAME_FOR_CODEGURU": "handler_module.handler_function"}
-
-                def test_it_returns_the_handler_function(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    assert load_handler(self.bootstrap, self.env) == handler_function
-
-                def test_it_resets_handler_env_variable(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    load_handler(self.bootstrap, self.env)
-                    assert self.env['_HANDLER'] == "handler_module.handler_function"
-
-            class TestWhenHandlerEnvIsMissing:
-                @before
-                def before(self):
-                    self.bootstrap = BootstrapModuleMock()
-                    self.env = {}
-
-                def test_it_throws_value_error(self):
-                    with pytest.raises(ValueError):
-                        from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                        load_handler(self.bootstrap, self.env)
+    class TestGetHandlerExtractor:
+        class TestWhenRicIsAvailable:
+            @pytest.fixture(autouse=True)
+            def around(self):
+                # simulate that we are in a lambda environment where the awslambdaric.bootstrap module is available
+                self.module_available = RicBootstrapModuleMock()
+                sys.modules['awslambdaric.bootstrap'] = self.module_available
+                yield
+                del sys.modules['awslambdaric.bootstrap']
+
+            def test_it_loads_the_ric_module_code(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                result = get_lambda_handler_extractor()
+                assert result == self.module_available._get_handler
+
+        class TestWhenLambdaBootstrapIsAvailable:
+            @pytest.fixture(autouse=True)
+            def around(self):
+                # simulate that we are in a lambda environment where the awslambdaric.bootstrap module is not available
+                # but bootstrap from lambda is available.
+                self.module_available = BootstrapModuleMock()
+                if 'awslambdaric.bootstrap' in sys.modules:
+                    del sys.modules['awslambdaric.bootstrap']
+                sys.modules['bootstrap'] = self.module_available
+                yield
+                del sys.modules['bootstrap']
+
+            def test_it_loads_the_lambda_module_code(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                result = get_lambda_handler_extractor()
+                assert result == self.module_available._get_handler
 
         class TestWhenPython36LambdaBootstrapCalls:
             class TestWhenHandlerEnvIsSetProperly:
-                @before
-                def before(self):
-                    self.bootstrap = BootstrapPython36ModuleMock()
-                    self.env = {"HANDLER_ENV_NAME_FOR_CODEGURU": "handler_module.handler_function"}
+                @pytest.fixture(autouse=True)
+                def around(self):
+                    # simulate that we are in a lambda environment where the awslambdaric.bootstrap module is available
+                    sys.modules['bootstrap'] = BootstrapPython36ModuleMock()
                     global init_handler_has_been_called
                     init_handler_has_been_called = False
+                    global python36_extractor_has_been_called
+                    python36_extractor_has_been_called = False
+                    yield
+                    del sys.modules['bootstrap']
 
-                def test_it_returns_the_handler_function(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    assert load_handler(self.bootstrap, self.env) == handler_function
+                def test_it_uses_the_old_bootstrap_code(self):
+                    from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                    # call extractor
+                    get_lambda_handler_extractor()("handler_module.handler_function")
+                    assert python36_extractor_has_been_called
 
                 def test_it_calls_the_init_handler(self):
-                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
-                    load_handler(self.bootstrap, self.env)
+                    from codeguru_profiler_agent.aws_lambda.lambda_handler import get_lambda_handler_extractor
+                    # call extractor
+                    get_lambda_handler_extractor()("handler_module.handler_function")
                     assert init_handler_has_been_called
+
+    class TestLoadHandlerFunction:
+        class TestWhenHandlerEnvIsSetProperly:
+            @before
+            def before(self):
+                self.extractor = BootstrapModuleMock()._get_handler
+                self.env = {"HANDLER_ENV_NAME_FOR_CODEGURU": "handler_module.handler_function"}
+
+            def test_it_returns_the_handler_function(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
+                assert load_handler(self.extractor, self.env) == handler_function
+
+            def test_it_resets_handler_env_variable(self):
+                from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
+                load_handler(self.extractor, self.env)
+                assert self.env['_HANDLER'] == "handler_module.handler_function"
+
+        class TestWhenHandlerEnvIsMissing:
+            @before
+            def before(self):
+                self.extractor = BootstrapModuleMock()._get_handler
+                self.env = {}
+
+            def test_it_throws_value_error(self):
+                with pytest.raises(ValueError):
+                    from codeguru_profiler_agent.aws_lambda.lambda_handler import load_handler
+                    load_handler(self.extractor, self.env)