RD-3530 - support events with list for step functions (#143)

* support events with list for step functions

Author: saartochner-lumigo

src/lumigo_tracer/parsers/utils.py

@@ -163,6 +163,8 @@ def parse_triggered_by(event: dict):
     """
     with lumigo_safe_execute("triggered by"):
         if not isinstance(event, dict):
+            if _is_step_function(event):
+                return _parse_step_function(event)
             return None
         if _is_supported_http_method(event):
             return parse_http_method(event)
@@ -183,9 +185,11 @@ def _parse_unknown(event: dict):
     return result
 
 
-def _is_step_function(event):
-    return Configuration.is_step_function and STEP_FUNCTION_UID_KEY in recursive_get_key(
-        event, LUMIGO_EVENT_KEY, default={}
+def _is_step_function(event: Union[List, Dict]):
+    return (
+        Configuration.is_step_function
+        and isinstance(event, (list, dict))  # noqa
+        and STEP_FUNCTION_UID_KEY in recursive_get_key(event, LUMIGO_EVENT_KEY, default={})  # noqa
     )
 
 
@@ -346,16 +350,22 @@ def str_to_tuple(val: str) -> Optional[Tuple]:
     return None
 
 
-def recursive_get_key(d: Dict[str, Union[Dict, str]], key, depth=None, default=None):
+def recursive_get_key(d: Union[List, Dict[str, Union[Dict, str]]], key, depth=None, default=None):
     if depth is None:
         depth = Configuration.get_key_depth
     if depth == 0:
         return default
     if key in d:
         return d[key]
-    for v in d.values():
-        if isinstance(v, dict):
+    if isinstance(d, list):
+        for v in d:
             recursive_result = recursive_get_key(v, key, depth - 1, default)
             if recursive_result:
                 return recursive_result
+    if isinstance(d, dict):
+        for v in d.values():
+            if isinstance(v, (list, dict)):
+                recursive_result = recursive_get_key(v, key, depth - 1, default)
+                if recursive_result:
+                    return recursive_result
     return default

src/lumigo_tracer/utils.py

@@ -51,6 +51,7 @@
 WARN_CLIENT_PREFIX = "Lumigo Warning"
 TRUNCATE_SUFFIX = "...[too long]"
 NUMBER_OF_SPANS_IN_REPORT_OPTIMIZATION = 200
+DEFAULT_KEY_DEPTH = 4
 
 _logger: Union[logging.Logger, None] = None
 
@@ -78,7 +79,7 @@ class Configuration:
     send_only_if_error: bool = False
     domains_scrubber: Optional[List] = None
     max_entry_size: int = DEFAULT_MAX_ENTRY_SIZE
-    get_key_depth: int = 3
+    get_key_depth: int = DEFAULT_KEY_DEPTH
 
 
 def config(
@@ -108,7 +109,7 @@ def config(
         The default is 10% of the duration of the lambda (with upper and lower bounds of 0.5 and 3 seconds).
     :param domains_scrubber: List of regexes. We will not collect data of requests with hosts that match it.
     :param max_entry_size: The maximum size of each entry when sending back the events.
-    :param get_key_depth: Max depth to search the lumigo key in the event (relevant to step functions). default 3.
+    :param get_key_depth: Max depth to search the lumigo key in the event (relevant to step functions). default 4.
     """
     if should_report is not None:
         Configuration.should_report = should_report
@@ -120,7 +121,9 @@ def config(
         enhance_print or os.environ.get("LUMIGO_ENHANCED_PRINT", "").lower() == "true"
     )
     Configuration.verbose = verbose and os.environ.get("LUMIGO_VERBOSE", "").lower() != "false"
-    Configuration.get_key_depth = get_key_depth or int(os.environ.get("LUMIGO_EVENT_KEY_DEPTH", 3))
+    Configuration.get_key_depth = get_key_depth or int(
+        os.environ.get("LUMIGO_EVENT_KEY_DEPTH", DEFAULT_KEY_DEPTH)
+    )
     Configuration.is_step_function = (
         step_function or os.environ.get("LUMIGO_STEP_FUNCTION", "").lower() == "true"
     )

src/test/unit/parsers/test_utils.py

@@ -208,6 +208,28 @@ def test_recursive_json_join(d1, d2, result):
             },
             {"triggeredBy": "stepFunction", "messageId": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"},
         ),
+        (  # Step Function from list
+            [
+                {
+                    "bla": "saart",
+                    "inner": {
+                        "_lumigo": {"step_function_uid": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"}
+                    },
+                },
+                {"something": "else"},
+            ],
+            {"triggeredBy": "stepFunction", "messageId": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"},
+        ),
+        (  # Step Function from inner list
+            {
+                "bla": "saart",
+                "inner": [
+                    {"_lumigo": {"step_function_uid": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"}},
+                    {"something": "else"},
+                ],
+            },
+            {"triggeredBy": "stepFunction", "messageId": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"},
+        ),
         (  # Step Function - too deep
             {
                 "bla": "saart",