|
| 1 | +from itertools import chain |
| 2 | +import logging |
| 3 | +from typing import List |
| 4 | + |
| 5 | +from ddtrace._trace.utils_botocore.span_pointers import ( |
| 6 | + _aws_s3_object_span_pointer_description, |
| 7 | +) |
| 8 | +from ddtrace._trace._span_pointer import _SpanPointerDirection |
| 9 | +from ddtrace._trace._span_pointer import _SpanPointerDescription |
| 10 | +from datadog_lambda.trigger import EventTypes |
| 11 | + |
| 12 | + |
| 13 | +logger = logging.getLogger(__name__) |
| 14 | + |
| 15 | + |
| 16 | +def calculate_span_pointers( |
| 17 | + event_source, |
| 18 | + event, |
| 19 | +) -> List[_SpanPointerDescription]: |
| 20 | + try: |
| 21 | + if event_source.equals(EventTypes.S3): |
| 22 | + return _calculate_s3_span_pointers_for_event(event) |
| 23 | + |
| 24 | + except Exception as e: |
| 25 | + logger.warning( |
| 26 | + "failed to calculate span pointers for event: %s", |
| 27 | + str(e), |
| 28 | + ) |
| 29 | + |
| 30 | + return [] |
| 31 | + |
| 32 | + |
| 33 | +def _calculate_s3_span_pointers_for_event(event) -> List[_SpanPointerDescription]: |
| 34 | + # Example event: |
| 35 | + # https://docs.aws.amazon.com/lambda/latest/dg/with-s3.html |
| 36 | + |
| 37 | + return list( |
| 38 | + chain.from_iterable( |
| 39 | + _calculate_s3_span_pointers_for_event_record(record) |
| 40 | + for record in event.get("Records", []) |
| 41 | + ) |
| 42 | + ) |
| 43 | + |
| 44 | + |
| 45 | +def _calculate_s3_span_pointers_for_event_record( |
| 46 | + record, |
| 47 | +) -> List[_SpanPointerDescription]: |
| 48 | + # Event types: |
| 49 | + # https://docs.aws.amazon.com/AmazonS3/latest/userguide/notification-how-to-event-types-and-destinations.html |
| 50 | + |
| 51 | + if record.get("eventName").startswith("ObjectCreated:"): |
| 52 | + s3_information = record.get("s3", None) |
| 53 | + if s3_information is not None: |
| 54 | + return _calculate_s3_span_pointers_for_object_created_s3_information( |
| 55 | + s3_information |
| 56 | + ) |
| 57 | + |
| 58 | + return [] |
| 59 | + |
| 60 | + |
| 61 | +def _calculate_s3_span_pointers_for_object_created_s3_information( |
| 62 | + s3_information, |
| 63 | +) -> List[_SpanPointerDescription]: |
| 64 | + try: |
| 65 | + bucket = s3_information["bucket"]["name"] |
| 66 | + key = s3_information["object"]["key"] |
| 67 | + etag = s3_information["object"]["eTag"] |
| 68 | + |
| 69 | + except KeyError as e: |
| 70 | + logger.warning( |
| 71 | + "missing s3 information required to make a span pointer: %s", |
| 72 | + str(e), |
| 73 | + ) |
| 74 | + return [] |
| 75 | + |
| 76 | + try: |
| 77 | + return [ |
| 78 | + _aws_s3_object_span_pointer_description( |
| 79 | + pointer_direction=_SpanPointerDirection.UPSTREAM, |
| 80 | + bucket=bucket, |
| 81 | + key=key, |
| 82 | + etag=etag, |
| 83 | + ) |
| 84 | + ] |
| 85 | + |
| 86 | + except Exception as e: |
| 87 | + logger.warning( |
| 88 | + "failed to generate S3 span pointer: %s", |
| 89 | + str(e), |
| 90 | + ) |
| 91 | + return [] |
0 commit comments