diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index 45c1a68fae..53b0a5c412 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -375,6 +375,18 @@ class SPANDATA: Example: "arn:aws:lambda:us-east-1:123456789012:function:my-function" """ + AWS_LAMBDA_RESOURCE_MAPPING_ID = "aws.lambda.resource_mapping.id" + """ + The UUID of the AWS Lambda event source mapping the request refers to. + Example: "587b4f6c-1d0e-4d48-8a3f-2c6f9a1e0b7d" + """ + + AWS_LAMBDA_TARGET_FUNCTION_NAME = "aws.lambda.target_function.name" + """ + The name of the AWS Lambda function an AWS SDK call acts on, without a qualifier. + Example: "billing" + """ + AWS_LOG_GROUP_NAMES = "aws.log.group.names" """ The name(s) of the AWS log group(s) an application is writing to. @@ -670,6 +682,24 @@ class SPANDATA: Example: "af9d5aa4-a685-4c5f-a22b-444f80b3cc28" """ + FAAS_INVOKED_NAME = "faas.invoked_name" + """ + The name of the invoked function, without a qualifier. + Example: "billing" + """ + + FAAS_INVOKED_PROVIDER = "faas.invoked_provider" + """ + The cloud provider of the invoked function. + Example: "aws" + """ + + FAAS_INVOKED_REGION = "faas.invoked_region" + """ + The cloud region of the invoked function. + Example: "eu-north-1" + """ + FAAS_NAME = "faas.name" """ The name of the serverless function. diff --git a/sentry_sdk/integrations/boto3/_client.py b/sentry_sdk/integrations/boto3/_client.py index 35ea846e22..25af44c02d 100644 --- a/sentry_sdk/integrations/boto3/_client.py +++ b/sentry_sdk/integrations/boto3/_client.py @@ -96,6 +96,9 @@ def sentry_patched_make_api_call( ctx.add_metadata(self) service_ext = _resolve_service(ctx.service_name) span = _start_client_span(ctx, service_ext) + with capture_internal_exceptions(): + # override the request params for e.g. trace propagation. + api_params = service_ext.get_request_params(ctx, span) or api_params if span is None: return orig_make_api_call(self, operation_name, api_params) @@ -107,7 +110,7 @@ def sentry_patched_make_api_call( try: parsed = orig_make_api_call(self, operation_name, api_params) except BaseException as error: - if service_ext is not None and isinstance(error, ClientError): + if isinstance(error, ClientError): with capture_internal_exceptions(): attributes.update( service_ext.get_response_attributes(ctx, error.response) @@ -117,11 +120,8 @@ def sentry_patched_make_api_call( with capture_internal_exceptions(): span.set_attributes(attributes) raise - if service_ext is not None: - with capture_internal_exceptions(): - attributes.update( - service_ext.get_response_attributes(ctx, parsed) - ) + with capture_internal_exceptions(): + attributes.update(service_ext.get_response_attributes(ctx, parsed)) with capture_internal_exceptions(): attributes.update(_get_response_attributes(parsed)) with capture_internal_exceptions(): @@ -136,7 +136,8 @@ def sentry_patched_make_api_call( streaming_body_instrumented = False with capture_internal_exceptions(): - streaming_body_instrumented = _instrument_streaming_body(span, parsed) + if service_ext.has_streaming_response(ctx): + streaming_body_instrumented = _instrument_streaming_body(span, parsed) # `StreamingBody`s finish their span when consumed or closed. if not streaming_body_instrumented: diff --git a/sentry_sdk/integrations/boto3/_instrumentation.py b/sentry_sdk/integrations/boto3/_instrumentation.py index 731dac6fa3..777780fa50 100644 --- a/sentry_sdk/integrations/boto3/_instrumentation.py +++ b/sentry_sdk/integrations/boto3/_instrumentation.py @@ -148,7 +148,7 @@ def _get_error_attributes(exception: "BaseException") -> "Attributes": def _start_client_span( ctx: "AwsCallContext", - service_ext: "Optional[_ServiceExtension]" = None, + service_ext: "_ServiceExtension", ) -> "Optional[Span]": if sentry_sdk.get_current_span() is None: return None @@ -171,14 +171,13 @@ def _start_client_span( # with `get_span_op()`. span_op = OP.RPC - if service_ext is not None: - with capture_internal_exceptions(): - service_op = service_ext.get_span_op(ctx) - if service_op is not None: - span_op = service_op + with capture_internal_exceptions(): + service_op = service_ext.get_span_op(ctx) + if service_op is not None: + span_op = service_op - with capture_internal_exceptions(): - attributes.update(service_ext.get_request_attributes(ctx)) + with capture_internal_exceptions(): + attributes.update(service_ext.get_request_attributes(ctx)) # `start_span()` evaluates `ignore_spans` against the initial attributes. # https://opentelemetry.io/docs/specs/semconv/rpc/rpc-spans/#rpc-client-span diff --git a/sentry_sdk/integrations/boto3/_services/base.py b/sentry_sdk/integrations/boto3/_services/base.py index ec7ac1df1c..910e7a2e0e 100644 --- a/sentry_sdk/integrations/boto3/_services/base.py +++ b/sentry_sdk/integrations/boto3/_services/base.py @@ -1,18 +1,19 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Dict, Optional from sentry_sdk._types import Attributes from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.traces import Span class _ServiceExtension: """ Optional hooks for adding service-specific behavior to AWS client span; non-overridden methods keep the generic instrumentation. - Services without a registered extension in `_SERVICE_EXTENSIONS` continue - to use the generic instrumentation. + Services without a registered extension in `_SERVICE_EXTENSIONS` use this + class as is, i.e. the generic instrumentation. """ __slots__ = () @@ -30,3 +31,19 @@ def get_response_attributes( ) -> "Attributes": """Return response attributes to add after the AWS request is made.""" return {} + + def get_request_params( + self, ctx: "AwsCallContext", span: "Optional[Span]" + ) -> "Optional[Dict[str, Any]]": + """ + Return params to send instead of `ctx.params`, e.g. with trace propagation added, + or `None` to send them unchanged; must not mutate `ctx.params`. + """ + return None + + def has_streaming_response(self, ctx: "AwsCallContext") -> bool: + """ + Return whether a `StreamingBody` response is read over time, so the span ends + when it's read or closed, or `False` to end it with the call. + """ + return True diff --git a/sentry_sdk/integrations/boto3/_services/lambda_.py b/sentry_sdk/integrations/boto3/_services/lambda_.py index c40b969762..9d1e65bcad 100644 --- a/sentry_sdk/integrations/boto3/_services/lambda_.py +++ b/sentry_sdk/integrations/boto3/_services/lambda_.py @@ -1,22 +1,113 @@ +import base64 +import json from typing import TYPE_CHECKING +from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations.boto3._services.base import _ServiceExtension -from sentry_sdk.integrations.boto3._utils import _get_aws_arn_attributes +from sentry_sdk.integrations.boto3._utils import ( + _get_aws_arn_attributes, + _get_propagation_headers, +) +from sentry_sdk.integrations.boto3.consts import CLOUD_PROVIDER from sentry_sdk.utils import capture_internal_exceptions if TYPE_CHECKING: + from typing import Any, Dict, Optional + from sentry_sdk._types import Attributes from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.traces import Span class _LambdaExtension(_ServiceExtension): def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": attributes: "Attributes" = {} + + # event source mapping apis. + if "UUID" in ctx.params: + attributes[SPANDATA.AWS_LAMBDA_RESOURCE_MAPPING_ID] = ctx.params["UUID"] + if "FunctionName" not in ctx.params: + return attributes + + # a name, a partial arn (`:function:`) or an + # arn (`arn::lambda:::function:`) + # each optionally followed by `:`. + function_name = ctx.params["FunctionName"] + name = function_name.split(":function:")[-1].split(":", 1)[0] + + is_invoke = ctx.operation_name in ( + "Invoke", + "InvokeAsync", + "InvokeWithResponseStream", + ) + if is_invoke: + attributes[SPANDATA.FAAS_INVOKED_NAME] = name + attributes[SPANDATA.FAAS_INVOKED_PROVIDER] = CLOUD_PROVIDER + if ctx.region_name is not None: + attributes[SPANDATA.FAAS_INVOKED_REGION] = ctx.region_name + else: + attributes[SPANDATA.AWS_LAMBDA_TARGET_FUNCTION_NAME] = name + with capture_internal_exceptions(): - attributes.update( - _get_aws_arn_attributes( - (ctx.params.get("FunctionName"),), - ("lambda",), - ) - ) + arn_attributes = _get_aws_arn_attributes((function_name,), ("lambda",)) + attributes.update(arn_attributes) + if is_invoke and arn_attributes: + attributes[SPANDATA.FAAS_INVOKED_REGION] = function_name.split(":")[3] + return attributes + + def get_response_attributes( + self, ctx: "AwsCallContext", response: "Any" + ) -> "Attributes": + attributes: "Attributes" = {} + + if "UUID" in response: + attributes[SPANDATA.AWS_LAMBDA_RESOURCE_MAPPING_ID] = response["UUID"] + + # `Invoke` returns `200 OK` when the function fails; reading it would leave + # the caller with an empty body. + if "FunctionError" in response: + attributes[SPANDATA.ERROR_TYPE] = "lambda_function_error" + + return attributes + + def get_request_params( + self, ctx: "AwsCallContext", span: "Optional[Span]" + ) -> "Optional[Dict[str, Any]]": + # lambda only passes `ClientContext` to the function on synchronous invocations. + # https://docs.aws.amazon.com/lambda/latest/api/API_Invoke.html#lambda-Invoke-request-ClientContext + if ( + ctx.operation_name not in ("Invoke", "InvokeWithResponseStream") + or ctx.params.get("InvocationType", "RequestResponse") != "RequestResponse" + ): + return None + + headers = _get_propagation_headers(ctx, span) + + # https://docs.aws.amazon.com/lambda/latest/dg/python-context.html + client_context = ( + json.loads(base64.b64decode(ctx.params["ClientContext"])) + if "ClientContext" in ctx.params + else {} + ) + custom = client_context.get("custom", {}) + # never mix with or replace trace data the caller propagates itself. + if not headers or any(key in custom for key in headers): + return None + + client_context["custom"] = {**custom, **headers} + encoded = base64.b64encode( + json.dumps(client_context, separators=(",", ":")).encode() + ).decode() + + # lambda rejects more than 3,583 bytes of base64; skip propagation rather + # than fail the call. + if len(encoded) > 3583: + return None + + return {**ctx.params, "ClientContext": encoded} + + def has_streaming_response(self, ctx: "AwsCallContext") -> bool: + # `Invoke`'s `Payload` is a `StreamingBody`, but span + # shouldn't be kept open until the caller reads it. + return ctx.operation_name != "Invoke" diff --git a/sentry_sdk/integrations/boto3/_services/registry.py b/sentry_sdk/integrations/boto3/_services/registry.py index 84dba8b8f8..ace34aa023 100644 --- a/sentry_sdk/integrations/boto3/_services/registry.py +++ b/sentry_sdk/integrations/boto3/_services/registry.py @@ -1,5 +1,6 @@ from typing import TYPE_CHECKING +from sentry_sdk.integrations.boto3._services.base import _ServiceExtension from sentry_sdk.integrations.boto3._services.dynamodb import _DynamoDBExtension from sentry_sdk.integrations.boto3._services.kinesis import _KinesisExtension from sentry_sdk.integrations.boto3._services.lambda_ import _LambdaExtension @@ -16,14 +17,12 @@ if TYPE_CHECKING: from typing import Dict, Optional - from sentry_sdk.integrations.boto3._services.base import _ServiceExtension - # add a ServiceExtension here when one is implemented. for example: # _SERVICE_EXTENSIONS = {"s3": _S3Extension()} # when py 3.15 drops, we might want to take a look at using # a lazy-loading approach using the new `lazy` keyword. -_SERVICE_EXTENSIONS: "Dict[str, _ServiceExtension]" = { +_SERVICE_EXTENSIONS: "Dict[Optional[str], _ServiceExtension]" = { "dynamodb": _DynamoDBExtension(), "kinesis": _KinesisExtension(), "lambda": _LambdaExtension(), @@ -37,8 +36,6 @@ def _resolve_service( service_name: "Optional[str]", -) -> "Optional[_ServiceExtension]": - """Return the extension for a service, or `None` for generic instrumentation.""" - if service_name is None: - return None - return _SERVICE_EXTENSIONS.get(service_name) +) -> "_ServiceExtension": + """Return the extension for a service, or `_ServiceExtension()` for generic instrumentation.""" + return _SERVICE_EXTENSIONS.get(service_name, _ServiceExtension()) diff --git a/sentry_sdk/integrations/boto3/_utils.py b/sentry_sdk/integrations/boto3/_utils.py index 9f124639a6..7567547587 100644 --- a/sentry_sdk/integrations/boto3/_utils.py +++ b/sentry_sdk/integrations/boto3/_utils.py @@ -1,11 +1,16 @@ from typing import TYPE_CHECKING +import sentry_sdk from sentry_sdk.consts import SPANDATA +from sentry_sdk.traces import NoOpSpan +from sentry_sdk.tracing_utils import should_propagate_trace if TYPE_CHECKING: from typing import Any, Dict, Iterable, Optional, Sequence, Tuple from sentry_sdk._types import Attributes + from sentry_sdk.integrations.boto3._context import AwsCallContext + from sentry_sdk.traces import Span # tuple of (request param name, span attribute name). _AttributeSpec = Tuple[str, str] @@ -38,3 +43,21 @@ def _get_aws_arn_attributes( return attributes return {} + + +def _get_propagation_headers( + ctx: "AwsCallContext", span: "Optional[Span]" +) -> "Dict[str, str]": + """Return the `sentry-trace` and `baggage` headers a call propagates to its consumer.""" + if ctx.endpoint_url is None or not should_propagate_trace( + sentry_sdk.get_client(), ctx.endpoint_url + ): + return {} + + # `NoOpSpan` is never sent, use scope's span instead. + if isinstance(span, NoOpSpan): + span = None + + return dict( + sentry_sdk.get_current_scope().iter_trace_propagation_headers(span=span) + ) diff --git a/tests/integrations/boto3/test_client.py b/tests/integrations/boto3/test_client.py index 999d3876f4..f8ccb8261d 100644 --- a/tests/integrations/boto3/test_client.py +++ b/tests/integrations/boto3/test_client.py @@ -177,39 +177,11 @@ def record_client_span(request, **kwargs): assert span.get("end_timestamp") is not None -@pytest.mark.parametrize( - "service_name,method_name,api_params,payload_field", - [ - pytest.param( - "lambda", - "invoke", - {"FunctionName": "function"}, - "Payload", - id="payload", - ), - pytest.param( - "s3", - "get_object_annotation", - { - "Bucket": "bucket", - "Key": "file.txt", - "AnnotationName": "annotation", - }, - "AnnotationPayload", - id="annotation-payload", - ), - ], -) -def test_non_body_stream_delays_client_span( - capture_items, - client_factory, - service_name, - method_name, - api_params, - payload_field, -): - client = client_factory(service_name) - require_botocore_model_fields(client, method_name, output_fields=(payload_field,)) +def test_non_body_stream_delays_client_span(capture_items, client_factory): + client = client_factory("s3") + require_botocore_model_fields( + client, "get_object_annotation", output_fields=("AnnotationPayload",) + ) request_client_spans = [] def record_client_span(request, **kwargs): @@ -220,8 +192,10 @@ def record_client_span(request, **kwargs): def invoke(): parent = sentry_sdk.get_current_span() with MockResponse(client, 200, {"content-length": "1"}, b"x"): - response = getattr(client, method_name)(**api_params) - body = response[payload_field] + response = client.get_object_annotation( + Bucket="bucket", Key="file.txt", AnnotationName="annotation" + ) + body = response["AnnotationPayload"] assert isinstance(body, StreamingBody) (request_client_span,) = request_client_spans assert request_client_span.end_timestamp is None diff --git a/tests/integrations/boto3/test_lambda.py b/tests/integrations/boto3/test_lambda.py index 00c5d58df1..94d8697723 100644 --- a/tests/integrations/boto3/test_lambda.py +++ b/tests/integrations/boto3/test_lambda.py @@ -1,22 +1,226 @@ +import base64 +import io +import json +from contextlib import nullcontext +from copy import deepcopy + +import pytest +from botocore.response import StreamingBody from botocore.stub import Stubber +import sentry_sdk from sentry_sdk.consts import OP, SPANDATA -from tests.integrations.boto3.helpers import capture_spans_by_op +from sentry_sdk.integrations.boto3 import Boto3Integration +from tests.integrations.boto3.helpers import ( + capture_spans_by_op, + require_botocore_model_fields, +) from tests.integrations.boto3.helpers import client_factory as client_factory -def test_request_attributes(client_factory, capture_items): +def _stubbed_span(client_factory, capture_items, method, params, response): + """Return the span and the params botocore sent.""" client = client_factory("lambda") - function_arn = "arn:aws:lambda:eu-north-1:123456789012:function:orders" - + require_botocore_model_fields( + client, method, input_fields=tuple(params), output_fields=tuple(response) + ) + sent = [] + client.meta.events.register( + "before-parameter-build", lambda params, **kwargs: sent.append(params) + ) with Stubber(client) as stubber: - stubber.add_response("delete_function", {}, {"FunctionName": function_arn}) + stubber.add_response(method, response) spans = capture_spans_by_op( - lambda: client.delete_function(**{"FunctionName": function_arn}), - capture_items, + lambda: getattr(client, method)(**params), capture_items ) - (span,) = spans[OP.RPC] + return span, sent[0] + + +@pytest.mark.parametrize( + "function_name,region,resource_id", + [ + # fall back to client's region. + pytest.param("billing:prod", "eu-north-1", None, id="name"), + pytest.param( + "arn:aws:lambda:us-east-1:123456789012:function:billing:7", + "us-east-1", + "arn:aws:lambda:us-east-1:123456789012:function:billing:7", + id="arn", + ), + ], +) +def test_invoke_attributes( + client_factory, capture_items, function_name, region, resource_id +): + span, _ = _stubbed_span( + client_factory, + capture_items, + "invoke", + {"FunctionName": function_name}, + {"StatusCode": 200}, + ) + + attributes = span["attributes"] + assert attributes[SPANDATA.FAAS_INVOKED_NAME] == "billing" + assert attributes[SPANDATA.FAAS_INVOKED_REGION] == region + assert attributes.get(SPANDATA.CLOUD_RESOURCE_ID) == resource_id + assert SPANDATA.AWS_LAMBDA_TARGET_FUNCTION_NAME not in attributes + + +def test_function_attributes(client_factory, capture_items): + arn = "arn:aws:lambda:eu-north-1:123456789012:function:billing" + span, _ = _stubbed_span( + client_factory, capture_items, "delete_function", {"FunctionName": arn}, {} + ) + attributes = span["attributes"] - assert attributes[SPANDATA.CLOUD_ACCOUNT_ID] == "123456789012" - assert attributes[SPANDATA.CLOUD_RESOURCE_ID] == function_arn + assert attributes[SPANDATA.AWS_LAMBDA_TARGET_FUNCTION_NAME] == "billing" + assert attributes[SPANDATA.CLOUD_RESOURCE_ID] == arn + assert SPANDATA.FAAS_INVOKED_NAME not in attributes + + +@pytest.mark.parametrize( + "method,params,response", + [ + pytest.param( + "create_event_source_mapping", + { + "FunctionName": "billing", + "EventSourceArn": "arn:aws:sqs:eu-north-1:123456789012:orders", + }, + {"UUID": "587b4f6c-1d0e-4d48-8a3f-2c6f9a1e0b7d"}, + id="from-response", + ), + pytest.param( + "get_event_source_mapping", + {"UUID": "587b4f6c-1d0e-4d48-8a3f-2c6f9a1e0b7d"}, + {}, + id="from-request", + ), + ], +) +def test_event_source_mapping_attributes( + client_factory, capture_items, method, params, response +): + span, _ = _stubbed_span(client_factory, capture_items, method, params, response) + + assert ( + span["attributes"][SPANDATA.AWS_LAMBDA_RESOURCE_MAPPING_ID] + == "587b4f6c-1d0e-4d48-8a3f-2c6f9a1e0b7d" + ) + + +def test_invoke_function_error(client_factory, capture_items): + span, _ = _stubbed_span( + client_factory, + capture_items, + "invoke", + {"FunctionName": "billing"}, + { + "StatusCode": 200, + "FunctionError": "Unhandled", + "Payload": StreamingBody(io.BytesIO(b"{}"), 2), + }, + ) + + assert span["status"] == "error" + assert span["attributes"][SPANDATA.ERROR_TYPE] == "lambda_function_error" + + +def test_invoke_propagates_trace(client_factory, capture_items): + client_context = { + "client": {"app_title": "shop"}, + "custom": {"customer_id": "1234567890"}, + } + params = { + "FunctionName": "billing", + "ClientContext": base64.b64encode(json.dumps(client_context).encode()).decode(), + } + original_params = deepcopy(params) + span, sent = _stubbed_span( + client_factory, capture_items, "invoke", params, {"StatusCode": 200} + ) + + sent_context = json.loads(base64.b64decode(sent["ClientContext"])) + custom = sent_context["custom"] + assert custom.pop("sentry-trace") == f"{span['trace_id']}-{span['span_id']}-1" + assert f"sentry-trace_id={span['trace_id']}" in custom.pop("baggage") + # caller's values and params should be left unchanged. + assert sent_context == client_context + assert params == original_params + + +@pytest.mark.parametrize( + "ignore_spans,parent", + [ + pytest.param([], nullcontext, id="no-parent"), + pytest.param( + ["Lambda.Invoke"], + lambda: sentry_sdk.start_span(name="parent"), + id="ignored-span", + ), + ], +) +def test_invoke_propagates_scope_trace( + sentry_init, client_factory, ignore_spans, parent +): + client = client_factory("lambda") + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + ignore_spans=ignore_spans, + ) + sent = [] + client.meta.events.register( + "before-parameter-build", lambda params, **kwargs: sent.append(params) + ) + with Stubber(client) as stubber: + stubber.add_response("invoke", {"StatusCode": 200}) + with parent(): + client.invoke(FunctionName="billing") + expected = dict( + sentry_sdk.get_current_scope().iter_trace_propagation_headers() + ) + + # the scope's trace is propagated instead. + (params,) = sent + custom = json.loads(base64.b64decode(params["ClientContext"]))["custom"] + assert custom["sentry-trace"] == expected["sentry-trace"] + + +@pytest.mark.parametrize( + "params", + [ + # lambda doesn't pass `ClientContext` to an async. invoked function. + pytest.param( + {"FunctionName": "billing", "InvocationType": "Event"}, id="async" + ), + # with trace data, it would exceed the 3,583 bytes lambda accepts. + pytest.param( + { + "FunctionName": "billing", + "ClientContext": base64.b64encode( + json.dumps({"custom": {"padding": "x" * 2600}}).encode() + ).decode(), + }, + id="too-large", + ), + # the caller propagates its own trace. + pytest.param( + { + "FunctionName": "billing", + "ClientContext": base64.b64encode( + json.dumps({"custom": {"sentry-trace": "set-by-caller"}}).encode() + ).decode(), + }, + id="caller-trace", + ), + ], +) +def test_invoke_without_propagation(client_factory, capture_items, params): + _, sent = _stubbed_span( + client_factory, capture_items, "invoke", params, {"StatusCode": 202} + ) + + assert sent == params