Skip to content

[SLS-2141] Fix keyError issue when parsing event source #226

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hghotra merged 1 commit into from
May 9, 2022
Merged

[SLS-2141] Fix keyError issue when parsing event source #226

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 24 additions & 24 deletions datadog_lambda/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -553,8 +553,8 @@ def create_inferred_span(event, context):
def create_inferred_span_from_lambda_function_url_event(event, context):
request_context = event.get("requestContext")
domain = request_context.get("domainName")
method = request_context.get("http").get("method")
path = request_context.get("http").get("path")
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",
Expand Down Expand Up @@ -583,11 +583,7 @@ 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.get("headers")
and event.get("headers").get("X-Amz-Invocation-Type") == "Event"
)
return event.get("headers", {}).get("X-Amz-Invocation-Type") == "Event"


def create_inferred_span_from_api_gateway_websocket_event(event, context):
Expand Down Expand Up @@ -663,17 +659,17 @@ def create_inferred_span_from_api_gateway_event(event, context):
def create_inferred_span_from_http_api_event(event, context):
request_context = event.get("requestContext")
domain = request_context.get("domainName")
method = request_context.get("http").get("method")
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.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"),
"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.get("apiId"),
Expand Down Expand Up @@ -710,10 +706,10 @@ def create_inferred_span_from_sqs_event(event, context):
"queuename": queue_name,
"event_source_arn": event_source_arn,
"receipt_handle": event_record.get("receiptHandle"),
"sender_id": event_record.get("attributes").get("SenderId"),
"sender_id": event_record.get("attributes", {}).get("SenderId"),
}
InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async")
request_time_epoch = event_record.get("attributes").get("SentTimestamp")
request_time_epoch = event_record.get("attributes", {}).get("SentTimestamp")
args = {
"service": "sqs",
"resource": queue_name,
Expand Down Expand Up @@ -756,7 +752,7 @@ 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.get("Sns")
topic_arn = event_record.get("Sns").get("TopicArn")
topic_arn = event_record.get("Sns", {}).get("TopicArn")
topic_name = topic_arn.split(":")[-1]
tags = {
"operation_name": "aws.sns",
Expand All @@ -773,7 +769,7 @@ def create_inferred_span_from_sns_event(event, context):

InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async")
sns_dt_format = "%Y-%m-%dT%H:%M:%S.%fZ"
timestamp = event_record.get("Sns").get("Timestamp")
timestamp = event_record.get("Sns", {}).get("Timestamp")
dt = datetime.strptime(timestamp, sns_dt_format)

args = {
Expand Down Expand Up @@ -804,10 +800,12 @@ def create_inferred_span_from_kinesis_event(event, context):
"event_id": event_id,
"event_name": event_record.get("eventName"),
"event_version": event_record.get("eventVersion"),
"partition_key": event_record.get("kinesis").get("partitionKey"),
"partition_key": event_record.get("kinesis", {}).get("partitionKey"),
}
InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async")
request_time_epoch = event_record.get("kinesis").get("approximateArrivalTimestamp")
request_time_epoch = event_record.get("kinesis", {}).get(
"approximateArrivalTimestamp"
)

args = {
"service": "kinesis",
Expand Down Expand Up @@ -839,7 +837,9 @@ def create_inferred_span_from_dynamodb_event(event, context):
"size_bytes": str(dynamodb_message.get("SizeBytes")),
}
InferredSpanInfo.set_tags(tags, synchronicity="async", tag_source="self")
request_time_epoch = event_record.get("dynamodb").get("ApproximateCreationDateTime")
request_time_epoch = event_record.get("dynamodb", {}).get(
"ApproximateCreationDateTime"
)
args = {
"service": "dynamodb",
"resource": table_name,
Expand All @@ -856,16 +856,16 @@ 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.get("s3").get("bucket").get("name")
bucket_name = event_record.get("s3", {}).get("bucket", {}).get("name")
tags = {
"operation_name": "aws.s3",
"resource_names": bucket_name,
"event_name": event_record.get("eventName"),
"bucketname": bucket_name,
"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"),
"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"
Expand Down
20 changes: 11 additions & 9 deletions datadog_lambda/trigger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ def parse_event_source(event: dict) -> _EventSource:
event_source.subtype = EventSubtypes.API_GATEWAY
if "routeKey" in event:
event_source.subtype = EventSubtypes.HTTP_API
if "requestContext" in event and "messageDirection" in event["requestContext"]:
if event.get("requestContext", {}).get("messageDirection"):
event_source.subtype = EventSubtypes.WEBSOCKET

if request_context and request_context.get("elb"):
Expand Down Expand Up @@ -186,15 +186,17 @@ def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> s
event_record = get_first_record(event)
# e.g. arn:aws:s3:::lambda-xyz123-abc890
if source.to_string() == "s3":
return event_record.get("s3")["bucket"]["arn"]
return event_record.get("s3", {}).get("bucket", {}).get("arn")

# e.g. arn:aws:sns:us-east-1:123456789012:sns-lambda
if source.to_string() == "sns":
return event_record.get("Sns")["TopicArn"]
return event_record.get("Sns", {}).get("TopicArn")

# e.g. arn:aws:cloudfront::123456789012:distribution/ABC123XYZ
if source.event_type == EventTypes.CLOUDFRONT:
distribution_id = event_record.get("cf")["config"]["distributionId"]
distribution_id = (
event_record.get("cf", {}).get("config", {}).get("distributionId")
)
return "arn:{}:cloudfront::{}:distribution/{}".format(
aws_arn, account_id, distribution_id
)
Expand All @@ -215,18 +217,18 @@ def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> s
if source.event_type == EventTypes.API_GATEWAY:
request_context = event.get("requestContext")
return "arn:{}:apigateway:{}::/restapis/{}/stages/{}".format(
aws_arn, region, request_context["apiId"], request_context["stage"]
aws_arn, region, request_context.get("apiId"), request_context.get("stage")
)

# e.g. arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/lambda-xyz/123
if source.event_type == EventTypes.ALB:
request_context = event.get("requestContext")
return request_context.get("elb")["targetGroupArn"]
return request_context.get("elb", {}).get("targetGroupArn")

# e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz
if source.event_type == EventTypes.CLOUDWATCH_LOGS:
with gzip.GzipFile(
fileobj=BytesIO(base64.b64decode(event["awslogs"]["data"]))
fileobj=BytesIO(base64.b64decode(event.get("awslogs", {}).get("data")))
) as decompress_stream:
data = b"".join(BufferedReader(decompress_stream))
logs = json.loads(data)
Expand Down Expand Up @@ -265,7 +267,7 @@ def extract_http_tags(event):
method = event.get("httpMethod")
if request_context and request_context.get("stage"):
if request_context.get("domainName"):
http_tags["http.url"] = request_context["domainName"]
http_tags["http.url"] = request_context.get("domainName")

path = request_context.get("path")
method = request_context.get("httpMethod")
Expand All @@ -282,7 +284,7 @@ def extract_http_tags(event):

headers = event.get("headers")
if headers and headers.get("Referer"):
http_tags["http.referer"] = headers["Referer"]
http_tags["http.referer"] = headers.get("Referer")

return http_tags

Expand Down