SSR changes

Author: zarirhamza

datadog_lambda/tracing.py

@@ -4,6 +4,7 @@
 # Copyright 2019 Datadog, Inc.
 import logging
 import os
+import re
 import traceback
 import ujson as json
 from datetime import datetime, timezone
@@ -856,10 +857,15 @@ def create_service_mapping(val):
     return new_service_mapping
 
 
-def determine_service_name(service_mapping, specific_key, generic_key, default_value):
+def determine_service_name(service_mapping,
+                           specific_key,
+                           generic_key,
+                           default_value,
+                           fallback=None):
     service_name = service_mapping.get(specific_key)
     if service_name is None:
-        service_name = service_mapping.get(generic_key, default_value)
+        service_name = service_mapping.get(generic_key,
+                                           default_value if default_value is not None else fallback)
     return service_name
 
 
@@ -988,6 +994,7 @@ def create_inferred_span_from_api_gateway_websocket_event(
         "http.url": http_url,
         "endpoint": endpoint,
         "resource_names": endpoint,
+        "span.kind": "server",
         "apiid": api_id,
         "apiname": api_id,
         "stage": request_context.get("stage"),
@@ -1046,6 +1053,7 @@ def create_inferred_span_from_api_gateway_event(
         "endpoint": path,
         "http.method": method,
         "resource_names": resource,
+        "span.kind": "server",
         "apiid": api_id,
         "apiname": api_id,
         "stage": request_context.get("stage"),
@@ -1150,12 +1158,13 @@ def create_inferred_span_from_sqs_event(event, context):
     event_source_arn = event_record.get("eventSourceARN")
     queue_name = event_source_arn.split(":")[-1]
     service_name = determine_service_name(
-        service_mapping, queue_name, "lambda_sqs", "sqs"
+        service_mapping, queue_name, "lambda_sqs", queue_name, "sqs"
     )
     attrs = event_record.get("attributes") or {}
     tags = {
         "operation_name": "aws.sqs",
         "resource_names": queue_name,
+        "span.kind": "server",
         "queuename": queue_name,
         "event_source_arn": event_source_arn,
         "receipt_handle": event_record.get("receiptHandle"),
@@ -1217,11 +1226,12 @@ def create_inferred_span_from_sns_event(event, context):
     topic_arn = sns_message.get("TopicArn")
     topic_name = topic_arn.split(":")[-1]
     service_name = determine_service_name(
-        service_mapping, topic_name, "lambda_sns", "sns"
+        service_mapping, topic_name, "lambda_sns", topic_name, "sns"
     )
     tags = {
         "operation_name": "aws.sns",
         "resource_names": topic_name,
+        "span.kind": "server",
         "topicname": topic_name,
         "topic_arn": topic_arn,
         "message_id": sns_message.get("MessageId"),
@@ -1252,15 +1262,16 @@ def create_inferred_span_from_kinesis_event(event, context):
     event_record = get_first_record(event)
     event_source_arn = event_record.get("eventSourceARN")
     event_id = event_record.get("eventID")
-    stream_name = event_source_arn.split(":")[-1]
+    stream_name = re.sub(r"^stream/", "", (event_source_arn or "").split(":")[-1])
     shard_id = event_id.split(":")[0]
     service_name = determine_service_name(
-        service_mapping, stream_name, "lambda_kinesis", "kinesis"
+        service_mapping, stream_name, "lambda_kinesis", stream_name, "kinesis"
     )
     kinesis = event_record.get("kinesis") or {}
     tags = {
         "operation_name": "aws.kinesis",
         "resource_names": stream_name,
+        "span.kind": "server",
         "streamname": stream_name,
         "shardid": shard_id,
         "event_source_arn": event_source_arn,
@@ -1287,12 +1298,13 @@ def create_inferred_span_from_dynamodb_event(event, context):
     event_source_arn = event_record.get("eventSourceARN")
     table_name = event_source_arn.split("/")[1]
     service_name = determine_service_name(
-        service_mapping, table_name, "lambda_dynamodb", "dynamodb"
+        service_mapping, table_name, "lambda_dynamodb", table_name, "dynamodb"
     )
     dynamodb_message = event_record.get("dynamodb") or {}
     tags = {
         "operation_name": "aws.dynamodb",
         "resource_names": table_name,
+        "span.kind": "server",
         "tablename": table_name,
         "event_source_arn": event_source_arn,
         "event_id": event_record.get("eventID"),
@@ -1321,11 +1333,12 @@ def create_inferred_span_from_s3_event(event, context):
     obj = s3.get("object") or {}
     bucket_name = bucket.get("name")
     service_name = determine_service_name(
-        service_mapping, bucket_name, "lambda_s3", "s3"
+        service_mapping, bucket_name, "lambda_s3", bucket_name, "s3"
     )
     tags = {
         "operation_name": "aws.s3",
         "resource_names": bucket_name,
+        "span.kind": "server",
         "event_name": event_record.get("eventName"),
         "bucketname": bucket_name,
         "bucket_arn": bucket.get("arn"),
@@ -1351,11 +1364,12 @@ def create_inferred_span_from_s3_event(event, context):
 def create_inferred_span_from_eventbridge_event(event, context):
     source = event.get("source")
     service_name = determine_service_name(
-        service_mapping, source, "lambda_eventbridge", "eventbridge"
+        service_mapping, source, "lambda_eventbridge", source, "eventbridge"
     )
     tags = {
         "operation_name": "aws.eventbridge",
         "resource_names": source,
+        "span.kind": "server",
         "detail_type": event.get("detail-type"),
     }
     InferredSpanInfo.set_tags(
@@ -1431,7 +1445,7 @@ def create_function_execution_span(
     tracer.set_tags(_dd_origin)
     span = tracer.trace(
         "aws.lambda",
-        service="aws.lambda",
+        service=config.service if config.service else function_name,
         resource=function_name,
         span_type="serverless",
     )

tests/event_samples/kinesisStream.json

@@ -0,0 +1,20 @@
+{
+    "Records": [
+        {
+            "kinesis": {
+                "kinesisSchemaVersion": "1.0",
+                "partitionKey": "partitionkey",
+                "sequenceNumber": "49624230154685806402418173680709770494154422022871973922",
+                "data": "eyJmb28iOiAiYmFyIiwgIl9kYXRhZG9nIjogeyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiNDk0ODM3NzMxNjM1NzI5MTQyMSIsICJ4LWRhdGFkb2ctcGFyZW50LWlkIjogIjI4NzYyNTMzODAwMTg2ODEwMjYiLCAieC1kYXRhZG9nLXNhbXBsaW5nLXByaW9yaXR5IjogIjEifX0=",
+                "approximateArrivalTimestamp": 1643638425.163
+            },
+            "eventSource": "aws:kinesis",
+            "eventVersion": "1.0",
+            "eventID": "shardId-000000000002:49624230154685806402418173680709770494154422022871973922",
+            "eventName": "aws:kinesis:record",
+            "invokeIdentityArn": "arn:aws:iam::601427279990:role/inferred-spans-python-dev-eu-west-1-lambdaRole",
+            "awsRegion": "eu-west-1",
+            "eventSourceARN": "arn:aws:kinesis:eu-west-1:601427279990:stream/kinesisStream"
+        }
+    ]
+}

tests/test_tracing.py

@@ -1478,7 +1478,7 @@ def test_remaps_specific_inferred_span_service_names_from_sqs_event(self):
         ] = "arn:aws:sqs:eu-west-1:123456789012:different-sqs-url"
         span2 = create_inferred_span(event2, ctx)
         self.assertEqual(span2.get_tag("operation_name"), "aws.sqs")
-        self.assertEqual(span2.service, "sqs")
+        self.assertEqual(span2.service, "different-sqs-url")
 
     def test_remaps_all_inferred_span_service_names_from_sns_event(self):
         self.set_service_mapping({"lambda_sns": "new-name"})
@@ -1524,7 +1524,7 @@ def test_remaps_specific_inferred_span_service_names_from_sns_event(self):
         ] = "arn:aws:sns:us-west-2:123456789012:different-sns-topic"
         span2 = create_inferred_span(event2, ctx)
         self.assertEqual(span2.get_tag("operation_name"), "aws.sns")
-        self.assertEqual(span2.service, "sns")
+        self.assertEqual(span2.service, "different-sns-topic")
 
     def test_remaps_all_inferred_span_service_names_from_kinesis_event(self):
         self.set_service_mapping({"lambda_kinesis": "new-name"})
@@ -1561,7 +1561,7 @@ def test_remaps_specific_inferred_span_service_names_from_kinesis_event(self):
 
         span1 = create_inferred_span(original_event, ctx)
         self.assertEqual(span1.get_tag("operation_name"), "aws.kinesis")
-        self.assertEqual(span1.service, "kinesis")
+        self.assertEqual(span1.service, "kinesisStream")
 
         # Testing the second event
         event2 = copy.deepcopy(original_event)
@@ -1570,7 +1570,7 @@ def test_remaps_specific_inferred_span_service_names_from_kinesis_event(self):
         ] = "arn:aws:kinesis:eu-west-1:601427279990:stream/DifferentKinesisStream"
         span2 = create_inferred_span(event2, ctx)
         self.assertEqual(span2.get_tag("operation_name"), "aws.kinesis")
-        self.assertEqual(span2.service, "kinesis")
+        self.assertEqual(span2.service, "DifferentKinesisStream")
 
     def test_remaps_all_inferred_span_service_names_from_s3_event(self):
         self.set_service_mapping({"lambda_s3": "new-name"})
@@ -1614,7 +1614,7 @@ def test_remaps_specific_inferred_span_service_names_from_s3_event(self):
         event2["Records"][0]["s3"]["bucket"]["name"] = "different-example-bucket"
         span2 = create_inferred_span(event2, ctx)
         self.assertEqual(span2.get_tag("operation_name"), "aws.s3")
-        self.assertEqual(span2.service, "s3")
+        self.assertEqual(span2.service, "different-example-bucket")
 
     def test_remaps_all_inferred_span_service_names_from_dynamodb_event(self):
         self.set_service_mapping({"lambda_dynamodb": "new-name"})
@@ -1660,7 +1660,7 @@ def test_remaps_specific_inferred_span_service_names_from_dynamodb_event(self):
         ] = "arn:aws:dynamodb:us-east-1:123456789012:table/DifferentExampleTableWithStream/stream/2015-06-27T00:48:05.899"
         span2 = create_inferred_span(event2, ctx)
         self.assertEqual(span2.get_tag("operation_name"), "aws.dynamodb")
-        self.assertEqual(span2.service, "dynamodb")
+        self.assertEqual(span2.service, "DifferentExampleTableWithStream")
 
     def test_remaps_all_inferred_span_service_names_from_eventbridge_event(self):
         self.set_service_mapping({"lambda_eventbridge": "new-name"})
@@ -1704,7 +1704,7 @@ def test_remaps_specific_inferred_span_service_names_from_eventbridge_event(
         event2["source"] = "different.eventbridge.custom.event.sender"
         span2 = create_inferred_span(event2, ctx)
         self.assertEqual(span2.get_tag("operation_name"), "aws.eventbridge")
-        self.assertEqual(span2.service, "eventbridge")
+        self.assertEqual(span2.service, "different.eventbridge.custom.event.sender")
 
 
 class _Span(object):
@@ -1927,7 +1927,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
     (
         "sqs-string-msg-attribute",
         _Span(
-            service="sqs",
+            service="InferredSpansQueueNode",
             start=1634662094.538,
             span_type="web",
             tags={
@@ -1945,7 +1945,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
     (
         "sns-string-msg-attribute",
         _Span(
-            service="sns",
+            service="serverlessTracingTopicPy",
             start=1643638421.637,
             span_type="web",
             tags={
@@ -1964,7 +1964,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
     (
         "sns-b64-msg-attribute",
         _Span(
-            service="sns",
+            service="serverlessTracingTopicPy",
             start=1643638421.637,
             span_type="web",
             tags={
@@ -1981,9 +1981,9 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
         ),
     ),
     (
-        "kinesis",
+        "kinesisStream",
         _Span(
-            service="kinesis",
+            service="kinesisStream",
             start=1643638425.163,
             span_type="web",
             tags={
@@ -2000,16 +2000,16 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
                 "operation_name": "aws.kinesis",
                 "partition_key": "partitionkey",
                 "request_id": None,
-                "resource_names": "stream/kinesisStream",
+                "resource_names": "kinesisStream",
                 "shardid": "shardId-000000000002",
-                "streamname": "stream/kinesisStream",
+                "streamname": "kinesisStream",
             },
         ),
     ),
     (
         "dynamodb",
         _Span(
-            service="dynamodb",
+            service="ExampleTableWithStream",
             start=1428537600.0,
             span_type="web",
             tags={
@@ -2035,7 +2035,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
     (
         "s3",
         _Span(
-            service="s3",
+            service="example-bucket",
             start=0.0,
             span_type="web",
             tags={
@@ -2060,7 +2060,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
     (
         "eventbridge-custom",
         _Span(
-            service="eventbridge",
+            service="eventbridge.custom.event.sender",
             start=1635989865.0,
             span_type="web",
             tags={
@@ -2079,7 +2079,7 @@ def __init__(self, service, start, span_type, parent_name=None, tags=None):
     (
         "eventbridge-sqs",
         _Span(
-            service="sqs",
+            service="eventbridge-sqs-queue",
             start=1691102943.638,
             span_type="web",
             parent_name="aws.eventbridge",