diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 487321a6..c6f7d694 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -80,9 +80,9 @@ def _get_xray_trace_context(): if xray_trace_entity is None: return None trace_context = { - "trace-id": _convert_xray_trace_id(xray_trace_entity["trace_id"]), - "parent-id": _convert_xray_entity_id(xray_trace_entity["parent_id"]), - "sampling-priority": _convert_xray_sampling(xray_trace_entity["sampled"]), + "trace-id": _convert_xray_trace_id(xray_trace_entity.get("trace_id")), + "parent-id": _convert_xray_entity_id(xray_trace_entity.get("parent_id")), + "sampling-priority": _convert_xray_sampling(xray_trace_entity.get("sampled")), } logger.debug( "Converted trace context %s from X-Ray segment %s", @@ -204,7 +204,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): Falls back to lambda context if no trace data is found in the SQS message attributes. """ try: - first_record = event["Records"][0] + first_record = event.get("Records")[0] # logic to deal with SNS => SQS event if "body" in first_record: @@ -215,7 +215,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): logger.debug("Found SNS message inside SQS event") first_record = get_first_record(create_sns_event(body)) except Exception: - first_record = event["Records"][0] + first_record = event.get("Records")[0] pass msg_attributes = first_record.get( @@ -259,7 +259,7 @@ def extract_context_from_eventbridge_event(event, lambda_context): This is only possible if Details is a JSON string. """ try: - detail = event["detail"] + detail = event.get("detail") dd_context = detail.get("_datadog") if not dd_context: return extract_context_from_lambda_context(lambda_context) @@ -409,8 +409,10 @@ def get_dd_trace_context(): context = xray_context elif xray_context and dd_trace_context: context = dd_trace_context.copy() - context["parent-id"] = xray_context["parent-id"] - logger.debug("Set parent id from xray trace context: %s", context["parent-id"]) + context["parent-id"] = xray_context.get("parent-id") + logger.debug( + "Set parent id from xray trace context: %s", context.get("parent-id") + ) if dd_tracing_enabled: dd_trace_py_context = _get_dd_trace_py_context() @@ -487,7 +489,7 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces): if merge_xray_traces: xray_context = _get_xray_trace_context() if xray_context is not None: - context["parent-id"] = xray_context["parent-id"] + context["parent-id"] = xray_context.get("parent-id") headers = _context_obj_to_headers(context) span_context = propagator.extract(headers) @@ -537,7 +539,6 @@ def create_inferred_span(event, context): elif event_source.equals(EventTypes.EVENTBRIDGE): logger.debug("Eventbridge event detected. Inferring a span") return create_inferred_span_from_eventbridge_event(event, context) - except Exception as e: logger.debug( "Unable to infer span. Detected type: {}. Reason: {}", @@ -550,10 +551,10 @@ def create_inferred_span(event, context): def create_inferred_span_from_lambda_function_url_event(event, context): - request_context = event["requestContext"] - domain = request_context["domainName"] - method = request_context["http"]["method"] - path = request_context["http"]["path"] + request_context = event.get("requestContext") + domain = request_context.get("domainName") + method = request_context.get("http").get("method") + path = request_context.get("http").get("path") resource = "{0} {1}".format(method, path) tags = { "operation_name": "aws.lambda.url", @@ -563,7 +564,7 @@ def create_inferred_span_from_lambda_function_url_event(event, context): "resource_names": domain + path, "request_id": context.aws_request_id, } - request_time_epoch = request_context["timeEpoch"] + request_time_epoch = request_context.get("timeEpoch") args = { "service": domain, "resource": resource, @@ -584,29 +585,29 @@ def create_inferred_span_from_lambda_function_url_event(event, context): def is_api_gateway_invocation_async(event): return ( "headers" in event - and "X-Amz-Invocation-Type" in event["headers"] - and event["headers"]["X-Amz-Invocation-Type"] == "Event" + and "X-Amz-Invocation-Type" in event.get("headers") + and event.get("headers").get("X-Amz-Invocation-Type") == "Event" ) def create_inferred_span_from_api_gateway_websocket_event(event, context): - request_context = event["requestContext"] - domain = request_context["domainName"] - endpoint = request_context["routeKey"] + request_context = event.get("requestContext") + domain = request_context.get("domainName") + endpoint = request_context.get("routeKey") tags = { "operation_name": "aws.apigateway.websocket", "http.url": domain + endpoint, "endpoint": endpoint, "resource_names": endpoint, - "apiid": request_context["apiId"], - "apiname": request_context["apiId"], - "stage": request_context["stage"], - "request_id": request_context["requestId"], - "connection_id": request_context["connectionId"], - "event_type": request_context["eventType"], - "message_direction": request_context["messageDirection"], + "apiid": request_context.get("apiId"), + "apiname": request_context.get("apiId"), + "stage": request_context.get("stage"), + "request_id": request_context.get("requestId"), + "connection_id": request_context.get("connectionId"), + "event_type": request_context.get("eventType"), + "message_direction": request_context.get("messageDirection"), } - request_time_epoch = request_context["requestTimeEpoch"] + request_time_epoch = request_context.get("requestTimeEpoch") if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -625,10 +626,10 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): def create_inferred_span_from_api_gateway_event(event, context): - request_context = event["requestContext"] - domain = request_context["domainName"] - method = event["httpMethod"] - path = event["path"] + request_context = event.get("requestContext") + domain = request_context.get("domainName") + method = event.get("httpMethod") + path = event.get("path") resource = "{0} {1}".format(method, path) tags = { "operation_name": "aws.apigateway.rest", @@ -636,12 +637,12 @@ def create_inferred_span_from_api_gateway_event(event, context): "endpoint": path, "http.method": method, "resource_names": resource, - "apiid": request_context["apiId"], - "apiname": request_context["apiId"], - "stage": request_context["stage"], - "request_id": request_context["requestId"], + "apiid": request_context.get("apiId"), + "apiname": request_context.get("apiId"), + "stage": request_context.get("stage"), + "request_id": request_context.get("requestId"), } - request_time_epoch = request_context["requestTimeEpoch"] + request_time_epoch = request_context.get("requestTimeEpoch") if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -660,26 +661,26 @@ def create_inferred_span_from_api_gateway_event(event, context): def create_inferred_span_from_http_api_event(event, context): - request_context = event["requestContext"] - domain = request_context["domainName"] - method = request_context["http"]["method"] - path = event["rawPath"] + request_context = event.get("requestContext") + domain = request_context.get("domainName") + method = request_context.get("http").get("method") + path = event.get("rawPath") resource = "{0} {1}".format(method, path) tags = { "operation_name": "aws.httpapi", "endpoint": path, "http.url": domain + path, - "http.method": request_context["http"]["method"], - "http.protocol": request_context["http"]["protocol"], - "http.source_ip": request_context["http"]["sourceIp"], - "http.user_agent": request_context["http"]["userAgent"], + "http.method": request_context.get("http").get("method"), + "http.protocol": request_context.get("http").get("protocol"), + "http.source_ip": request_context.get("http").get("sourceIp"), + "http.user_agent": request_context.get("http").get("userAgent"), "resource_names": resource, "request_id": context.aws_request_id, - "apiid": request_context["apiId"], - "apiname": request_context["apiId"], - "stage": request_context["stage"], + "apiid": request_context.get("apiId"), + "apiname": request_context.get("apiId"), + "stage": request_context.get("stage"), } - request_time_epoch = request_context["timeEpoch"] + request_time_epoch = request_context.get("timeEpoch") if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -701,18 +702,18 @@ def create_inferred_span_from_sqs_event(event, context): trace_ctx = tracer.current_trace_context() event_record = get_first_record(event) - event_source_arn = event_record["eventSourceARN"] + event_source_arn = event_record.get("eventSourceARN") queue_name = event_source_arn.split(":")[-1] tags = { "operation_name": "aws.sqs", "resource_names": queue_name, "queuename": queue_name, "event_source_arn": event_source_arn, - "receipt_handle": event_record["receiptHandle"], - "sender_id": event_record["attributes"]["SenderId"], + "receipt_handle": event_record.get("receiptHandle"), + "sender_id": event_record.get("attributes").get("SenderId"), } InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") - request_time_epoch = event_record["attributes"]["SentTimestamp"] + request_time_epoch = event_record.get("attributes").get("SentTimestamp") args = { "service": "sqs", "resource": queue_name, @@ -754,25 +755,25 @@ def create_inferred_span_from_sqs_event(event, context): def create_inferred_span_from_sns_event(event, context): event_record = get_first_record(event) - sns_message = event_record["Sns"] - topic_arn = event_record["Sns"]["TopicArn"] + sns_message = event_record.get("Sns") + topic_arn = event_record.get("Sns").get("TopicArn") topic_name = topic_arn.split(":")[-1] tags = { "operation_name": "aws.sns", "resource_names": topic_name, "topicname": topic_name, "topic_arn": topic_arn, - "message_id": sns_message["MessageId"], - "type": sns_message["Type"], + "message_id": sns_message.get("MessageId"), + "type": sns_message.get("Type"), } # Subject not available in SNS => SQS scenario if "Subject" in sns_message and sns_message["Subject"]: - tags["subject"] = sns_message["Subject"] + tags["subject"] = sns_message.get("Subject") InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") sns_dt_format = "%Y-%m-%dT%H:%M:%S.%fZ" - timestamp = event_record["Sns"]["Timestamp"] + timestamp = event_record.get("Sns").get("Timestamp") dt = datetime.strptime(timestamp, sns_dt_format) args = { @@ -790,8 +791,8 @@ def create_inferred_span_from_sns_event(event, context): def create_inferred_span_from_kinesis_event(event, context): event_record = get_first_record(event) - event_source_arn = event_record["eventSourceARN"] - event_id = event_record["eventID"] + event_source_arn = event_record.get("eventSourceARN") + event_id = event_record.get("eventID") stream_name = event_source_arn.split(":")[-1] shard_id = event_id.split(":")[0] tags = { @@ -801,12 +802,12 @@ def create_inferred_span_from_kinesis_event(event, context): "shardid": shard_id, "event_source_arn": event_source_arn, "event_id": event_id, - "event_name": event_record["eventName"], - "event_version": event_record["eventVersion"], - "partition_key": event_record["kinesis"]["partitionKey"], + "event_name": event_record.get("eventName"), + "event_version": event_record.get("eventVersion"), + "partition_key": event_record.get("kinesis").get("partitionKey"), } InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") - request_time_epoch = event_record["kinesis"]["approximateArrivalTimestamp"] + request_time_epoch = event_record.get("kinesis").get("approximateArrivalTimestamp") args = { "service": "kinesis", @@ -823,22 +824,22 @@ def create_inferred_span_from_kinesis_event(event, context): def create_inferred_span_from_dynamodb_event(event, context): event_record = get_first_record(event) - event_source_arn = event_record["eventSourceARN"] + event_source_arn = event_record.get("eventSourceARN") table_name = event_source_arn.split("/")[1] - dynamodb_message = event_record["dynamodb"] + dynamodb_message = event_record.get("dynamodb") tags = { "operation_name": "aws.dynamodb", "resource_names": table_name, "tablename": table_name, "event_source_arn": event_source_arn, - "event_id": event_record["eventID"], - "event_name": event_record["eventName"], - "event_version": event_record["eventVersion"], - "stream_view_type": dynamodb_message["StreamViewType"], - "size_bytes": str(dynamodb_message["SizeBytes"]), + "event_id": event_record.get("eventID"), + "event_name": event_record.get("eventName"), + "event_version": event_record.get("eventVersion"), + "stream_view_type": dynamodb_message.get("StreamViewType"), + "size_bytes": str(dynamodb_message.get("SizeBytes")), } InferredSpanInfo.set_tags(tags, synchronicity="async", tag_source="self") - request_time_epoch = event_record["dynamodb"]["ApproximateCreationDateTime"] + request_time_epoch = event_record.get("dynamodb").get("ApproximateCreationDateTime") args = { "service": "dynamodb", "resource": table_name, @@ -855,20 +856,20 @@ def create_inferred_span_from_dynamodb_event(event, context): def create_inferred_span_from_s3_event(event, context): event_record = get_first_record(event) - bucket_name = event_record["s3"]["bucket"]["name"] + bucket_name = event_record.get("s3").get("bucket").get("name") tags = { "operation_name": "aws.s3", "resource_names": bucket_name, - "event_name": event_record["eventName"], + "event_name": event_record.get("eventName"), "bucketname": bucket_name, - "bucket_arn": event_record["s3"]["bucket"]["arn"], - "object_key": event_record["s3"]["object"]["key"], - "object_size": str(event_record["s3"]["object"]["size"]), - "object_etag": event_record["s3"]["object"]["eTag"], + "bucket_arn": event_record.get("s3").get("bucket").get("arn"), + "object_key": event_record.get("s3").get("object").get("key"), + "object_size": str(event_record.get("s3").get("object").get("size")), + "object_etag": event_record.get("s3").get("object").get("eTag"), } InferredSpanInfo.set_tags(tags, synchronicity="async", tag_source="self") dt_format = "%Y-%m-%dT%H:%M:%S.%fZ" - timestamp = event_record["eventTime"] + timestamp = event_record.get("eventTime") dt = datetime.strptime(timestamp, dt_format) args = { @@ -885,11 +886,11 @@ def create_inferred_span_from_s3_event(event, context): def create_inferred_span_from_eventbridge_event(event, context): - source = event["source"] + source = event.get("source") tags = { "operation_name": "aws.eventbridge", "resource_names": source, - "detail_type": event["detail-type"], + "detail_type": event.get("detail-type"), } InferredSpanInfo.set_tags( tags, @@ -897,7 +898,7 @@ def create_inferred_span_from_eventbridge_event(event, context): tag_source="self", ) dt_format = "%Y-%m-%dT%H:%M:%SZ" - timestamp = event["time"] + timestamp = event.get("time") dt = datetime.strptime(timestamp, dt_format) args = { diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index a2d2792e..86f2100a 100644 --- a/datadog_lambda/trigger.py +++ b/datadog_lambda/trigger.py @@ -312,8 +312,8 @@ def extract_trigger_tags(event: dict, context: Any) -> dict: def extract_http_status_code_tag(trigger_tags, response): """ - If the Lambda was triggered by API Gateway, Lambda Function URL, or ALB add the returned status code - as a tag to the function execution span. + If the Lambda was triggered by API Gateway, Lambda Function URL, or ALB, + add the returned status code as a tag to the function execution span. """ if trigger_tags is None: return diff --git a/scripts/check_format.sh b/scripts/check_format.sh index 984b588f..df5a6718 100755 --- a/scripts/check_format.sh +++ b/scripts/check_format.sh @@ -2,7 +2,7 @@ set -e PYTHON_VERSION=$(python -c 'import sys; print(sys.version_info.major)') -pip install -Iv black==21.5b2 +pip install -Iv black==22.3.0 python -m black --check datadog_lambda/ --diff python -m black --check tests --diff diff --git a/tests/event_samples/api-gateway-no-apiid.json b/tests/event_samples/api-gateway-no-apiid.json new file mode 100644 index 00000000..0b8f91f1 --- /dev/null +++ b/tests/event_samples/api-gateway-no-apiid.json @@ -0,0 +1,126 @@ +{ + "body": "eyJ0ZXN0IjoiYm9keSJ9", + "resource": "/{proxy+}", + "path": "/path/to/resource", + "httpMethod": "POST", + "isBase64Encoded": true, + "queryStringParameters": { + "foo": "bar" + }, + "multiValueQueryStringParameters": { + "foo": [ + "bar" + ] + }, + "pathParameters": { + "proxy": "/path/to/resource" + }, + "stageVariables": { + "baz": "qux" + }, + "headers": { + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8", + "Accept-Encoding": "gzip, deflate, sdch", + "Accept-Language": "en-US,en;q=0.8", + "Cache-Control": "max-age=0", + "CloudFront-Forwarded-Proto": "https", + "CloudFront-Is-Desktop-Viewer": "true", + "CloudFront-Is-Mobile-Viewer": "false", + "CloudFront-Is-SmartTV-Viewer": "false", + "CloudFront-Is-Tablet-Viewer": "false", + "CloudFront-Viewer-Country": "US", + "Host": "1234567890.execute-api.us-east-1.amazonaws.com", + "Upgrade-Insecure-Requests": "1", + "User-Agent": "Custom User Agent String", + "Via": "1.1 08f323deadbeefa7af34d5feb414ce27.cloudfront.net (CloudFront)", + "X-Amz-Cf-Id": "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA==", + "X-Forwarded-For": "127.0.0.1, 127.0.0.2", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https", + "X-Datadog-Trace-Id": "12345", + "X-Datadog-Parent-Id": "67890", + "x-datadog-sampling-priority": "2" + }, + "multiValueHeaders": { + "Accept": [ + "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8" + ], + "Accept-Encoding": [ + "gzip, deflate, sdch" + ], + "Accept-Language": [ + "en-US,en;q=0.8" + ], + "Cache-Control": [ + "max-age=0" + ], + "CloudFront-Forwarded-Proto": [ + "https" + ], + "CloudFront-Is-Desktop-Viewer": [ + "true" + ], + "CloudFront-Is-Mobile-Viewer": [ + "false" + ], + "CloudFront-Is-SmartTV-Viewer": [ + "false" + ], + "CloudFront-Is-Tablet-Viewer": [ + "false" + ], + "CloudFront-Viewer-Country": [ + "US" + ], + "Host": [ + "0123456789.execute-api.us-east-1.amazonaws.com" + ], + "Upgrade-Insecure-Requests": [ + "1" + ], + "User-Agent": [ + "Custom User Agent String" + ], + "Via": [ + "1.1 08f323deadbeefa7af34d5feb414ce27.cloudfront.net (CloudFront)" + ], + "X-Amz-Cf-Id": [ + "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA==" + ], + "X-Forwarded-For": [ + "127.0.0.1, 127.0.0.2" + ], + "X-Forwarded-Port": [ + "443" + ], + "X-Forwarded-Proto": [ + "https" + ] + }, + "requestContext": { + "accountId": "123456789012", + "resourceId": "123456", + "stage": "prod", + "requestId": "c6af9ac6-7b61-11e6-9a41-93e8deadbeef", + "requestTime": "09/Apr/2015:12:34:56 +0000", + "requestTimeEpoch": 1428582896000, + "identity": { + "cognitoIdentityPoolId": null, + "accountId": null, + "cognitoIdentityId": null, + "caller": null, + "accessKey": null, + "sourceIp": "127.0.0.1", + "cognitoAuthenticationType": null, + "cognitoAuthenticationProvider": null, + "userArn": null, + "userAgent": "Custom User Agent String", + "user": null + }, + "domainName": "70ixmpl4fl.execute-api.us-east-2.amazonaws.com", + "path": "/prod/path/to/resource", + "resourcePath": "/{proxy+}", + "httpMethod": "POST", + "protocol": "HTTP/1.1" + } +} diff --git a/tests/test_tag_object.py b/tests/test_tag_object.py index cdb2941c..67622afe 100644 --- a/tests/test_tag_object.py +++ b/tests/test_tag_object.py @@ -70,7 +70,7 @@ def test_json_tag_object(self): def test_unicode_tag_object(self): payload = { "token": "world", - u"jsonString": u'{"stringifyThisJson":[{"here":"is","an":"object","number":1}]}', + "jsonString": '{"stringifyThisJson":[{"here":"is","an":"object","number":1}]}', } spanMock = MagicMock() tag_object(spanMock, "function.request", payload) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0d5afe26..be52697e 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1155,3 +1155,39 @@ def test_extract_context_from_kinesis_batch_event(self): self.assertEqual(context["trace-id"], "4948377316357291421") self.assertEqual(context["parent-id"], "2876253380018681026") self.assertEqual(context["sampling-priority"], "1") + + def test_create_inferred_span_from_api_gateway_event_no_apiid(self): + event_sample_source = "api-gateway-no-apiid" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "123" + print(event) + span = create_inferred_span(event, ctx) + print(span) + self.assertEqual(span.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual( + span.service, + "70ixmpl4fl.execute-api.us-east-2.amazonaws.com", + ) + self.assertEqual( + span.get_tag("http.url"), + "70ixmpl4fl.execute-api.us-east-2.amazonaws.com/path/to/resource", + ) + self.assertEqual(span.get_tag("endpoint"), "/path/to/resource") + self.assertEqual(span.get_tag("http.method"), "POST") + self.assertEqual( + span.get_tag("resource_names"), + "POST /path/to/resource", + ) + self.assertEqual( + span.get_tag("request_id"), "c6af9ac6-7b61-11e6-9a41-93e8deadbeef" + ) + self.assertEqual(span.get_tag("apiid"), "None") + self.assertEqual(span.get_tag("apiname"), "None") + self.assertEqual(span.get_tag("stage"), "prod") + self.assertEqual(span.start, 1428582896.0) + self.assertEqual(span.span_type, "http") + self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") + self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync")