diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index e9443f92..3174cd6c 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -47,3 +47,17 @@ def extract_durable_function_tags(event): "durable_function_execution_name": execution_name, "durable_function_execution_id": execution_id, } + + +VALID_DURABLE_STATUSES = {"SUCCEEDED", "FAILED", "STOPPED", "TIMED_OUT"} + + +def extract_durable_execution_status(response, event): + if not isinstance(event, dict) or "DurableExecutionArn" not in event: + return None + if not isinstance(response, dict): + return None + status = response.get("Status") + if status not in VALID_DURABLE_STATUSES: + return None + return status diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index d022988b..b5f3b7d5 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -42,7 +42,10 @@ tracer, propagator, ) -from datadog_lambda.durable import extract_durable_function_tags +from datadog_lambda.durable import ( + extract_durable_function_tags, + extract_durable_execution_status, +) from datadog_lambda.trigger import ( extract_trigger_tags, extract_http_status_code_tag, @@ -153,7 +156,7 @@ def __init__(self, func): if config.trace_extractor: extractor_parts = config.trace_extractor.rsplit(".", 1) if len(extractor_parts) == 2: - (mod_name, extractor_name) = extractor_parts + mod_name, extractor_name = extractor_parts modified_extractor_name = modify_module_name(mod_name) extractor_module = import_module(modified_extractor_name) self.trace_extractor = getattr(extractor_module, extractor_name) @@ -340,6 +343,12 @@ def _after(self, event, context): if status_code: self.span.set_tag("http.status_code", status_code) + durable_status = extract_durable_execution_status(self.response, event) + if durable_status: + self.span.set_tag( + "durable_function_execution_status", durable_status + ) + self.span.finish() if status_code: diff --git a/tests/test_durable.py b/tests/test_durable.py index 60914934..e668b45f 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -7,6 +7,7 @@ from datadog_lambda.durable import ( _parse_durable_execution_arn, extract_durable_function_tags, + extract_durable_execution_status, ) @@ -89,3 +90,72 @@ def test_returns_empty_dict_when_durable_execution_arn_cannot_be_parsed(self): def test_returns_empty_dict_when_event_is_empty(self): result = extract_durable_function_tags({}) self.assertEqual(result, {}) + + +class TestExtractDurableExecutionStatus(unittest.TestCase): + def test_returns_succeeded(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + response = {"Status": "SUCCEEDED", "Result": "some-result"} + self.assertEqual(extract_durable_execution_status(response, event), "SUCCEEDED") + + def test_returns_failed(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + response = {"Status": "FAILED", "Error": "some-error"} + self.assertEqual(extract_durable_execution_status(response, event), "FAILED") + + def test_returns_stopped(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + response = {"Status": "STOPPED"} + self.assertEqual(extract_durable_execution_status(response, event), "STOPPED") + + def test_returns_timed_out(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + response = {"Status": "TIMED_OUT"} + self.assertEqual(extract_durable_execution_status(response, event), "TIMED_OUT") + + def test_returns_none_for_non_durable_event(self): + event = {"key": "value"} + response = {"Status": "SUCCEEDED"} + self.assertIsNone(extract_durable_execution_status(response, event)) + + def test_returns_none_for_non_dict_response(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + self.assertIsNone(extract_durable_execution_status("string", event)) + + def test_returns_none_for_missing_status(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + response = {"Result": "some-result"} + self.assertIsNone(extract_durable_execution_status(response, event)) + + def test_returns_none_for_invalid_status(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + response = {"Status": "INVALID"} + self.assertIsNone(extract_durable_execution_status(response, event)) + + def test_returns_none_for_non_dict_event(self): + response = {"Status": "SUCCEEDED"} + self.assertIsNone(extract_durable_execution_status(response, "not-a-dict")) + + def test_returns_none_for_none_event(self): + response = {"Status": "SUCCEEDED"} + self.assertIsNone(extract_durable_execution_status(response, None)) + + def test_returns_none_for_none_response(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id" + } + self.assertIsNone(extract_durable_execution_status(None, event))