botocore: Add Lambda extension (#760)

* botocore: Add Lambda extension

* add extension to add additional attributes for lambda invoke
  calls
* move lambda specific tests to separate module

* changelog

Co-authored-by: Diego Hurtado <[email protected]>
Co-authored-by: Owais Lone <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -74,6 +74,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#735](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/735))
 - `opentelemetry-sdk-extension-aws` & `opentelemetry-propagator-aws` Remove unnecessary dependencies on `opentelemetry-test`
   ([#752](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/752))
+- `opentelemetry-instrumentation-botocore` Add Lambda extension
+  ([#760](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/760))
 
 ## [1.6.0-0.25b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.6.0-0.25b0) - 2021-10-13
 ### Added

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py

@@ -78,7 +78,6 @@ def response_hook(span, service_name, operation_name, result):
     ec2.describe_instances()
 """
 
-import json
 import logging
 from typing import Any, Callable, Collection, Dict, Optional, Tuple
 
@@ -161,27 +160,6 @@ def _uninstrument(self, **kwargs):
         unwrap(BaseClient, "_make_api_call")
         unwrap(Endpoint, "prepare_request")
 
-    @staticmethod
-    def _is_lambda_invoke(call_context: _AwsSdkCallContext):
-        return (
-            call_context.service == "lambda"
-            and call_context.operation == "Invoke"
-            and isinstance(call_context.params, dict)
-            and "Payload" in call_context.params
-        )
-
-    @staticmethod
-    def _patch_lambda_invoke(api_params):
-        try:
-            payload_str = api_params["Payload"]
-            payload = json.loads(payload_str)
-            headers = payload.get("headers", {})
-            inject(headers)
-            payload["headers"] = headers
-            api_params["Payload"] = json.dumps(payload)
-        except ValueError:
-            pass
-
     # pylint: disable=too-many-branches
     def _patched_api_call(self, original_func, instance, args, kwargs):
         if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
@@ -210,10 +188,6 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
             kind=call_context.span_kind,
             attributes=attributes,
         ) as span:
-            # inject trace context into payload headers for lambda Invoke
-            if BotocoreInstrumentor._is_lambda_invoke(call_context):
-                BotocoreInstrumentor._patch_lambda_invoke(call_context.params)
-
             _safe_invoke(extension.before_service_call, span)
             self._call_request_hook(span, call_context)
 

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/__init__.py

@@ -33,6 +33,7 @@ def loader():
 
 _KNOWN_EXTENSIONS = {
     "dynamodb": _lazy_load(".dynamodb", "_DynamoDbExtension"),
+    "lambda": _lazy_load(".lmbd", "_LambdaExtension"),
     "sqs": _lazy_load(".sqs", "_SqsExtension"),
 }
 

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/lmbd.py

@@ -0,0 +1,126 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import inspect
+import json
+import re
+from typing import Dict
+
+from opentelemetry.instrumentation.botocore.extensions.types import (
+    _AttributeMapT,
+    _AwsSdkCallContext,
+    _AwsSdkExtension,
+)
+from opentelemetry.propagate import inject
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace.span import Span
+
+
+class _LambdaOperation(abc.ABC):
+    @classmethod
+    @abc.abstractmethod
+    def operation_name(cls):
+        pass
+
+    @classmethod
+    def prepare_attributes(
+        cls, call_context: _AwsSdkCallContext, attributes: _AttributeMapT
+    ):
+        pass
+
+    @classmethod
+    def before_service_call(cls, call_context: _AwsSdkCallContext, span: Span):
+        pass
+
+
+class _OpInvoke(_LambdaOperation):
+    # https://docs.aws.amazon.com/lambda/latest/dg/API_Invoke.html#API_Invoke_RequestParameters
+    ARN_LAMBDA_PATTERN = re.compile(
+        "(?:arn:(?:aws[a-zA-Z-]*)?:lambda:)?"
+        "(?:[a-z]{2}(?:-gov)?-[a-z]+-\\d{1}:)?(?:\\d{12}:)?"
+        "(?:function:)?([a-zA-Z0-9-_\\.]+)(?::(?:\\$LATEST|[a-zA-Z0-9-_]+))?"
+    )
+
+    @classmethod
+    def operation_name(cls):
+        return "Invoke"
+
+    @classmethod
+    def extract_attributes(
+        cls, call_context: _AwsSdkCallContext, attributes: _AttributeMapT
+    ):
+        attributes[SpanAttributes.FAAS_INVOKED_PROVIDER] = "aws"
+        attributes[
+            SpanAttributes.FAAS_INVOKED_NAME
+        ] = cls._parse_function_name(call_context)
+        attributes[SpanAttributes.FAAS_INVOKED_REGION] = call_context.region
+
+    @classmethod
+    def _parse_function_name(cls, call_context: _AwsSdkCallContext):
+        function_name_or_arn = call_context.params.get("FunctionName")
+        matches = cls.ARN_LAMBDA_PATTERN.match(function_name_or_arn)
+        function_name = matches.group(1)
+        return function_name_or_arn if function_name is None else function_name
+
+    @classmethod
+    def before_service_call(cls, call_context: _AwsSdkCallContext, span: Span):
+        cls._inject_current_span(call_context)
+
+    @classmethod
+    def _inject_current_span(cls, call_context: _AwsSdkCallContext):
+        payload_str = call_context.params.get("Payload")
+        if payload_str is None:
+            return
+
+        # TODO: reconsider propagation via payload as it manipulates input of the called lambda function
+        try:
+            payload = json.loads(payload_str)
+            headers = payload.get("headers", {})
+            inject(headers)
+            payload["headers"] = headers
+            call_context.params["Payload"] = json.dumps(payload)
+        except ValueError:
+            pass
+
+
+################################################################################
+# Lambda extension
+################################################################################
+
+_OPERATION_MAPPING = {
+    op.operation_name(): op
+    for op in globals().values()
+    if inspect.isclass(op)
+    and issubclass(op, _LambdaOperation)
+    and not inspect.isabstract(op)
+}  # type: Dict[str, _LambdaOperation]
+
+
+class _LambdaExtension(_AwsSdkExtension):
+    def __init__(self, call_context: _AwsSdkCallContext):
+        super().__init__(call_context)
+        self._op = _OPERATION_MAPPING.get(call_context.operation)
+
+    def extract_attributes(self, attributes: _AttributeMapT):
+        if self._op is None:
+            return
+
+        self._op.extract_attributes(self._call_context, attributes)
+
+    def before_service_call(self, span: Span):
+        if self._op is None:
+            return
+
+        self._op.before_service_call(self._call_context, span)

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

@@ -11,26 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import io
 import json
-import sys
-import zipfile
 from unittest.mock import Mock, patch
 
 import botocore.session
 from botocore.exceptions import ParamValidationError
 from moto import (  # pylint: disable=import-error
     mock_ec2,
-    mock_iam,
     mock_kinesis,
     mock_kms,
-    mock_lambda,
     mock_s3,
     mock_sqs,
     mock_sts,
     mock_xray,
 )
-from pytest import mark
 
 from opentelemetry import trace as trace_api
 from opentelemetry.context import attach, detach, set_value
@@ -44,24 +38,6 @@
 _REQUEST_ID_REGEX_MATCH = r"[A-Z0-9]{52}"
 
 
-def get_as_zip_file(file_name, content):
-    zip_output = io.BytesIO()
-    with zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED) as zip_file:
-        zip_file.writestr(file_name, content)
-    zip_output.seek(0)
-    return zip_output.read()
-
-
-def return_headers_lambda_str():
-    pfunc = """
-def lambda_handler(event, context):
-    print("custom log event")
-    headers = event.get('headers', event.get('attributes', {}))
-    return headers
-"""
-    return pfunc
-
-
 # pylint:disable=too-many-public-methods
 class TestBotocoreInstrumentor(TestBase):
     """Botocore integration testsuite"""
@@ -278,79 +254,6 @@ def test_double_patch(self):
             "SQS", "ListQueues", request_id=_REQUEST_ID_REGEX_MATCH
         )
 
-    @mock_lambda
-    def test_lambda_client(self):
-        lamb = self._make_client("lambda")
-
-        lamb.list_functions()
-        self.assert_span("Lambda", "ListFunctions")
-
-    @mock_iam
-    def get_role_name(self):
-        iam = self._make_client("iam")
-        return iam.create_role(
-            RoleName="my-role",
-            AssumeRolePolicyDocument="some policy",
-            Path="/my-path/",
-        )["Role"]["Arn"]
-
-    @mark.skipif(
-        sys.platform == "win32",
-        reason="requires docker and Github CI Windows does not have docker installed by default",
-    )
-    @mock_lambda
-    def test_lambda_invoke_propagation(self):
-
-        previous_propagator = get_global_textmap()
-        try:
-            set_global_textmap(MockTextMapPropagator())
-
-            lamb = self._make_client("lambda")
-            lamb.create_function(
-                FunctionName="testFunction",
-                Runtime="python3.8",
-                Role=self.get_role_name(),
-                Handler="lambda_function.lambda_handler",
-                Code={
-                    "ZipFile": get_as_zip_file(
-                        "lambda_function.py", return_headers_lambda_str()
-                    )
-                },
-                Description="test lambda function",
-                Timeout=3,
-                MemorySize=128,
-                Publish=True,
-            )
-            # 2 spans for create IAM + create lambda
-            self.assertEqual(2, len(self.memory_exporter.get_finished_spans()))
-            self.memory_exporter.clear()
-
-            response = lamb.invoke(
-                Payload=json.dumps({}),
-                FunctionName="testFunction",
-                InvocationType="RequestResponse",
-            )
-
-            span = self.assert_span(
-                "Lambda", "Invoke", request_id=_REQUEST_ID_REGEX_MATCH
-            )
-            span_context = span.get_span_context()
-
-            # assert injected span
-            results = response["Payload"].read().decode("utf-8")
-            headers = json.loads(results)
-
-            self.assertEqual(
-                str(span_context.trace_id),
-                headers[MockTextMapPropagator.TRACE_ID_KEY],
-            )
-            self.assertEqual(
-                str(span_context.span_id),
-                headers[MockTextMapPropagator.SPAN_ID_KEY],
-            )
-        finally:
-            set_global_textmap(previous_propagator)
-
     @mock_kms
     def test_kms_client(self):
         kms = self._make_client("kms")