implements context propagation for lambda invoke + tests (#458)

Author: kuba-wu

Diff for: CHANGELOG.md

@@ -6,6 +6,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.2.0-0.21b0...HEAD)
 
+### Added
+- `opentelemetry-instrumentation-botocore` now supports
+  context propagation for lambda invoke via Payload embedded headers. 
+  ([#458](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/458))
+  
 ## [0.21b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.2.0-0.21b0) - 2021-05-11
 ### Changed
 

Diff for: instrumentation/opentelemetry-instrumentation-boto/setup.cfg

@@ -47,7 +47,7 @@ install_requires =
 [options.extras_require]
 test =
     boto~=2.0
-    moto~=1.0
+    moto~=2.0
     opentelemetry-test == 0.22.dev0
 
 [options.packages.find]

Diff for: instrumentation/opentelemetry-instrumentation-botocore/setup.cfg

@@ -45,7 +45,7 @@ install_requires =
 
 [options.extras_require]
 test =
-    moto ~= 1.0
+    moto[all] ~= 2.0
     opentelemetry-test == 0.22.dev0
 
 [options.packages.find]

Diff for: instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py

@@ -46,6 +46,7 @@
 ---
 """
 
+import json
 import logging
 
 from botocore.client import BaseClient
@@ -99,6 +100,27 @@ def _instrument(self, **kwargs):
     def _uninstrument(self, **kwargs):
         unwrap(BaseClient, "_make_api_call")
 
+    @staticmethod
+    def _is_lambda_invoke(service_name, operation_name, api_params):
+        return (
+            service_name == "lambda"
+            and operation_name == "Invoke"
+            and isinstance(api_params, dict)
+            and "Payload" in api_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"):
@@ -111,6 +133,12 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
         error = None
         result = None
 
+        # inject trace context into payload headers for lambda Invoke
+        if BotocoreInstrumentor._is_lambda_invoke(
+            service_name, operation_name, api_params
+        ):
+            BotocoreInstrumentor._patch_lambda_invoke(api_params)
+
         with self._tracer.start_as_current_span(
             "{}".format(service_name), kind=SpanKind.CLIENT,
         ) as span:

Diff for: instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

@@ -11,14 +11,17 @@
 # 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 zipfile
 from unittest.mock import Mock, patch
 
 import botocore.session
 from botocore.exceptions import ParamValidationError
 from moto import (  # pylint: disable=import-error
     mock_dynamodb2,
     mock_ec2,
+    mock_iam,
     mock_kinesis,
     mock_kms,
     mock_lambda,
@@ -37,6 +40,24 @@
 from opentelemetry.test.test_base import TestBase
 
 
+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
+
+
 class TestBotocoreInstrumentor(TestBase):
     """Botocore integration testsuite"""
 
@@ -328,6 +349,64 @@ def test_lambda_client(self):
             },
         )
 
+    @mock_iam
+    def get_role_name(self):
+        iam = self.session.create_client("iam", "us-east-1")
+        return iam.create_role(
+            RoleName="my-role",
+            AssumeRolePolicyDocument="some policy",
+            Path="/my-path/",
+        )["Role"]["Arn"]
+
+    @mock_lambda
+    def test_lambda_invoke_propagation(self):
+
+        previous_propagator = get_global_textmap()
+        try:
+            set_global_textmap(MockTextMapPropagator())
+
+            lamb = self.session.create_client(
+                "lambda", region_name="us-east-1"
+            )
+            lamb.create_function(
+                FunctionName="testFunction",
+                Runtime="python2.7",
+                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,
+            )
+            response = lamb.invoke(
+                Payload=json.dumps({}),
+                FunctionName="testFunction",
+                InvocationType="RequestResponse",
+            )
+
+            spans = self.memory_exporter.get_finished_spans()
+            assert spans
+            self.assertEqual(len(spans), 3)
+
+            results = response["Payload"].read().decode("utf-8")
+            headers = json.loads(results)
+
+            self.assertIn(MockTextMapPropagator.TRACE_ID_KEY, headers)
+            self.assertEqual(
+                "0", headers[MockTextMapPropagator.TRACE_ID_KEY],
+            )
+            self.assertIn(MockTextMapPropagator.SPAN_ID_KEY, headers)
+            self.assertEqual(
+                "0", headers[MockTextMapPropagator.SPAN_ID_KEY],
+            )
+        finally:
+            set_global_textmap(previous_propagator)
+
     @mock_kms
     def test_kms_client(self):
         kms = self.session.create_client("kms", region_name="us-east-1")