diff --git a/src/lumigo_tracer/parsers/utils.py b/src/lumigo_tracer/parsers/utils.py index 157da2dd..b3cf6e32 100644 --- 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 index 07fc3fe6..faa219bb 100644 --- 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 index 28ee013d..8f5b8700 100644 --- 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",