Parse function version from function arn (DataDog#69)

* Parse function version from function arn

* Add unit tests for create_function_execution_span

* Show what's wrong with format

* Fix formatting issues

* Default function version to $Latest

* $Latest to $LATEST

* $Latest to $LATEST

Author: hghotra

datadog_lambda/tracing.py

@@ -239,9 +239,13 @@ def create_function_execution_span(
     tags = {}
     if context:
         function_arn = (context.invoked_function_arn or "").lower()
+        tk = function_arn.split(":")
+        function_arn = ":".join(tk[0:7]) if len(tk) > 7 else function_arn
+        function_version = tk[7] if len(tk) > 7 else "$LATEST"
         tags = {
             "cold_start": str(is_cold_start).lower(),
             "function_arn": function_arn,
+            "function_version": function_version,
             "request_id": context.aws_request_id,
             "resource_names": context.function_name,
         }

scripts/check_format.sh

@@ -8,8 +8,8 @@ if [ "$PYTHON_VERSION" = "2" ]; then
 fi
 pip install -Iv black==19.10b0
 
-python -m black --check datadog_lambda/
-python -m black --check tests
+python -m black --check datadog_lambda/ --diff
+python -m black --check tests --diff
 
 
 

tests/test_tracing.py

@@ -10,13 +10,30 @@
 from datadog_lambda.constants import SamplingPriority, TraceHeader, XraySubsegment
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
+    create_function_execution_span,
     get_dd_trace_context,
     set_correlation_ids,
     _convert_xray_trace_id,
     _convert_xray_entity_id,
     _convert_xray_sampling,
 )
 
+function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test"
+
+
+def get_mock_context(
+    aws_request_id="request-id-1",
+    memory_limit_in_mb="256",
+    invoked_function_arn=function_arn,
+    function_version="1",
+):
+    lambda_context = MagicMock()
+    lambda_context.aws_request_id = aws_request_id
+    lambda_context.memory_limit_in_mb = memory_limit_in_mb
+    lambda_context.invoked_function_arn = invoked_function_arn
+    lambda_context.function_version = function_version
+    return lambda_context
+
 
 class TestExtractAndGetDDTraceContext(unittest.TestCase):
     def setUp(self):
@@ -195,3 +212,27 @@ def test_set_correlation_ids(self):
         trace_id, span_id = get_correlation_ids()
         self.assertEqual(trace_id, "123")
         self.assertEqual(span_id, "456")
+
+
+class TestFunctionSpanTags(unittest.TestCase):
+    def test_function(self):
+        ctx = get_mock_context()
+        span = create_function_execution_span(ctx, "", False, {"source": ""}, False)
+        self.assertEqual(span.get_tag("function_arn"), function_arn)
+        self.assertEqual(span.get_tag("function_version"), "$LATEST")
+
+    def test_function_with_version(self):
+        function_version = "1"
+        ctx = get_mock_context(
+            invoked_function_arn=function_arn + ":" + function_version
+        )
+        span = create_function_execution_span(ctx, "", False, {"source": ""}, False)
+        self.assertEqual(span.get_tag("function_arn"), function_arn)
+        self.assertEqual(span.get_tag("function_version"), function_version)
+
+    def test_function_with_alias(self):
+        function_alias = "alias"
+        ctx = get_mock_context(invoked_function_arn=function_arn + ":" + function_alias)
+        span = create_function_execution_span(ctx, "", False, {"source": ""}, False)
+        self.assertEqual(span.get_tag("function_arn"), function_arn)
+        self.assertEqual(span.get_tag("function_version"), function_alias)