search lumigo's key recursively in the event (#141)

lumigo-io · Aug 18, 2020 · ac6e9a4 · ac6e9a4
1 parent 124161f
commit ac6e9a4
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 6 deletions.
diff --git a/src/lumigo_tracer/parsers/utils.py b/src/lumigo_tracer/parsers/utils.py
@@ -184,15 +184,15 @@ def _parse_unknown(event: dict):
 
 
 def _is_step_function(event):
-    return Configuration.is_step_function and STEP_FUNCTION_UID_KEY in event.get(
-        LUMIGO_EVENT_KEY, {}
+    return Configuration.is_step_function and STEP_FUNCTION_UID_KEY in recursive_get_key(
+        event, LUMIGO_EVENT_KEY, default={}
     )
 
 
 def _parse_step_function(event: dict):
     result = {
         "triggeredBy": "stepFunction",
-        "messageId": event[LUMIGO_EVENT_KEY][STEP_FUNCTION_UID_KEY],
+        "messageId": recursive_get_key(event, LUMIGO_EVENT_KEY)[STEP_FUNCTION_UID_KEY],
     }
     return result
 
@@ -344,3 +344,18 @@ def str_to_tuple(val: str) -> Optional[Tuple]:
     except Exception as e:
         get_logger().debug("Error while convert str to tuple", exc_info=e)
     return None
+
+
+def recursive_get_key(d: 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):
+            recursive_result = recursive_get_key(v, key, depth - 1, default)
+            if recursive_result:
+                return recursive_result
+    return default
diff --git a/src/lumigo_tracer/utils.py b/src/lumigo_tracer/utils.py
@@ -78,6 +78,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
 
 
 def config(
@@ -91,6 +92,7 @@ def config(
     timeout_timer_buffer: Optional[float] = None,
     domains_scrubber: Optional[List[str]] = None,
     max_entry_size: int = DEFAULT_MAX_ENTRY_SIZE,
+    get_key_depth: int = None,
 ) -> None:
     """
     This function configure the lumigo wrapper.
@@ -106,6 +108,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.
     """
     if should_report is not None:
         Configuration.should_report = should_report
@@ -117,6 +120,7 @@ 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.is_step_function = (
         step_function or os.environ.get("LUMIGO_STEP_FUNCTION", "").lower() == "true"
     )

diff --git a/src/test/unit/parsers/test_utils.py b/src/test/unit/parsers/test_utils.py
@@ -201,7 +201,31 @@ def test_recursive_json_join(d1, d2, result):
             },
             {"triggeredBy": "stepFunction", "messageId": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"},
         ),
-        (
+        (  # Inner Step Function
+            {
+                "bla": "saart",
+                "inner": {"_lumigo": {"step_function_uid": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"}},
+            },
+            {"triggeredBy": "stepFunction", "messageId": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"},
+        ),
+        (  # Step Function - too deep
+            {
+                "bla": "saart",
+                "a": {
+                    "b": {
+                        "c": {
+                            "d": {
+                                "_lumigo": {
+                                    "step_function_uid": "54589cfc-5ed8-4799-8fc0-5b45f6f225d1"
+                                }
+                            }
+                        }
+                    }
+                },
+            },
+            {"triggeredBy": "unknown"},
+        ),
+        (  # cloudwatch
             {
                 "id": "cdc73f9d-aea9-11e3-9d5a-835b769c0d9c",
                 "detail-type": "Scheduled Event",
@@ -218,7 +242,7 @@ def test_recursive_json_join(d1, d2, result):
                 "detailType": "Scheduled Event",
             },
         ),
-        (
+        (  # unknown
             {
                 "id": "cdc73f9d-aea9-11e3-9d5a-835b769c0d9c",
                 "detail-type": "Unknown",
@@ -230,7 +254,7 @@ def test_recursive_json_join(d1, d2, result):
             },
             {"triggeredBy": "unknown"},
         ),
-        (
+        (  # cloudwatch
             {
                 "id": "cdc73f9d-aea9-11e3-9d5a-835b769c0d9c",
                 "detail-type": "Scheduled Event",