diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index e9443f92..3acc7c0c 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -41,9 +41,13 @@ def extract_durable_function_tags(event): if not parsed: logger.error("Failed to parse DurableExecutionArn: %s", durable_execution_arn) return {} - execution_name, execution_id = parsed + # Use the number of operations to determine if it's the first invocation. This is + # what the durable execution SDK does to determine the replay status. + operations = event.get("InitialExecutionState", {}).get("Operations", []) + is_first_invocation = len(operations) == 1 return { "durable_function_execution_name": execution_name, "durable_function_execution_id": execution_id, + "durable_function_first_invocation": str(is_first_invocation).lower(), } diff --git a/tests/test_durable.py b/tests/test_durable.py index 60914934..0887e2d9 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -45,11 +45,12 @@ def test_works_with_numeric_version_qualifier(self): class TestExtractDurableFunctionTags(unittest.TestCase): - def test_extracts_tags_from_event_with_durable_execution_arn(self): + def test_sets_first_invocation_true_when_only_execution_operation(self): + # One operation (the current EXECUTION operation itself) → not replaying → first invocation event = { "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", "CheckpointToken": "some-token", - "InitialExecutionState": {"Operations": []}, + "InitialExecutionState": {"Operations": [{"OperationType": "EXECUTION"}]}, } result = extract_durable_function_tags(event) self.assertEqual( @@ -57,6 +58,29 @@ def test_extracts_tags_from_event_with_durable_execution_arn(self): { "durable_function_execution_name": "my-execution", "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", + "durable_function_first_invocation": "true", + }, + ) + + def test_sets_first_invocation_false_when_multiple_operations(self): + # More than one operation → replaying → not first invocation + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + "CheckpointToken": "some-token", + "InitialExecutionState": { + "Operations": [ + {"OperationType": "EXECUTION"}, + {"OperationType": "STEP"}, + ] + }, + } + result = extract_durable_function_tags(event) + self.assertEqual( + result, + { + "durable_function_execution_name": "my-execution", + "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", + "durable_function_first_invocation": "false", }, )