Repository: DataDog/dd-trace-py

Issue:
The aws_lambda patch module can't handle callable object event handlers
### Summary of problem
NOTE: I'll be cross-posting this in [datadog-lambda-python](https://github.com/DataDog/datadog-lambda-python) project because I'm not sure whether this is in the domain of the Lambda layer or this library itself. This could either be solved in this library by checking if the function being wrapped is a callable class instance, or by the Lambda layer passing the `__call__` function reference instead of the callable class.

The issue stems from ddtrace (specifically the `aws_lambda` module) trying to patch the outer most Lambda handler functions, which can be any callable that accepts an event and context arguments, so either a standard function or a callable object, like a class instance that defines `__call__`.

The internal wrapper is passed the class instance itself instead of the `__call__` function however, so as a result, it fails as a class instance isn't going to have a `__code__` property. The `patch_all` the Lambda layer is executing is in a try/catch block and only runs once, so there's no errors being raised, but on every cold start there's a stack trace popping up in the logs. 

Platform: AWS Lambda
Runtime: Python 3.9
Serverless Framework Version: 3.28.0
Serverless Plugin Version: 5.20.0
DataDog Lambda Layer Version: 68
ddtrace Version: 1.9.1

### Which version of dd-trace-py are you using?
ddtrace@1.9.1

### Which version of pip are you using?
pip@21.2.4

### Which libraries and their versions are you using?
Output of `poetry show --no-dev`, excluded some internal libraries irrelevant to this and added `chalice` since that's packaged manually using their internal packaging mechanism.

attrs@21.4.0
aws-lambda-context@1.1.0
boto3@1.24.57
botocore@1.27.57
bytecode@0.14.0
cattrs@22.2.0
certifi@2022.6.15
charset-normalizer@2.1.1
ddsketch@2.0.4
ddtrace@1.9.1
elasticsearch@7.9.1
elasticsearch-dsl@7.3.0
envier@0.4.0
exceptiongroup@1.1.0
idna@3.3
jmespath@0.10.0
jsonschema@3.2.0
marshmallow@3.15.0
packaging@23.0
protobuf@4.22.0
pyrsistent@0.18.1
python-dateutil@2.8.2
python-json-logger@2.0.4
requests@2.28.1
requests-aws4auth@1.0
s3transfer@0.6.0
simplejson@3.17.2
six@1.16.0
tenacity@8.2.2
typing-extensions@4.3.0
urllib3@1.26.12
xmltodict@0.13.0
chalice@1.26.6

### How can we reproduce your problem?
Deploy a lambda function where the event handler is a callable class instance, whether through the [chalice](https://github.com/aws/chalice) framework or something like this:
```
# app.py
from typing import Any

class LambdaHandler:
    def __call__(self, event: dict[str, Any], context: Any) -> dict[str, Any]:
        return {"hello": "world"}

app = LambdaHandler()
```
With `app.app` being the handler of the lambda function.

### What is the result that you get?
The following stack trace is thrown on every cold start as `patch_all` fails:
```
Traceback (most recent call last):
	File "/opt/python/lib/python3.9/site-packages/datadog_lambda/wrapper.py", line 177, in __init__
		patch_all()
	File "/opt/python/lib/python3.9/site-packages/datadog_lambda/patch.py", line 36, in patch_all
		patch_all_dd()
	File "/var/task/ddtrace/_monkey.py", line 190, in patch_all
		patch(raise_errors=False, **modules)
	File "/var/task/ddtrace/_monkey.py", line 230, in patch
		when_imported(module)(_on_import_factory(contrib, raise_errors=False))
	File "/var/task/ddtrace/vendor/wrapt/importer.py", line 284, in register
		register_post_import_hook(hook, name)
	File "/var/task/ddtrace/vendor/wrapt/decorators.py", line 470, in _synchronized
		return wrapped(*args, **kwargs)
	File "/var/task/ddtrace/vendor/wrapt/importer.py", line 82, in register_post_import_hook
		hook(module)
	File "/var/task/ddtrace/_monkey.py", line 155, in on_import
		imported_module.patch()
	File "/var/task/ddtrace/contrib/aws_lambda/patch.py", line 156, in patch
		wrap(handler, wrapper)
	File "/var/task/ddtrace/internal/wrapping.py", line 428, in wrap
		f.__code__,
	AttributeError: 'EventSourceHandler' object has no attribute '__code__'
```

### What is the result that you expected?
This operation should not fail, and either [ddtrace.internal.wrapping.wrap](https://github.com/DataDog/dd-trace-py/blob/1.x/ddtrace/internal/wrapping.py#L417) should account for callable classes, or `datadog_lambda` should be figuring out that the handler is a callable class instance and passing the `__call__` function reference instead of the class instance.



Discussion:
Hello @michaeldimchuk,

Thank you for raising this issue. I will work on multiple POCs to reproduce this and will come with a bug fix.

Resolving patch:
diff --git a/ddtrace/contrib/aws_lambda/patch.py b/ddtrace/contrib/aws_lambda/patch.py
--- a/ddtrace/contrib/aws_lambda/patch.py
+++ b/ddtrace/contrib/aws_lambda/patch.py
@@ -140,16 +140,42 @@ def _get_handler_and_module():
     if path is None:
         from datadog_lambda.wrapper import datadog_lambda_wrapper
 
-        handler = getattr(datadog_lambda_wrapper, "__call__")
+        wrapper_module = datadog_lambda_wrapper
+        wrapper_handler = getattr(datadog_lambda_wrapper, "__call__")
 
-        return handler, datadog_lambda_wrapper, _datadog_instrumentation
+        return wrapper_handler, wrapper_module, _datadog_instrumentation
     else:
         parts = path.rsplit(".", 1)
         (mod_name, handler_name) = parts
         modified_mod_name = _modify_module_name(mod_name)
         handler_module = import_module(modified_mod_name)
         handler = getattr(handler_module, handler_name)
-        return handler, handler_module, _datadog_instrumentation
+
+        if callable(handler):
+            class_name = type(handler).__name__
+            is_function = not isinstance(handler, type) and hasattr(handler, "__code__") and class_name == "function"
+            # handler is a function
+            #
+            # note: this is a best effort to identify function based handlers
+            # this will not cover all cases
+            if is_function:
+                return handler, handler_module, _datadog_instrumentation
+
+            # handler must be either a class or an instance of a class
+            #
+            # note: if handler is a class instance with `__code__` defined,
+            # we will prioritize the `__call__` method, ignoring `__code__`.
+            class_module = getattr(handler_module, class_name)
+            class_handler = getattr(class_module, "__call__")
+
+            if isinstance(handler, type):
+                # class handler is a metaclass
+                if hasattr(class_handler, "__func__"):
+                    class_handler = class_handler.__func__
+
+            return class_handler, class_module, _datadog_instrumentation
+        else:
+            raise TypeError("Handler type is not supported to patch.")
 
 
 def _has_patch_module():


Test patch:
diff --git a/tests/contrib/aws_lambda/handlers.py b/tests/contrib/aws_lambda/handlers.py
--- a/tests/contrib/aws_lambda/handlers.py
+++ b/tests/contrib/aws_lambda/handlers.py
@@ -7,12 +7,9 @@
 
 @datadog_lambda_wrapper
 def manually_wrapped_handler(event, context, success=True):
-    @tracer.wrap("result-trace")
-    def get_result():
+    with tracer.trace("result-trace"):
         return {"success": success}
 
-    return get_result()
-
 
 def handler(event, context, success=True):
     return {"success": success}
@@ -37,6 +34,45 @@ def finishing_spans_early_handler(event, context):
     return {"success": False}
 
 
+class StaticMeta(type):
+    @staticmethod
+    def __call__(event, context):
+        with tracer.trace("result-trace"):
+            return {"success": True}
+
+
+class ClassMeta(type):
+    @classmethod
+    def __call__(cls, event, context):
+        with tracer.trace("result-trace"):
+            return {"success": True}
+
+
+class ClassHandler(metaclass=ClassMeta):
+    pass
+
+
+class StaticHandler(metaclass=StaticMeta):
+    pass
+
+
+class InstanceHandler:
+    def __call__(self, event, context):
+        with tracer.trace("result-trace"):
+            return {"success": True}
+
+
+class InstanceHandlerWithCode(InstanceHandler):
+    def __code__(self):
+        return
+
+
+static_handler = StaticHandler
+class_handler = ClassHandler
+instance_handler = InstanceHandler()
+instance_handler_with_code = InstanceHandlerWithCode()
+
+
 def datadog(_handler):
     @tracer.wrap("aws.lambda")
     def wrapped(*args):
diff --git a/tests/contrib/aws_lambda/test_aws_lambda.py b/tests/contrib/aws_lambda/test_aws_lambda.py
--- a/tests/contrib/aws_lambda/test_aws_lambda.py
+++ b/tests/contrib/aws_lambda/test_aws_lambda.py
@@ -1,14 +1,17 @@
-import os
-
 import pytest
 
 from ddtrace.contrib.aws_lambda import patch
 from ddtrace.contrib.aws_lambda import unpatch
+from tests.contrib.aws_lambda.handlers import class_handler
 from tests.contrib.aws_lambda.handlers import datadog
 from tests.contrib.aws_lambda.handlers import finishing_spans_early_handler
 from tests.contrib.aws_lambda.handlers import handler
+from tests.contrib.aws_lambda.handlers import instance_handler
+from tests.contrib.aws_lambda.handlers import instance_handler_with_code
 from tests.contrib.aws_lambda.handlers import manually_wrapped_handler
+from tests.contrib.aws_lambda.handlers import static_handler
 from tests.contrib.aws_lambda.handlers import timeout_handler
+from tests.utils import override_env
 
 
 class LambdaContext:
@@ -32,22 +35,24 @@ def create_context(remaining_time_in_millis=300):
     return create_context
 
 
+def get_env(env={}):
+    common_env = {
+        "DD_TRACE_AGENT_URL": "http://localhost:9126/",
+        "DD_TRACE_ENABLED": "true",
+    }
+    return {**common_env, **env}
+
+
 @pytest.fixture(autouse=True)
 def setup():
-    os.environ.update(
-        {
-            "DD_TRACE_AGENT_URL": "http://localhost:9126/",
-            "DD_TRACE_ENABLED": "true",
-        }
-    )
     yield
     unpatch()
 
 
 @pytest.mark.parametrize("customApmFlushDeadline", [("-100"), ("10"), ("100"), ("200")])
-@pytest.mark.snapshot()
+@pytest.mark.snapshot
 def test_timeout_traces(context, customApmFlushDeadline):
-    os.environ.update(
+    env = get_env(
         {
             "AWS_LAMBDA_FUNCTION_NAME": "timeout_handler",
             "DD_LAMBDA_HANDLER": "tests.contrib.aws_lambda.handlers.timeout_handler",
@@ -55,44 +60,46 @@ def test_timeout_traces(context, customApmFlushDeadline):
         }
     )
 
-    patch()
+    with override_env(env):
+        patch()
 
-    datadog(timeout_handler)({}, context())
+        datadog(timeout_handler)({}, context())
 
 
-@pytest.mark.snapshot()
+@pytest.mark.snapshot
 def test_continue_on_early_trace_ending(context):
     """
     These scenario expects no timeout error being tagged on the root span
     when closing all spans in the customers code and reaching a timeout.
     """
-    os.environ.update(
+    env = get_env(
         {
             "AWS_LAMBDA_FUNCTION_NAME": "finishing_spans_early_handler",
             "DD_LAMBDA_HANDLER": "tests.contrib.aws_lambda.handlers.finishing_spans_early_handler",
         }
     )
 
-    patch()
+    with override_env(env):
+        patch()
 
-    datadog(finishing_spans_early_handler)({}, context())
+        datadog(finishing_spans_early_handler)({}, context())
 
 
 @pytest.mark.snapshot
 async def test_file_patching(context):
-    os.environ.update(
+    env = get_env(
         {
             "AWS_LAMBDA_FUNCTION_NAME": "handler",
             "DD_LAMBDA_HANDLER": "tests.contrib.aws_lambda.handlers.handler",
         }
     )
 
-    patch()
+    with override_env(env):
+        patch()
 
-    result = datadog(handler)({}, context())
+        result = datadog(handler)({}, context())
 
-    assert result == {"success": True}
-    return
+        assert result == {"success": True}
 
 
 @pytest.mark.snapshot
@@ -100,16 +107,40 @@ async def test_module_patching(mocker, context):
     mocker.patch("datadog_lambda.wrapper._LambdaDecorator._before")
     mocker.patch("datadog_lambda.wrapper._LambdaDecorator._after")
 
-    os.environ.update(
+    env = get_env(
         {
             "AWS_LAMBDA_FUNCTION_NAME": "manually_wrapped_handler",
         }
     )
 
-    os.environ.pop("DD_LAMBDA_HANDLER")
-    patch()
+    with override_env(env):
+        patch()
+
+        result = manually_wrapped_handler({}, context())
+
+        assert result == {"success": True}
+
+
+@pytest.mark.parametrize(
+    "handler,function_name",
+    [
+        (static_handler, "static_handler"),
+        (class_handler, "class_handler"),
+        (instance_handler, "instance_handler"),
+        (instance_handler_with_code, "instance_handler_with_code"),
+    ],
+)
+@pytest.mark.snapshot
+def test_class_based_handlers(context, handler, function_name):
+    env = get_env(
+        {
+            "AWS_LAMBDA_FUNCTION_NAME": function_name,
+            "DD_LAMBDA_HANDLER": "tests.contrib.aws_lambda.handlers." + function_name,
+        }
+    )
 
-    result = manually_wrapped_handler({}, context())
+    with override_env(env):
+        patch()
 
-    assert result == {"success": True}
-    return
+        result = datadog(handler)({}, context())
+        assert result == {"success": True}
diff --git a/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[ClassHandler-class_handler].json b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[ClassHandler-class_handler].json
new file mode 100644
--- /dev/null
+++ b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[ClassHandler-class_handler].json
@@ -0,0 +1,37 @@
+[[
+  {
+    "name": "aws.lambda",
+    "service": "",
+    "resource": "aws.lambda",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "language": "python",
+      "runtime-id": "a39fc906196b4fa6a4e778f9dcb0568a"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 91275
+    },
+    "duration": 112000,
+    "start": 1680106790650063000
+  },
+     {
+       "name": "result-trace",
+       "service": "",
+       "resource": "result-trace",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "type": "",
+       "error": 0,
+       "duration": 12000,
+       "start": 1680106790650139000
+     }]]
diff --git a/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[StaticHandler-static_handler].json b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[StaticHandler-static_handler].json
new file mode 100644
--- /dev/null
+++ b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[StaticHandler-static_handler].json
@@ -0,0 +1,37 @@
+[[
+  {
+    "name": "aws.lambda",
+    "service": "",
+    "resource": "aws.lambda",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "language": "python",
+      "runtime-id": "a39fc906196b4fa6a4e778f9dcb0568a"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 91275
+    },
+    "duration": 111000,
+    "start": 1680106790636572000
+  },
+     {
+       "name": "result-trace",
+       "service": "",
+       "resource": "result-trace",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "type": "",
+       "error": 0,
+       "duration": 11000,
+       "start": 1680106790636647000
+     }]]
diff --git a/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[handler2-instance_handler].json b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[handler2-instance_handler].json
new file mode 100644
--- /dev/null
+++ b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[handler2-instance_handler].json
@@ -0,0 +1,37 @@
+[[
+  {
+    "name": "aws.lambda",
+    "service": "",
+    "resource": "aws.lambda",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "language": "python",
+      "runtime-id": "a39fc906196b4fa6a4e778f9dcb0568a"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 91275
+    },
+    "duration": 110000,
+    "start": 1680106790663281000
+  },
+     {
+       "name": "result-trace",
+       "service": "",
+       "resource": "result-trace",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "type": "",
+       "error": 0,
+       "duration": 11000,
+       "start": 1680106790663356000
+     }]]
diff --git a/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[handler3-instance_handler_with_code].json b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[handler3-instance_handler_with_code].json
new file mode 100644
--- /dev/null
+++ b/tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_class_based_handlers[handler3-instance_handler_with_code].json
@@ -0,0 +1,37 @@
+[[
+  {
+    "name": "aws.lambda",
+    "service": "",
+    "resource": "aws.lambda",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "language": "python",
+      "runtime-id": "dc30016e848e4fdb84e4781e5448a48a"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 99955
+    },
+    "duration": 118000,
+    "start": 1680533051713574000
+  },
+     {
+       "name": "result-trace",
+       "service": "",
+       "resource": "result-trace",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "type": "",
+       "error": 0,
+       "duration": 11000,
+       "start": 1680533051713656000
+     }]]