From 4498781d17d41971670e58a73e1781c07fcac9a3 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Tue, 7 Nov 2023 11:09:06 -0800 Subject: [PATCH 01/10] [SVLS-4126] Patch botocore stepfunctions --- ddtrace/contrib/botocore/patch.py | 10 ++++++++++ tests/contrib/botocore/test.py | 21 +++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/ddtrace/contrib/botocore/patch.py b/ddtrace/contrib/botocore/patch.py index 3dea81cffa6..9281e201a7d 100644 --- a/ddtrace/contrib/botocore/patch.py +++ b/ddtrace/contrib/botocore/patch.py @@ -577,6 +577,16 @@ def patched_api_call(original_func, instance, args, kwargs): span.name = schematize_cloud_messaging_operation( trace_operation, cloud_provider="aws", cloud_service="sns", direction=SpanDirection.OUTBOUND ) + if endpoint_name == "stepfunctions" and ( + operation == "start-execution" or operation == "start-sync-execution" + ): + # TODO: add inject_trace method + span.name = schematize_cloud_messaging_operation( + trace_operation, + cloud_provider="aws", + cloud_service="stepfunctions", + direction=SpanDirection.OUTBOUND, + ) except Exception: log.warning("Unable to inject trace context", exc_info=True) diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index 46a738f326a..7027cb17909 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -10,6 +10,7 @@ import botocore.session import mock from moto import mock_dynamodb +from moto import mock_stepfunctions from moto import mock_ec2 from moto import mock_events from moto import mock_kinesis @@ -885,6 +886,26 @@ def test_schematized_unspecified_service_sqs_client_v1(self): assert spans[2].service == DEFAULT_SPAN_SERVICE_NAME assert spans[2].name == "aws.sqs.receive" + @mock_stepfunctions + def test_stepfunctions_client(self): + sf = self.session.create_client("stepfunctions", region_name="us-west-2") + Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) + sf.start_execution(stateMachineArn="foo", name="bar", input='{"baz":1}') + spans = self.get_spans() + assert spans + span = spans[0] + assert len(spans) == 1 + assert span.get_tag("aws.region") == "us-west-2" + assert span.get_tag("region") == "us-west-2" + assert span.get_tag("aws.operation") == "StartExecution" + assert span.get_tag("component") == "botocore" + assert span.get_tag("span.kind"), "client" + assert_is_measured(span) + assert_span_http_status_code(span, 200) + assert span.service == "test-botocore-tracing.stepfunctions" + assert span.resource == "stepfunctions.startexecution" + assert span.get_tag("params.input") == '{"baz":1}' + def _test_kinesis_client(self): client = self.session.create_client("kinesis", region_name="us-east-1") stream_name = "test" From 476f5522fad3f40b0986ba5a1f29db7809b72f4a Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Mon, 13 Nov 2023 13:37:34 -0800 Subject: [PATCH 02/10] Inject trace context into step function input object --- ddtrace/contrib/botocore/patch.py | 42 ++++++++++++++++++++++++++++++- tests/contrib/botocore/test.py | 10 ++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/ddtrace/contrib/botocore/patch.py b/ddtrace/contrib/botocore/patch.py index 9281e201a7d..956b6cb9204 100644 --- a/ddtrace/contrib/botocore/patch.py +++ b/ddtrace/contrib/botocore/patch.py @@ -238,6 +238,46 @@ def inject_trace_to_sqs_or_sns_message(params, span, endpoint_service=None, pin= inject_trace_data_to_message_attributes(trace_data, params, endpoint_service) +def inject_trace_to_stepfunction_input(params, span): + # type: (Any, Span) -> None + """ + :params: contains the params for the current botocore action + :span: the span which provides the trace context to be propagated + + Inject the trace headers into the StepFunction input if the input is a JSON string + """ + if "input" not in params: + log.warning("Unable to inject context. The StepFunction input had no input.") + return + + if params["input"] is None: + log.warning("Unable to inject context. The StepFunction input was None.") + return + + if isinstance(params["input"], dict): + if "_datadog" in params["input"]: + log.warning("Input already has trace context.") + return + params["input"]["_datadog"] = {} + HTTPPropagator.inject(span.context, params["input"]["_datadog"]) + return + + if isinstance(params["input"], str): + try: + input_obj = json.loads(params["input"]) + except ValueError: + log.warning("Input is not a valid JSON string") + return + + input_obj["_datadog"] = {} + HTTPPropagator.inject(span.context, input_obj["_datadog"]) + input_json = json.dumps(input_obj) + + params["input"] = input_json + + log.warning("Unable to inject context. The StepFunction input was not a dict or a JSON string.") + + def inject_trace_to_eventbridge_detail(params, span): # type: (Any, Span) -> None """ @@ -580,7 +620,7 @@ def patched_api_call(original_func, instance, args, kwargs): if endpoint_name == "stepfunctions" and ( operation == "start-execution" or operation == "start-sync-execution" ): - # TODO: add inject_trace method + inject_trace_to_stepfunction_input(params, span) span.name = schematize_cloud_messaging_operation( trace_operation, cloud_provider="aws", diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index 7027cb17909..e9d2e2d889c 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -906,6 +906,16 @@ def test_stepfunctions_client(self): assert span.resource == "stepfunctions.startexecution" assert span.get_tag("params.input") == '{"baz":1}' + def test_stepfunctions_send_start_execution_trace_injection(self): + sf = self.session.create_client("stepfunctions", region_name="us-west-2") + Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) + sf.start_execution(stateMachineArn="foo", name="bar", input='{"baz":1}') + spans = self.get_spans() + assert spans + span = spans[0] + input_obj = json.loads(span.get_tag("params.input")) + assert input_obj["_datadog"][HTTP_HEADER_TRACE_ID] == str(span.trace_id) + def _test_kinesis_client(self): client = self.session.create_client("kinesis", region_name="us-east-1") stream_name = "test" From 3c3f3917c45e0a0ac28aaf2f5170b1703e1bbcf5 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Mon, 13 Nov 2023 14:46:15 -0800 Subject: [PATCH 03/10] Add releasenote --- ...-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml diff --git a/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml b/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml new file mode 100644 index 00000000000..a634937c156 --- /dev/null +++ b/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml @@ -0,0 +1,5 @@ +--- +features: + - | + Add the ability to inject trace context into the input field of botocore stepfunction start_execution and + start_sync_execution calls. From b65ca3eb8b40e05f01353a9f49987243dd9736cd Mon Sep 17 00:00:00 2001 From: Christopher Agocs Date: Mon, 27 Nov 2023 09:58:13 -0800 Subject: [PATCH 04/10] Update releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml Co-authored-by: Tahir H. Butt --- ...ore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml b/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml index a634937c156..c1c226113a0 100644 --- a/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml +++ b/releasenotes/notes/inject-botocore-stepfunction-start_execution-calls-95bed0ca2e1d006e.yaml @@ -1,5 +1,5 @@ --- features: - | - Add the ability to inject trace context into the input field of botocore stepfunction start_execution and + botocore: Add the ability to inject trace context into the input field of botocore stepfunction start_execution and start_sync_execution calls. From 9e3c4b0f0d4f55784dd7b7639ff37a38feee8459 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Mon, 27 Nov 2023 10:55:47 -0800 Subject: [PATCH 05/10] ruff --- tests/contrib/botocore/test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index f361d232214..efd84404aff 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -10,7 +10,6 @@ import botocore.session import mock from moto import mock_dynamodb -from moto import mock_stepfunctions from moto import mock_ec2 from moto import mock_events from moto import mock_kinesis @@ -19,6 +18,7 @@ from moto import mock_s3 from moto import mock_sns from moto import mock_sqs +from moto import mock_stepfunctions import pytest from tests.utils import get_128_bit_trace_id_from_headers From b9932c0f6832ac61be174909ccaf2f4db0a9a65f Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Thu, 30 Nov 2023 11:13:37 -0800 Subject: [PATCH 06/10] Fix statemachine arns --- tests/contrib/botocore/test.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index efd84404aff..4aa29aba432 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -43,6 +43,7 @@ from ddtrace.internal.schema import DEFAULT_SPAN_SERVICE_NAME from ddtrace.internal.utils.version import parse_version from ddtrace.propagation.http import HTTP_HEADER_PARENT_ID +from ddtrace.propagation.http import HTTP_HEADER_TRACE_ID from tests.opentracer.utils import init_tracer from tests.utils import TracerTestCase from tests.utils import assert_is_measured @@ -890,7 +891,9 @@ def test_schematized_unspecified_service_sqs_client_v1(self): def test_stepfunctions_client(self): sf = self.session.create_client("stepfunctions", region_name="us-west-2") Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) - sf.start_execution(stateMachineArn="foo", name="bar", input='{"baz":1}') + sf.start_execution( + stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}' + ) spans = self.get_spans() assert spans span = spans[0] @@ -909,7 +912,9 @@ def test_stepfunctions_client(self): def test_stepfunctions_send_start_execution_trace_injection(self): sf = self.session.create_client("stepfunctions", region_name="us-west-2") Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) - sf.start_execution(stateMachineArn="foo", name="bar", input='{"baz":1}') + sf.start_execution( + stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}' + ) spans = self.get_spans() assert spans span = spans[0] From 2a9408b76fb1f1543f0e028b896381da53ec7e90 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Fri, 1 Dec 2023 14:47:37 -0800 Subject: [PATCH 07/10] Create and delete the state machines --- tests/contrib/botocore/test.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index 4aa29aba432..49e00b7c931 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -890,6 +890,12 @@ def test_schematized_unspecified_service_sqs_client_v1(self): @mock_stepfunctions def test_stepfunctions_client(self): sf = self.session.create_client("stepfunctions", region_name="us-west-2") + sf.create_state_machine( + name="foo", + definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', + roleArn="arn:aws:iam::012345678901:role/DummyRole", + ) + Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) sf.start_execution( stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}' @@ -908,9 +914,15 @@ def test_stepfunctions_client(self): assert span.service == "test-botocore-tracing.stepfunctions" assert span.resource == "stepfunctions.startexecution" assert span.get_tag("params.input") == '{"baz":1}' + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo") def test_stepfunctions_send_start_execution_trace_injection(self): sf = self.session.create_client("stepfunctions", region_name="us-west-2") + sf.create_state_machine( + name="foo", + definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', + roleArn="arn:aws:iam::012345678901:role/DummyRole", + ) Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) sf.start_execution( stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}' @@ -920,6 +932,7 @@ def test_stepfunctions_send_start_execution_trace_injection(self): span = spans[0] input_obj = json.loads(span.get_tag("params.input")) assert input_obj["_datadog"][HTTP_HEADER_TRACE_ID] == str(span.trace_id) + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo") def _test_kinesis_client(self): client = self.session.create_client("kinesis", region_name="us-east-1") From b521496f7367a3fa8de2a158447e447fb2282e79 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Wed, 6 Dec 2023 12:56:38 -0800 Subject: [PATCH 08/10] Fix some details about how Patch works and get tests working as well --- ddtrace/contrib/botocore/patch.py | 12 +++++++----- tests/contrib/botocore/test.py | 26 +++++++++++--------------- 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/ddtrace/contrib/botocore/patch.py b/ddtrace/contrib/botocore/patch.py index ee2ed4d80ac..254a061f0be 100644 --- a/ddtrace/contrib/botocore/patch.py +++ b/ddtrace/contrib/botocore/patch.py @@ -251,7 +251,7 @@ def inject_trace_to_stepfunction_input(params, span): log.warning("Unable to inject context. The StepFunction input was None.") return - if isinstance(params["input"], dict): + elif isinstance(params["input"], dict): if "_datadog" in params["input"]: log.warning("Input already has trace context.") return @@ -259,7 +259,7 @@ def inject_trace_to_stepfunction_input(params, span): HTTPPropagator.inject(span.context, params["input"]["_datadog"]) return - if isinstance(params["input"], str): + elif isinstance(params["input"], str): try: input_obj = json.loads(params["input"]) except ValueError: @@ -271,8 +271,10 @@ def inject_trace_to_stepfunction_input(params, span): input_json = json.dumps(input_obj) params["input"] = input_json + return - log.warning("Unable to inject context. The StepFunction input was not a dict or a JSON string.") + else: + log.warning("Unable to inject context. The StepFunction input was not a dict or a JSON string.") def inject_trace_to_eventbridge_detail(params, span): @@ -614,8 +616,8 @@ def patched_api_call(original_func, instance, args, kwargs): span.name = schematize_cloud_messaging_operation( trace_operation, cloud_provider="aws", cloud_service="sns", direction=SpanDirection.OUTBOUND ) - if endpoint_name == "stepfunctions" and ( - operation == "start-execution" or operation == "start-sync-execution" + if endpoint_name == "states" and ( + operation == "StartExecution" or operation == "StartSyncExecution" ): inject_trace_to_stepfunction_input(params, span) span.name = schematize_cloud_messaging_operation( diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index 0a354a0ac9e..b81a01c70be 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -889,7 +889,7 @@ def test_schematized_unspecified_service_sqs_client_v1(self): @mock_stepfunctions def test_stepfunctions_client(self): - sf = self.session.create_client("stepfunctions", region_name="us-west-2") + sf = self.session.create_client("stepfunctions", region_name="us-west-2", endpoint_url="http://localhost:4566") sf.create_state_machine( name="foo", definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', @@ -897,9 +897,7 @@ def test_stepfunctions_client(self): ) Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) - sf.start_execution( - stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}' - ) + sf.start_execution(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo", input='{"baz":1}') spans = self.get_spans() assert spans span = spans[0] @@ -911,28 +909,26 @@ def test_stepfunctions_client(self): assert span.get_tag("span.kind"), "client" assert_is_measured(span) assert_span_http_status_code(span, 200) - assert span.service == "test-botocore-tracing.stepfunctions" - assert span.resource == "stepfunctions.startexecution" - assert span.get_tag("params.input") == '{"baz":1}' - sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo") + assert span.service == "test-botocore-tracing.states" + assert span.resource == "states.startexecution" + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo") + @mock_stepfunctions def test_stepfunctions_send_start_execution_trace_injection(self): - sf = self.session.create_client("stepfunctions", region_name="us-west-2") + sf = self.session.create_client("stepfunctions", region_name="us-west-2", endpoint_url="http://localhost:4566") sf.create_state_machine( name="foo", definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', roleArn="arn:aws:iam::012345678901:role/DummyRole", ) Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) - sf.start_execution( - stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo", name="bar", input='{"baz":1}' - ) + sf.start_execution(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo", input='{"baz":1}') + # I've tried to find a way to make Moto show me the input to the execution, but can't get that to work. spans = self.get_spans() assert spans span = spans[0] - input_obj = json.loads(span.get_tag("params.input")) - assert input_obj["_datadog"][HTTP_HEADER_TRACE_ID] == str(span.trace_id) - sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:425362996713:stateMachine:foo") + assert span.name == "states.command" # This confirms our patch is working + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo") def _test_kinesis_client(self): client = self.session.create_client("kinesis", region_name="us-east-1") From 38705cbdc1481fd8b87cdd5d5dc29f2c296dd1f9 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Mon, 11 Dec 2023 12:18:53 -0800 Subject: [PATCH 09/10] Add more tests for inject_trace_to_stepfunction_input --- ddtrace/contrib/botocore/patch.py | 14 +++++---- tests/contrib/botocore/test.py | 49 +++++++++++++++++++------------ 2 files changed, 39 insertions(+), 24 deletions(-) diff --git a/ddtrace/contrib/botocore/patch.py b/ddtrace/contrib/botocore/patch.py index 254a061f0be..b530a18c221 100644 --- a/ddtrace/contrib/botocore/patch.py +++ b/ddtrace/contrib/botocore/patch.py @@ -266,12 +266,16 @@ def inject_trace_to_stepfunction_input(params, span): log.warning("Input is not a valid JSON string") return - input_obj["_datadog"] = {} - HTTPPropagator.inject(span.context, input_obj["_datadog"]) - input_json = json.dumps(input_obj) + if isinstance(input_obj, dict): + input_obj["_datadog"] = {} + HTTPPropagator.inject(span.context, input_obj["_datadog"]) + input_json = json.dumps(input_obj) - params["input"] = input_json - return + params["input"] = input_json + return + else: + log.warning("Unable to inject context. The StepFunction input was not a dict.") + return else: log.warning("Unable to inject context. The StepFunction input was not a dict or a JSON string.") diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index b81a01c70be..64d887948b1 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -43,7 +43,6 @@ from ddtrace.internal.schema import DEFAULT_SPAN_SERVICE_NAME from ddtrace.internal.utils.version import parse_version from ddtrace.propagation.http import HTTP_HEADER_PARENT_ID -from ddtrace.propagation.http import HTTP_HEADER_TRACE_ID from tests.opentracer.utils import init_tracer from tests.utils import TracerTestCase from tests.utils import assert_is_measured @@ -888,47 +887,59 @@ def test_schematized_unspecified_service_sqs_client_v1(self): assert spans[2].name == "aws.sqs.receive" @mock_stepfunctions - def test_stepfunctions_client(self): + def test_stepfunctions_send_start_execution_trace_injection(self): sf = self.session.create_client("stepfunctions", region_name="us-west-2", endpoint_url="http://localhost:4566") sf.create_state_machine( - name="foo", + name="lincoln", definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', roleArn="arn:aws:iam::012345678901:role/DummyRole", ) + Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) + sf.start_execution( + stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:lincoln", input='{"baz":1}' + ) + # I've tried to find a way to make Moto show me the input to the execution, but can't get that to work. + spans = self.get_spans() + assert spans + span = spans[0] + assert span.name == "states.command" # This confirms our patch is working + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:lincoln") + @mock_stepfunctions + def test_stepfunctions_send_start_execution_trace_injection_with_array_input(self): + sf = self.session.create_client("stepfunctions", region_name="us-west-2", endpoint_url="http://localhost:4566") + sf.create_state_machine( + name="miller", + definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', + roleArn="arn:aws:iam::012345678901:role/DummyRole", + ) Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) - sf.start_execution(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo", input='{"baz":1}') + sf.start_execution( + stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:miller", input='["one", "two", "three"]' + ) + # I've tried to find a way to make Moto show me the input to the execution, but can't get that to work. spans = self.get_spans() assert spans span = spans[0] - assert len(spans) == 1 - assert span.get_tag("aws.region") == "us-west-2" - assert span.get_tag("region") == "us-west-2" - assert span.get_tag("aws.operation") == "StartExecution" - assert span.get_tag("component") == "botocore" - assert span.get_tag("span.kind"), "client" - assert_is_measured(span) - assert_span_http_status_code(span, 200) - assert span.service == "test-botocore-tracing.states" - assert span.resource == "states.startexecution" - sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo") + assert span.name == "states.command" # This confirms our patch is working + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:miller") @mock_stepfunctions - def test_stepfunctions_send_start_execution_trace_injection(self): + def test_stepfunctions_send_start_execution_trace_injection_with_true_input(self): sf = self.session.create_client("stepfunctions", region_name="us-west-2", endpoint_url="http://localhost:4566") sf.create_state_machine( - name="foo", + name="hobart", definition='{"StartAt": "HelloWorld","States": {"HelloWorld": {"Type": "Pass","End": true}}}', roleArn="arn:aws:iam::012345678901:role/DummyRole", ) Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(sf) - sf.start_execution(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo", input='{"baz":1}') + sf.start_execution(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:hobart", input="true") # I've tried to find a way to make Moto show me the input to the execution, but can't get that to work. spans = self.get_spans() assert spans span = spans[0] assert span.name == "states.command" # This confirms our patch is working - sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:foo") + sf.delete_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:hobart") def _test_kinesis_client(self): client = self.session.create_client("kinesis", region_name="us-east-1") From 186bf7fce031954ab8e57e2b7163c55619601ab2 Mon Sep 17 00:00:00 2001 From: "chris.agocs" Date: Thu, 11 Jan 2024 14:56:52 -0800 Subject: [PATCH 10/10] run the linter --- ddtrace/contrib/botocore/patch.py | 3 ++- .../contrib/botocore/services/stepfunctions.py | 17 ++++++++++++----- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/ddtrace/contrib/botocore/patch.py b/ddtrace/contrib/botocore/patch.py index d8bded1a64b..8f4c6b68b08 100644 --- a/ddtrace/contrib/botocore/patch.py +++ b/ddtrace/contrib/botocore/patch.py @@ -35,7 +35,8 @@ from .services.sqs import inject_trace_to_sqs_or_sns_batch_message from .services.sqs import inject_trace_to_sqs_or_sns_message from .services.sqs import patched_sqs_api_call -from .services.stepfunctions import inject_trace_to_stepfunction_input, patched_stepfunction_api_call +from .services.stepfunctions import inject_trace_to_stepfunction_input +from .services.stepfunctions import patched_stepfunction_api_call from .utils import inject_trace_to_client_context from .utils import inject_trace_to_eventbridge_detail from .utils import set_patched_api_call_span_tags diff --git a/ddtrace/contrib/botocore/services/stepfunctions.py b/ddtrace/contrib/botocore/services/stepfunctions.py index 276cac708f7..7aa32a514e3 100644 --- a/ddtrace/contrib/botocore/services/stepfunctions.py +++ b/ddtrace/contrib/botocore/services/stepfunctions.py @@ -1,15 +1,22 @@ import json +from typing import Any # noqa:F401 +from typing import Dict # noqa:F401 import botocore.exceptions +from ddtrace import Span # noqa:F401 +from ddtrace import config +from ddtrace.ext import http from ddtrace.propagation.http import HTTPPropagator -from ..utils import set_patched_api_call_span_tags, set_response_metadata_tags + from ....ext import SpanTypes from ....internal.logger import get_logger -from typing import Any, Dict # noqa:F401 -from ddtrace import Span # noqa:F401 -from ....internal.schema import schematize_service_name, SpanDirection, schematize_cloud_messaging_operation -from ddtrace import config +from ....internal.schema import SpanDirection +from ....internal.schema import schematize_cloud_messaging_operation +from ....internal.schema import schematize_service_name +from ..utils import set_patched_api_call_span_tags +from ..utils import set_response_metadata_tags + log = get_logger(__name__)