|
| 1 | +import json |
| 2 | +from typing import Any # noqa:F401 |
| 3 | +from typing import Dict # noqa:F401 |
| 4 | + |
| 5 | +import botocore.exceptions |
| 6 | + |
| 7 | +from ddtrace import Span # noqa:F401 |
| 8 | +from ddtrace import config |
| 9 | +from ddtrace.ext import http |
| 10 | +from ddtrace.propagation.http import HTTPPropagator |
| 11 | + |
| 12 | +from ....ext import SpanTypes |
| 13 | +from ....internal.logger import get_logger |
| 14 | +from ....internal.schema import SpanDirection |
| 15 | +from ....internal.schema import schematize_cloud_messaging_operation |
| 16 | +from ....internal.schema import schematize_service_name |
| 17 | +from ..utils import set_patched_api_call_span_tags |
| 18 | +from ..utils import set_response_metadata_tags |
| 19 | + |
| 20 | + |
| 21 | +log = get_logger(__name__) |
| 22 | + |
| 23 | + |
| 24 | +def inject_trace_to_stepfunction_input(params, span): |
| 25 | + # type: (Any, Span) -> None |
| 26 | + """ |
| 27 | + :params: contains the params for the current botocore action |
| 28 | + :span: the span which provides the trace context to be propagated |
| 29 | +
|
| 30 | + Inject the trace headers into the StepFunction input if the input is a JSON string |
| 31 | + """ |
| 32 | + if "input" not in params: |
| 33 | + log.warning("Unable to inject context. The StepFunction input had no input.") |
| 34 | + return |
| 35 | + |
| 36 | + if params["input"] is None: |
| 37 | + log.warning("Unable to inject context. The StepFunction input was None.") |
| 38 | + return |
| 39 | + |
| 40 | + elif isinstance(params["input"], dict): |
| 41 | + if "_datadog" in params["input"]: |
| 42 | + log.warning("Input already has trace context.") |
| 43 | + return |
| 44 | + params["input"]["_datadog"] = {} |
| 45 | + HTTPPropagator.inject(span.context, params["input"]["_datadog"]) |
| 46 | + return |
| 47 | + |
| 48 | + elif isinstance(params["input"], str): |
| 49 | + try: |
| 50 | + input_obj = json.loads(params["input"]) |
| 51 | + except ValueError: |
| 52 | + log.warning("Input is not a valid JSON string") |
| 53 | + return |
| 54 | + |
| 55 | + if isinstance(input_obj, dict): |
| 56 | + input_obj["_datadog"] = {} |
| 57 | + HTTPPropagator.inject(span.context, input_obj["_datadog"]) |
| 58 | + input_json = json.dumps(input_obj) |
| 59 | + |
| 60 | + params["input"] = input_json |
| 61 | + return |
| 62 | + else: |
| 63 | + log.warning("Unable to inject context. The StepFunction input was not a dict.") |
| 64 | + return |
| 65 | + |
| 66 | + else: |
| 67 | + log.warning("Unable to inject context. The StepFunction input was not a dict or a JSON string.") |
| 68 | + |
| 69 | + |
| 70 | +def patched_stepfunction_api_call(original_func, instance, args, kwargs: Dict, function_vars: Dict): |
| 71 | + params = function_vars.get("params") |
| 72 | + trace_operation = function_vars.get("trace_operation") |
| 73 | + pin = function_vars.get("pin") |
| 74 | + endpoint_name = function_vars.get("endpoint_name") |
| 75 | + operation = function_vars.get("operation") |
| 76 | + |
| 77 | + with pin.tracer.trace( |
| 78 | + trace_operation, |
| 79 | + service=schematize_service_name("{}.{}".format(pin.service, endpoint_name)), |
| 80 | + span_type=SpanTypes.HTTP, |
| 81 | + ) as span: |
| 82 | + set_patched_api_call_span_tags(span, instance, args, params, endpoint_name, operation) |
| 83 | + |
| 84 | + if args: |
| 85 | + if config.botocore["distributed_tracing"]: |
| 86 | + try: |
| 87 | + if endpoint_name == "states" and operation in {"StartExecution", "StartSyncExecution"}: |
| 88 | + inject_trace_to_stepfunction_input(params, span) |
| 89 | + span.name = schematize_cloud_messaging_operation( |
| 90 | + trace_operation, |
| 91 | + cloud_provider="aws", |
| 92 | + cloud_service="stepfunctions", |
| 93 | + direction=SpanDirection.OUTBOUND, |
| 94 | + ) |
| 95 | + except Exception: |
| 96 | + log.warning("Unable to inject trace context", exc_info=True) |
| 97 | + |
| 98 | + try: |
| 99 | + return original_func(*args, **kwargs) |
| 100 | + except botocore.exceptions.ClientError as e: |
| 101 | + set_response_metadata_tags(span, e.response) |
| 102 | + |
| 103 | + # If we have a status code, and the status code is not an error, |
| 104 | + # then ignore the exception being raised |
| 105 | + status_code = span.get_tag(http.STATUS_CODE) |
| 106 | + if status_code and not config.botocore.operations[span.resource].is_error_code(int(status_code)): |
| 107 | + span._ignore_exception(botocore.exceptions.ClientError) |
| 108 | + raise |
0 commit comments