diff --git a/codeguru_profiler_agent/aws_lambda/lambda_handler.py b/codeguru_profiler_agent/aws_lambda/lambda_handler.py index f922465..9b1be9e 100644 --- a/codeguru_profiler_agent/aws_lambda/lambda_handler.py +++ b/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 diff --git a/test/unit/aws_lambda/test_lambda_handler.py b/test/unit/aws_lambda/test_lambda_handler.py index b925d21..a29bed9 100644 --- a/test/unit/aws_lambda/test_lambda_handler.py +++ b/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)