diff --git a/src/lumigo_tracer/parsers/utils.py b/src/lumigo_tracer/parsers/utils.py index b3cf6e32..74ef1994 100644 --- a/src/lumigo_tracer/parsers/utils.py +++ b/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 diff --git a/src/lumigo_tracer/utils.py b/src/lumigo_tracer/utils.py index faa219bb..2c98e056 100644 --- a/src/lumigo_tracer/utils.py +++ b/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" ) diff --git a/src/test/unit/parsers/test_utils.py b/src/test/unit/parsers/test_utils.py index 8f5b8700..cbbcfd11 100644 --- a/src/test/unit/parsers/test_utils.py +++ b/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",