From c2bb1e50044ccbd36444106eba14ee45d3d666c2 Mon Sep 17 00:00:00 2001 From: IceS2 Date: Fri, 2 Feb 2024 14:23:09 +0100 Subject: [PATCH] MINOR: Fix 'calculate_execution_time_generator' (#14987) * Refactor the calculate_execution_time for generators decorator in order to fix it * Add small comment in the code --- ingestion/src/metadata/utils/helpers.py | 27 +++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/ingestion/src/metadata/utils/helpers.py b/ingestion/src/metadata/utils/helpers.py index 94692ed1475..3cb857e66ed 100644 --- a/ingestion/src/metadata/utils/helpers.py +++ b/ingestion/src/metadata/utils/helpers.py @@ -125,12 +125,27 @@ def calculate_execution_time_generator(func): """ def calculate_debug_time(*args, **kwargs): - start = perf_counter() - yield from func(*args, **kwargs) - end = perf_counter() - logger.debug( - f"{func.__name__} executed in { pretty_print_time_duration(end - start)}" - ) + # NOTE: We are basically implementing by hand a simplified version of 'yield from' + # in order to be able to calculate the time difference correctly. + # The 'while True' loop allows us to guarantee we are iterating over all thje values + # from func(*args, **kwargs). + generator = func(*args, **kwargs) + + while True: + start = perf_counter() + + try: + element = next(generator) + except StopIteration: + return + + end = perf_counter() + + logger.debug( + f"{func.__name__} executed in { pretty_print_time_duration(end - start)}" + ) + + yield element return calculate_debug_time