Merge branch 'main' into emmett.butler/nine-more-kilobytes

Author: emmettbutler

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

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,13 @@ 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(
+                        "aws_lambda.durable_function.execution_status",
+                        durable_status,
+                    )
+
                 self.span.finish()
 
             if status_code:

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))