From b1afe17777c77527928fb6b75d22b4d23694022f Mon Sep 17 00:00:00 2001 From: ZCode Date: Sun, 30 Aug 2026 10:34:32 +0800 Subject: [PATCH] fix: resolve logging.py issues from code audit Three fixes in rdsa_utils/logging.py: 1. & -> and operator (line 237): Changed bitwise AND to logical AND in show_records & stop_pipeline condition. 2. try/finally for cache cleanup (log_rows_in_spark_df): Moved kwargs['df'].unpersist() into a finally block so the input DataFrame cache is always released, even if the decorated function raises an exception. Prevents Spark cache leak / OOM. 3. Added isinstance check before unpersist to avoid AttributeError when kwargs['df'] is not a Spark DataFrame. See issue #254 for full audit report. --- rdsa_utils/logging.py | 47 ++++++++++++++++++++++--------------------- 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/rdsa_utils/logging.py b/rdsa_utils/logging.py index 6ae72b4..bcf86a6 100644 --- a/rdsa_utils/logging.py +++ b/rdsa_utils/logging.py @@ -240,7 +240,7 @@ def print_full_table_and_raise_error( ValueError Raises error and stops pipeline if switch applied. """ - if show_records & stop_pipeline: + if show_records and stop_pipeline: logger.error(df.to_string()) logger.error(message) raise ValueError @@ -490,34 +490,35 @@ def wrapper_decorator(*args, **kwargs): ), ) + try: # Run the decorated function in its normal way, but catch its output # so it can be counted. df_return = func(*args, **kwargs) + finally: + # Ensure input df cache is always released, even if func raises + if kwargs.get("df") and isinstance(kwargs["df"], SparkDF): + kwargs["df"].unpersist() - # Check to ensure that the function returns a single value that is a - # spark dataframe, as otherwise the count operation will fail. - if isinstance(df_return, SparkDF): - # Persist the dataframe to be returned prior to counting to allow - # more efficient processing downstream. We persist here as we will - # not be removing this dataframe from the in-memory cache in this - # decorator. Therefore, we don't want it to get pushed onto disk - # (and incur an expensive swap operation). - df_return.persist(StorageLevel.MEMORY_ONLY) - logger.info( - f"Rows in dataframe after {func_name} : {df_return.count()}", - ) # noqa: E501 - - else: - logger.warning( - f"{func_name} should return a spark dataframe for decorator, " - f"but returned {type(df_return)}", - ) + # Check to ensure that the function returns a single value that is a + # spark dataframe, as otherwise the count operation will fail. + if isinstance(df_return, SparkDF): + # Persist the dataframe to be returned prior to counting to allow + # more efficient processing downstream. We persist here as we will + # not be removing this dataframe from the in-memory cache in this + # decorator. Therefore, we don't want it to get pushed onto disk + # (and incur an expensive swap operation). + df_return.persist(StorageLevel.MEMORY_ONLY) + logger.info( + f"Rows in dataframe after {func_name} : {df_return.count()}", + ) # noqa: E501 - if kwargs.get("df"): - # Unpersist the cached input df to manage memory. - kwargs["df"].unpersist() + else: + logger.warning( + f"{func_name} should return a spark dataframe for decorator, " + f"but returned {type(df_return)}", + ) - return df_return + return df_return return wrapper_decorator