Skip to content

Commit 0a51475

Browse files
avedmalaavedmala
committed
fix circular import
1 parent a38c77b commit 0a51475

File tree

4 files changed

+91
-92
lines changed

4 files changed

+91
-92
lines changed

‎datadog_lambda/tracing.py

+1-26
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
_EventSource,
4040
parse_event_source,
4141
get_first_record,
42+
is_step_function_event,
4243
EventTypes,
4344
EventSubtypes,
4445
)
@@ -479,32 +480,6 @@ def extract_context_from_step_functions(event, lambda_context):
479480
return extract_context_from_lambda_context(lambda_context)
480481

481482

482-
def is_step_function_event(event):
483-
"""
484-
Check if the event is a step function that invoked the current lambda.
485-
486-
The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a
487-
"_datadog" for JSONata style context propagation.
488-
489-
The actual event must contain "Execution", "StateMachine", and "State" fields.
490-
"""
491-
event = event.get("Payload", event)
492-
493-
# JSONPath style
494-
if all(field in event for field in ("Execution", "StateMachine", "State")):
495-
return True
496-
497-
# JSONata style
498-
if "_datadog" in event:
499-
event = event["_datadog"]
500-
return all(
501-
field in event
502-
for field in ("Execution", "StateMachine", "State", "serverless-version")
503-
)
504-
505-
return False
506-
507-
508483
def extract_context_custom_extractor(extractor, event, lambda_context):
509484
"""
510485
Extract Datadog trace context using a custom trace extractor function

‎datadog_lambda/trigger.py

+25-2
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@
1010
from enum import Enum
1111
from typing import Any
1212

13-
from datadog_lambda.tracing import is_step_function_event
14-
1513

1614
class _stringTypedEnum(Enum):
1715
"""
@@ -369,3 +367,28 @@ def extract_http_status_code_tag(trigger_tags, response):
369367
status_code = response.status_code
370368

371369
return str(status_code)
370+
371+
def is_step_function_event(event):
372+
"""
373+
Check if the event is a step function that invoked the current lambda.
374+
375+
The whole event can be wrapped in "Payload" in Legacy Lambda cases. There may also be a
376+
"_datadog" for JSONata style context propagation.
377+
378+
The actual event must contain "Execution", "StateMachine", and "State" fields.
379+
"""
380+
event = event.get("Payload", event)
381+
382+
# JSONPath style
383+
if all(field in event for field in ("Execution", "StateMachine", "State")):
384+
return True
385+
386+
# JSONata style
387+
if "_datadog" in event:
388+
event = event["_datadog"]
389+
return all(
390+
field in event
391+
for field in ("Execution", "StateMachine", "State", "serverless-version")
392+
)
393+
394+
return False

‎tests/test_tracing.py

-64
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@
4242
service_mapping as global_service_mapping,
4343
propagator,
4444
emit_telemetry_on_exception_outside_of_handler,
45-
is_step_function_event,
4645
)
4746
from datadog_lambda.trigger import EventTypes
4847

@@ -2331,69 +2330,6 @@ def test_deterministic_m5_hash__always_leading_with_zero(self):
23312330
if len(result_in_binary) == 66: # "0b" + 64 bits.
23322331
self.assertTrue(result_in_binary.startswith("0b0"))
23332332

2334-
def test_is_step_function_event_jsonata(self):
2335-
event = {
2336-
"_datadog": {
2337-
"Execution": {
2338-
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
2339-
"RedriveCount": 0,
2340-
},
2341-
"StateMachine": {},
2342-
"State": {
2343-
"Name": "my-awesome-state",
2344-
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
2345-
"RetryCount": 0,
2346-
},
2347-
"x-datadog-trace-id": "5821803790426892636",
2348-
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
2349-
"serverless-version": "v1",
2350-
}
2351-
}
2352-
self.assertTrue(is_step_function_event(event))
2353-
2354-
def test_is_step_function_event_jsonpath(self):
2355-
event = {
2356-
"Execution": {
2357-
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
2358-
"RedriveCount": 0,
2359-
},
2360-
"StateMachine": {},
2361-
"State": {
2362-
"Name": "my-awesome-state",
2363-
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
2364-
"RetryCount": 0,
2365-
},
2366-
}
2367-
self.assertTrue(is_step_function_event(event))
2368-
2369-
def test_is_step_function_event_legacy_lambda(self):
2370-
event = {
2371-
"Payload": {
2372-
"Execution": {
2373-
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
2374-
"RedriveCount": 0,
2375-
},
2376-
"StateMachine": {},
2377-
"State": {
2378-
"Name": "my-awesome-state",
2379-
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
2380-
"RetryCount": 0,
2381-
},
2382-
}
2383-
}
2384-
self.assertTrue(is_step_function_event(event))
2385-
2386-
def test_is_step_function_event_dd_header(self):
2387-
event = {
2388-
"_datadog": {
2389-
"x-datadog-trace-id": "5821803790426892636",
2390-
"x-datadog-parent-id": "5821803790426892636",
2391-
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
2392-
"x-datadog-sampling-priority": "1",
2393-
}
2394-
}
2395-
self.assertFalse(is_step_function_event(event))
2396-
23972333

23982334
class TestExceptionOutsideHandler(unittest.TestCase):
23992335
@patch("datadog_lambda.tracing.dd_tracing_enabled", True)

‎tests/test_trigger.py

+65
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
get_event_source_arn,
1010
extract_trigger_tags,
1111
extract_http_status_code_tag,
12+
is_step_function_event,
1213
)
1314

1415
from tests.utils import get_mock_context
@@ -543,3 +544,67 @@ def test_extract_http_status_code_tag_from_response_object(self):
543544
response.status_code = 403
544545
status_code = extract_http_status_code_tag(trigger_tags, response)
545546
self.assertEqual(status_code, "403")
547+
548+
class IsStepFunctionEvent(unittest.TestCase):
549+
def test_is_step_function_event_jsonata(self):
550+
event = {
551+
"_datadog": {
552+
"Execution": {
553+
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
554+
"RedriveCount": 0,
555+
},
556+
"StateMachine": {},
557+
"State": {
558+
"Name": "my-awesome-state",
559+
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
560+
"RetryCount": 0,
561+
},
562+
"x-datadog-trace-id": "5821803790426892636",
563+
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
564+
"serverless-version": "v1",
565+
}
566+
}
567+
self.assertTrue(is_step_function_event(event))
568+
569+
def test_is_step_function_event_jsonpath(self):
570+
event = {
571+
"Execution": {
572+
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
573+
"RedriveCount": 0,
574+
},
575+
"StateMachine": {},
576+
"State": {
577+
"Name": "my-awesome-state",
578+
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
579+
"RetryCount": 0,
580+
},
581+
}
582+
self.assertTrue(is_step_function_event(event))
583+
584+
def test_is_step_function_event_legacy_lambda(self):
585+
event = {
586+
"Payload": {
587+
"Execution": {
588+
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
589+
"RedriveCount": 0,
590+
},
591+
"StateMachine": {},
592+
"State": {
593+
"Name": "my-awesome-state",
594+
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
595+
"RetryCount": 0,
596+
},
597+
}
598+
}
599+
self.assertTrue(is_step_function_event(event))
600+
601+
def test_is_step_function_event_dd_header(self):
602+
event = {
603+
"_datadog": {
604+
"x-datadog-trace-id": "5821803790426892636",
605+
"x-datadog-parent-id": "5821803790426892636",
606+
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=672a7cb100000000",
607+
"x-datadog-sampling-priority": "1",
608+
}
609+
}
610+
self.assertFalse(is_step_function_event(event))

0 commit comments

Comments
 (0)