cache spark config read

Author: yankomo

spark_expectations/core/context.py

@@ -136,6 +136,8 @@ def __post_init__(self) -> None:
         self._debugger_mode: bool = False
         self._supported_df_query_dq: DataFrame = self.set_supported_df_query_dq()
 
+        self._ansi_enabled: Optional[bool] = None
+
         self._source_agg_dq_start_time: Optional[datetime] = None
         self._final_agg_dq_start_time: Optional[datetime] = None
         self._source_query_dq_start_time: Optional[datetime] = None
@@ -3056,3 +3058,20 @@ def get_df_dq_obs_report_dataframe(self) -> DataFrame:
             DataFrame: DataFrame containing DQ observability report data
         """
         return self._df_dq_obs_report_dataframe
+
+    @property
+    def get_ansi_enabled(self) -> bool:
+        """
+        Returns whether Spark's ANSI mode is enabled. Value is cached after the first read.
+
+        Returns:
+            bool: Returns the ANSI mode enabled boolean
+        """
+        if self._ansi_enabled is None:
+            try:
+                self._ansi_enabled = self.spark.conf.get("spark.sql.ansi.enabled", "false").lower() == "true"
+            except Exception as e:
+                # Catch-all for any unexpected errors; return safe default
+                _log.warning(f"Failed to retrieve Spark ANSI mode, setting to 'false'. Error: {e}")
+                self._ansi_enabled = False
+        return self._ansi_enabled

spark_expectations/sinks/utils/report.py

@@ -36,6 +36,7 @@ def __post_init__(self) -> None:
     def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
         try:
             context = self._context
+            ansi_enabled = context.get_ansi_enabled
 
             print("dq_obs_report_data_insert method called stats_detailed table")
             df_stats_detailed = context.get_stats_detailed_dataframe
@@ -174,7 +175,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                 .withColumn(
                     "total_records_only_nbr",
                     when(col("_total_records_str") == "", lit(None).cast("bigint"))
-                    .otherwise(safe_cast(self.spark, "_total_records_str", "bigint")),
+                    .otherwise(safe_cast(ansi_enabled, "_total_records_str", "bigint")),
                 )
                 .withColumn(
                     "_valid_records_str",
@@ -183,7 +184,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                 .withColumn(
                     "valid_records_only_nbr",
                     when(col("_valid_records_str") == "", lit(None).cast("bigint"))
-                    .otherwise(safe_cast(self.spark, "_valid_records_str", "bigint")),
+                    .otherwise(safe_cast(ansi_enabled, "_valid_records_str", "bigint")),
                 )
                 .drop("_total_records_str", "_valid_records_str")
                 .withColumn(
@@ -199,7 +200,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                     .otherwise(
                         coalesce(
                             safe_cast(
-                                self.spark,
+                                ansi_enabled,
                                 """
                                 100 
                                 * least(
@@ -330,7 +331,7 @@ def dq_obs_report_data_insert(self) -> tuple[bool, DataFrame]:
                 .withColumnRenamed("source_dq_error_row_count", "failed_records")
                 .withColumn(
                     "success_percentage",
-                    (safe_cast(self.spark, "valid_records", "double") / safe_cast(self.spark, "total_records", "double")) * 100,
+                    (safe_cast(ansi_enabled, "valid_records", "double") / safe_cast(ansi_enabled, "total_records", "double")) * 100,
                 )
             )
 

spark_expectations/utils/actions.py

@@ -445,7 +445,7 @@ def execute_sql_and_get_result(
                 row_count,
             )
         except Exception as e:
-            raise SparkExpectationsMiscException(f"error occurred while running agg_query_dq_detailed_result {e}")
+            raise SparkExpectationsMiscException(f"error occurred while running agg_query_dq_detailed_result: {e}")
 
     @staticmethod
     def create_agg_dq_results(
@@ -479,7 +479,7 @@ def create_agg_dq_results(
                     return meta_results
             return None
         except Exception as e:
-            raise SparkExpectationsMiscException(f"error occurred while running create agg dq results {e}")
+            raise SparkExpectationsMiscException(f"error occurred while running create_agg_dq_results: {e}")
 
     @staticmethod
     def run_dq_rules(
@@ -632,7 +632,7 @@ def run_dq_rules(
             return df
 
         except Exception as e:
-            raise SparkExpectationsMiscException(f"error occurred while running expectations {e}")
+            raise SparkExpectationsMiscException(f"error occurred while running expectations: {e}")
 
     @staticmethod
     def action_on_rules(

spark_expectations/utils/udf.py

@@ -1,4 +1,4 @@
-from pyspark.sql import SparkSession, Column
+from pyspark.sql import Column
 from pyspark.sql.functions import filter, size, transform, when, lit, array, expr
 
 
@@ -41,18 +41,21 @@ def get_actions_list(column: Column) -> Column:
     action_if_failed = transform(column, lambda x: x["action_if_failed"])
     return when(size(action_if_failed) == 0, array(lit("ignore"))).otherwise(action_if_failed)  # pragma: no cover
 
-def safe_cast(spark: SparkSession, column: str, target_type: str) -> Column:
+
+def safe_cast(ansi_enabled: bool, column: str, target_type: str) -> Column:
     """
-    Checks if ANSI mode is enabled. If enabled, uses try_cast to cast the column to the target type. If not, uses cast.
+    If ANSI mode is enabled, uses try_cast to cast the column to the target type. If not, uses cast.
     Args:
-        spark: SparkSession
-        column: column to cast (provided as a string)
-        target_type: target type to cast to
+        ansi_enabled: bool for if ANSI mode is enabled or not
+        column_expr: column expression to cast (provided as a string that gets parsed as SQL)
+        target_type: target type to cast to (also gets parsed as SQL)
+
+        "column_expr" and "target_type" are interpolated to SQL and parsed by Spark. Never pass user-controlled input, as this is a SQL injection risk.
+        Both must be hardcoded literals or values that have been validated against an allow list.
 
     Returns:
         Column: the casted column
     """
-    ansi_enabled = spark.conf.get("spark.sql.ansi.enabled", "false").lower() == "true"
     if ansi_enabled:
         return expr(f"try_cast({column} as {target_type})")
     else: 

tests/integration/core/test_context.py

@@ -3036,3 +3036,41 @@ def test_get_kafka_write_error_message_default():
     context = SparkExpectationsContext(product_id="product1", spark=spark)
     # Test default value when attribute not set
     assert context.get_kafka_write_error_message == ""
+
+
+def test_get_ansi_enabled_happy_path():
+    """Test getting ANSI mode from spark conf"""
+    from unittest.mock import MagicMock
+
+    mock_spark = MagicMock()
+    mock_spark.conf.get.return_value = "True"
+
+    context = SparkExpectationsContext(product_id="product1", spark=mock_spark)
+    result = context.get_ansi_enabled
+    assert result == True
+
+def test_get_ansi_enabled_already_set():
+    """Test getting ANSI mode when it's already cached--should not call the conf"""
+    from unittest.mock import MagicMock
+
+    mock_spark = MagicMock()
+    context = SparkExpectationsContext(product_id="product1", spark=mock_spark)
+    context._ansi_enabled = False
+
+    result = context.get_ansi_enabled
+    assert result == False
+    mock_spark.conf.get.assert_not_called()
+
+def test_get_ansi_enabled_catches_exception():
+    """Test an exception when getting ANSI mode"""
+    class BadSpark:
+        @property
+        def version(self):
+            raise Exception("boom")
+
+    context = SparkExpectationsContext(product_id="test_product", spark=spark)
+    context.spark = BadSpark()
+
+    result = context.get_ansi_enabled
+    assert result == False
+  

tests/integration/utils/test_actions.py

@@ -806,7 +806,7 @@ def test_agg_query_dq_detailed_result_exception_v2(_fixture_df, _query_dq_rule_e
     _fixture_df.createOrReplaceTempView("query_test_table_target")
     with pytest.raises(
         SparkExpectationsMiscException,
-        match=r"(error occurred while running agg_query_dq_detailed_result Sql query is invalid. *)|(error occurred while running agg_query_dq_detailed_result Regex match not found. *)",
+        match=r"(error occurred while running agg_query_dq_detailed_result: Sql query is invalid. *)|(error occurred while running agg_query_dq_detailed_result: Regex match not found. *)",
     ):
         SparkExpectationsActions().agg_query_dq_detailed_result(
             _fixture_mock_context, _query_dq_rule_exception, _fixture_df, []
@@ -1028,7 +1028,7 @@ def test_run_dq_rules_query(
 def test_run_dq_rules_negative_case(_fixture_df, _fixture_mock_context):
     expectations = {"row_dq_rules": []}
 
-    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running expectations .*"):
+    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running expectations: .*"):
         SparkExpectationsActions.run_dq_rules(_fixture_mock_context, _fixture_df, expectations, "row_dq")
 
 @pytest.mark.parametrize(
@@ -1054,7 +1054,7 @@ def test_run_dq_rules_negative_case(_fixture_df, _fixture_mock_context):
 )
 def test_run_dq_rules_exception(_fixture_df, _fixture_mock_context, expectations, expected_exception):
     # test the exception functionality in run_dq_rules with faulty user input
-    with pytest.raises(expected_exception, match=r"error occurred while running expectations .*"):
+    with pytest.raises(expected_exception, match=r"error occurred while running expectations: .*"):
         SparkExpectationsActions.run_dq_rules(_fixture_mock_context, _fixture_df, expectations, "row_dq")
 
 
@@ -1080,7 +1080,7 @@ def test_run_dq_rules_condition_expression_exception(
     }
     _fixture_df.createOrReplaceTempView("query_test_table")
 
-    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running expectations .*"):
+    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running expectations: .*"):
         SparkExpectationsActions.run_dq_rules(
             _fixture_mock_context, _fixture_df, _expectations, "query_dq", False, True
         )
@@ -1113,7 +1113,7 @@ def test_run_dq_rules_condition_expression_dynamic_exception(
     }
     _fixture_df.createOrReplaceTempView("query_test_table")
 
-    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running expectations .*"):
+    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running expectations: .*"):
         _rule_type = _rule_test.get("rule_type")
         SparkExpectationsActions.run_dq_rules(
             _fixture_mock_context, _fixture_df, _expectations, _rule_type, False, True
@@ -1129,7 +1129,7 @@ def collect(self):
                     return [[["unexpected", "list"]]]  # Not int, float, str, or date
             return DummyRow()
     dummy_df = DummyDF()
-    with pytest.raises(SparkExpectationsMiscException, match="error occurred while running agg_query_dq_detailed_result .*"):
+    with pytest.raises(SparkExpectationsMiscException, match="error occurred while running agg_query_dq_detailed_result: .*"):
         SparkExpectationsActions.agg_query_dq_detailed_result(
             _fixture_mock_context, _fixture_agg_dq_rule, dummy_df, []
         )
@@ -1215,7 +1215,7 @@ def test_create_agg_dq_results(input_df, rule_type_name, expected_output, _fixtu
 )
 def test_create_agg_dq_results_exception(input_df, _fixture_mock_context):
     # faulty user input is given to test the exception functionality of the agg_dq_result
-    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running create agg dq results .*"):
+    with pytest.raises(SparkExpectationsMiscException, match=r"error occurred while running create_agg_dq_results: .*"):
         SparkExpectationsActions().create_agg_dq_results(
             _fixture_mock_context,
             input_df,
@@ -1961,7 +1961,7 @@ def test_agg_query_dq_detailed_result_type_error_line_210(_fixture_agg_dq_rule,
         mock_result = Mock()
         mock_result.collect.return_value = [[[]]]  # Return list to trigger TypeError
         mock_agg.return_value = mock_result
-        with pytest.raises(SparkExpectationsMiscException, match="error occurred while running agg_query_dq_detailed_result .*"):
+        with pytest.raises(SparkExpectationsMiscException, match="error occurred while running agg_query_dq_detailed_result: .*"):
             SparkExpectationsActions.agg_query_dq_detailed_result(
                 _fixture_mock_context, _fixture_agg_dq_rule, df, []
             )

tests/integration/utils/test_udf.py

@@ -62,13 +62,13 @@ def test_get_actions_list():
 def test_safe_cast_ansi_disabled():
     # with default spark.sql.ansi.enabled=false, safe cast should return
     #  a column equavalent to cast(col as type)
-    spark.conf.set("spark.sql.ansi.enabled", "false")
 
+    ansi_enabled = "false"
     columns = ["Name", "Value"]
     data = [("thing", "123")]
     df = spark.createDataFrame(data, schema=columns)
 
-    result_df = df.withColumn("casted_value", safe_cast(spark, "Value", "bigint"))
+    result_df = df.withColumn("casted_value", safe_cast(ansi_enabled, "Value", "bigint"))
     result = result_df.select("casted_value").collect()[0]["casted_value"]
 
     assert result == 123
@@ -78,16 +78,16 @@ def test_safe_cast_ansi_disabled():
 def test_safe_cast_ansi_enabled():
     # with spark.sql.ansi.enabled=true, safe cast should succeed if it's a castable value and should 
     #  return Null if it's not a valid cast (without safe cast it would throw an exception insead)
-    spark.conf.set("spark.sql.ansi.enabled", "true")
 
+    ansi_enabled = "true"
     columns = ["Name", "Value", "Color"]
     data = [("mug", "10", "red")]
     df = spark.createDataFrame(data, schema=columns)
 
-    success_result_df = df.withColumn("success_casted_value", safe_cast(spark, "Value", "double"))
+    success_result_df = df.withColumn("success_casted_value", safe_cast(ansi_enabled, "Value", "double"))
     success_result = success_result_df.select("success_casted_value").collect()[0]["success_casted_value"]
 
-    fail_result_df = df.withColumn("fail_casted_value", safe_cast(spark, "Color", "double"))
+    fail_result_df = df.withColumn("fail_casted_value", safe_cast(ansi_enabled, "Color", "double"))
     fail_result = fail_result_df.select("fail_casted_value").collect()[0]["fail_casted_value"]
 
     # re-set ansi mode so it doesn't affect other tests

tests/unit/utils/test_actions.py

@@ -77,7 +77,7 @@ def test_agg_query_dq_detailed_result_exception():
     # faulty user input is given to test the exception functionality of the agg_query_dq_detailed_result
 
     with pytest.raises(
-        SparkExpectationsMiscException, match=r"error occurred while running agg_query_dq_detailed_result .*"
+        SparkExpectationsMiscException, match=r"error occurred while running agg_query_dq_detailed_result: .*"
     ):
         SparkExpectationsActions().agg_query_dq_detailed_result(
             _mock_object_context, "_fixture_query_dq_rule", "<df>", []