Guard check filter and row_filter against unsafe SQL

docs/dqx/docs/guide/quality_checks_definition.mdx

@@ -342,7 +342,7 @@ Checks defined using declarative syntax must contain the following fields:
    - `arguments`: keyword arguments passed to the check function. Every parameter without a default in the function’s Python signature must appear here.
    - `for_each_column`: (optional) list of column names or expressions for which the same check should be applied. When saving with `DQEngine.save_checks` to Delta or Lakebase tables, DQX stores such checks in compact format (one row per check, `for_each_column` preserved).
 - (optional) `name`: name of the check: autogenerated if not provided.
-- (optional) `filter`: spark expression to filter the rows for which the check is applied (e.g. `"business_unit = 'Finance'"`). The check function will run only on the rows matching the filter condition. The condition can reference any column of the validated dataset, not only the one where you apply the check function. When using dataset-level checks, the filter condition is pushed down as `row_filter` to the check function and applied before aggregation, ensuring that the check operates only on the relevant subset of rows rather than on the aggregated results.
+- (optional) `filter`: spark expression to filter the rows for which the check is applied (e.g. `"business_unit = 'Finance'"`). The check function will run only on the rows matching the filter condition. The condition can reference any column of the validated dataset, not only the one where you apply the check function. When using dataset-level checks, the filter condition is pushed down as `row_filter` to the check function and applied before aggregation, ensuring that the check operates only on the relevant subset of rows rather than on the aggregated results. For security, `filter` and `row_filter` are restricted to simple predicates: they cannot contain a `SELECT` statement, and an unsafe filter raises an error. If you need arbitrary SQL (including subqueries), use the `sql_expression` check instead, and only accept check definitions from trusted sources in automated or multi-tenant pipelines.
 - (optional) `user_metadata`: key-value pairs added to the row-level warnings and errors
 
 ## YAML format (declarative approach)
@@ -699,7 +699,7 @@ The table used to store checks has the following structure. You can use `for_eac
 | └─ `function`          | `string`              | Name of the DQX check function to apply.                                                                                                                  |
 | └─ `for_each_column`   | `array<string>`       | (Optional) List of columns for which the same check applies. Stored in compact format (one row per check). Manual tables can use compact or expanded (individual rules instead of for_each_column) format. |
 | └─ `arguments`         | `map<string, string>` | Keyword arguments for the check function. All parameters without a default must be present (or supplied via `for_each_column` for `column` / `columns`).  |
-| `filter`               | `string`              | (Optional) Spark SQL expression to filter rows to which the check is applied, e.g. `"business_unit = 'Finance'"`. The check function will run only on the rows matching the filter condition. The condition can reference any column of the validated dataset, not only the one where you apply the check function. When using dataset-level checks, the filter condition is pushed down as `row_filter` to the check function and applied before aggregation, ensuring that the check operates only on the relevant subset of rows rather than on the aggregated results. |
+| `filter`               | `string`              | (Optional) Spark SQL expression to filter rows to which the check is applied, e.g. `"business_unit = 'Finance'"`. The check function will run only on the rows matching the filter condition. The condition can reference any column of the validated dataset, not only the one where you apply the check function. When using dataset-level checks, the filter condition is pushed down as `row_filter` to the check function and applied before aggregation, ensuring that the check operates only on the relevant subset of rows rather than on the aggregated results. For security, `filter` and `row_filter` are restricted to simple predicates: they cannot contain a `SELECT` statement, and an unsafe filter raises an error. If you need arbitrary SQL (including subqueries), use the `sql_expression` check instead, and only accept check definitions from trusted sources in automated or multi-tenant pipelines.|
 | `run_config_name`      | `string`              | Name of the run config name. Could be any string such as input table or job name (use "default" if not provided). Useful for selecting applicable checks. |
 | `user_metadata`        | `map<string, string>` | (Optional) Custom metadata to add to any row-level warnings or errors generated by the check.                                                             |
 | `created_at`           | `timestamp`           | Current local date and time in UTC when the rule set was saved.                                                        |

src/databricks/labs/dqx/anomaly/scoring_utils.py

@@ -13,6 +13,7 @@
 from databricks.labs.dqx.anomaly.anomaly_info_schema import ai_explanation_struct_schema, anomaly_info_struct_schema
 from databricks.labs.dqx.anomaly.segment_utils import canonicalize_segment_values
 from databricks.labs.dqx.errors import InvalidParameterError
+from databricks.labs.dqx.utils import safe_filter_expr
 from databricks.labs.dqx.schema.dq_info_schema import (
     build_dq_info_struct,
     register_dq_info_field,
@@ -248,4 +249,4 @@ def apply_row_filter(df: DataFrame, row_filter: str | None) -> DataFrame:
     row_filter is a SQL expression (e.g. \"region = 'US'\"). Only these rows are run
     through anomaly detection; elsewhere we join results back so output has same row count.
     """
-    return df.filter(F.expr(row_filter)) if row_filter else df
+    return df.filter(safe_filter_expr(row_filter)) if row_filter else df

src/databricks/labs/dqx/check_funcs.py

@@ -18,6 +18,7 @@
 from databricks.labs.dqx.utils import (
     get_column_name_or_alias,
     is_sql_query_safe,
+    safe_filter_expr,
     normalize_col_str,
     get_columns_as_strings,
     to_lowercase,
@@ -430,6 +431,11 @@ def sql_expression(
 
     Args:
         expression: SQL expression. Fail if expression evaluates to False, pass if it evaluates to True.
+            Security note: this parameter accepts arbitrary SQL and is evaluated as-is, so it may
+            include subqueries and run with the permissions of the process executing the checks.
+            Only use check definitions from trusted sources, especially in automated or multi-tenant
+            pipelines. Unlike *filter* and *row_filter*, SELECT is not blocked here because subqueries
+            can be legitimate in an SQL expression.
         msg: optional message of the *Column* type, automatically generated if None
         name: optional name of the resulting column, automatically generated if None
         negate: if the condition should be negated (true) or not. For example, "col is not null" will mark null
@@ -1258,7 +1264,7 @@ def apply(df: DataFrame) -> DataFrame:
                 f"Column '{col_expr_str}' must be of numeric type to perform outlier detection using MAD method, "
                 f"but got type '{column_type.simpleString()}' instead."
             )
-        filter_condition = F.expr(row_filter) if row_filter else F.lit(True)
+        filter_condition = safe_filter_expr(row_filter)
         median, mad = _calculate_median_absolute_deviation(df, col_expr_str, row_filter)
         if median is not None and mad is not None:
             median = float(median)
@@ -1349,7 +1355,7 @@ def apply(df: DataFrame) -> DataFrame:
 
         filter_condition = F.lit(True)
         if row_filter:
-            filter_condition = filter_condition & F.expr(row_filter)
+            filter_condition = filter_condition & safe_filter_expr(row_filter)
 
         if nulls_distinct:
             # All columns must be non-null
@@ -1478,7 +1484,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) ->
         ref_alias = f"__ref_{col_str_norm}_{unique_str}"
         ref_df_distinct = ref_df.select(ref_col_expr.alias(ref_alias)).distinct()
 
-        filter_expr = F.expr(row_filter) if row_filter else F.lit(True)
+        filter_expr = safe_filter_expr(row_filter)
 
         # col_expr.isNotNull() only filters rows in the single-column non-null-safe path;
         # when col_expr is a struct (composite keys or null_safe=True), the struct is never NULL
@@ -1595,7 +1601,7 @@ def _replace_template(sql: str, replacements: dict[str, str]) -> str:
     def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> DataFrame:
         filtered_df = df
         if row_filter:
-            filtered_df = df.filter(F.expr(row_filter))
+            filtered_df = df.filter(safe_filter_expr(row_filter))
 
         # since the check could be applied multiple times, the views created here must be unique
         filtered_df.createOrReplaceTempView(unique_input_view)
@@ -1968,7 +1974,7 @@ def apply(df: DataFrame) -> DataFrame:
                 f"but got type '{time_col_type.simpleString()}' instead."
             )
 
-        filter_col = F.expr(row_filter) if row_filter else F.lit(True)
+        filter_col = filter_col = safe_filter_expr(row_filter)
         filtered_expr = F.when(filter_col, aggr_col_expr) if row_filter else aggr_col_expr
         aggr_expr = _build_aggregate_expression(aggr_type, filtered_expr, aggr_params)
 
@@ -2202,7 +2208,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) ->
         skipped_columns = [col for col in df.columns if col not in compare_columns and col not in pk_column_names]
 
         # apply filter before aliasing to avoid ambiguity
-        df = df.withColumn(filter_col, F.expr(row_filter) if row_filter else F.lit(True))
+        df = df.withColumn(filter_col, safe_filter_expr(row_filter))
 
         df = df.alias("df")
         ref_df = ref_df.alias("ref_df")
@@ -2309,7 +2315,7 @@ def apply(df: DataFrame) -> DataFrame:
         # Build filter condition
         filter_condition = F.lit(True)
         if row_filter:
-            filter_condition = filter_condition & F.expr(row_filter)
+            filter_condition = filter_condition & safe_filter_expr(row_filter)
 
         # Limit checking to be within the lookback window if needed
         if lookback_windows is not None:
@@ -3372,7 +3378,7 @@ def apply(df: DataFrame) -> DataFrame:
         Returns:
             The DataFrame with additional condition and metric columns for aggregation validation.
         """
-        filter_col = F.expr(row_filter) if row_filter else F.lit(True)
+        filter_col = safe_filter_expr(row_filter)
         filtered_expr = F.when(filter_col, aggr_col_expr) if row_filter else aggr_col_expr
 
         # Build aggregation expression
@@ -3890,7 +3896,7 @@ def _apply_dataset_level_sql_check(
     # - If row_filter is provided, only rows matching the filter get the condition value
     # - Rows not matching the filter get None (consistent with row-level checks)
     if row_filter:
-        filter_expr = F.expr(row_filter)
+        filter_expr = safe_filter_expr(row_filter)
         result_df = df.withColumn(
             unique_condition_column, F.when(filter_expr, F.lit(condition_value)).otherwise(F.lit(None))
         )

src/databricks/labs/dqx/geo/check_funcs.py

@@ -9,6 +9,7 @@
 from databricks.labs.dqx.rule import register_rule
 from databricks.labs.dqx.check_funcs import make_condition, get_normalized_column_and_expr, get_limit_expr
 from databricks.labs.dqx.errors import InvalidParameterError
+from databricks.labs.dqx.utils import safe_filter_expr
 
 POINT_TYPE = "ST_Point"
 LINESTRING_TYPE = "ST_LineString"
@@ -915,7 +916,7 @@ def apply(df: DataFrame) -> DataFrame:
         df_with_hash = df.withColumn(row_id_col, F.xxhash64(F.struct("*")))
 
         # Apply filter after hashing (so hash is available for final join to preserve all rows)
-        filtered_df = df_with_hash.filter(F.expr(row_filter)) if row_filter is not None else df_with_hash
+        filtered_df = df_with_hash.filter(safe_filter_expr(row_filter)) if row_filter is not None else df_with_hash
 
         duplicates_df = filtered_df.groupBy(row_id_col).count().filter("count > 1").select(row_id_col)
 

src/databricks/labs/dqx/manager.py

@@ -13,7 +13,7 @@
     DQRule,
 )
 from databricks.labs.dqx.schema.dq_result_schema import dq_result_item_schema
-from databricks.labs.dqx.utils import get_column_name_or_alias
+from databricks.labs.dqx.utils import get_column_name_or_alias, safe_filter_expr
 
 logger = logging.getLogger(__name__)
 
@@ -68,7 +68,7 @@ def filter_condition(self) -> Column:
         """
         Returns the filter condition for the check.
         """
-        return F.expr(self.check.filter) if self.check.filter else F.lit(True)
+        return safe_filter_expr(self.check.filter)
 
     @cached_property
     def invalid_columns(self) -> list[str]:

src/databricks/labs/dqx/utils.py

@@ -23,7 +23,7 @@
 import pyspark.sql.functions as F
 from databricks.sdk import WorkspaceClient
 from databricks.labs.blueprint.limiter import rate_limited
-from databricks.labs.dqx.errors import InvalidParameterError
+from databricks.labs.dqx.errors import InvalidParameterError, UnsafeSqlQueryError
 from databricks.labs.dqx.table_manager import SparkTableDataProvider
 from databricks.sdk.errors import NotFound
 
@@ -214,7 +214,7 @@ def normalize_col_str(col_str: str) -> str:
     return re.sub(COLUMN_NORMALIZE_EXPRESSION, "_", col_str[:max_chars].lower()).rstrip("_")
 
 
-def is_sql_query_safe(query: str) -> bool:
+def is_sql_query_safe(query: str, forbid_select: bool = False) -> bool:
     # Normalize the query by removing extra whitespace and converting to lowercase
     normalized_query = re.sub(r"\s+", " ", query).strip().lower()
 
@@ -237,9 +237,32 @@ def is_sql_query_safe(query: str) -> bool:
         "optimize",
         "zorder",
     ]
+    if forbid_select:
+        forbidden_statements = forbidden_statements + ["select"]
     return not any(re.search(rf"\b{kw}\b", normalized_query) for kw in forbidden_statements)
 
 
+def safe_filter_expr(filter_expr: str | None) -> Column:
+    """Build a Spark column from a filter expression, rejecting unsafe SQL.
+
+    Validates the filter with *is_sql_query_safe* using *forbid_select=True* before
+    compiling it, since a filter must be a simple predicate and never a full query
+    (e.g. a subquery). Used for both check filters and *row_filter* parameters.
+
+    Args:
+        filter_expr: The filter predicate as a string, or None.
+
+    Returns:
+        The compiled filter column, or a literal true column when no filter is given.
+
+    Raises:
+        UnsafeSqlQueryError: If the filter contains a forbidden statement such as SELECT.
+    """
+    if filter_expr and not is_sql_query_safe(filter_expr, forbid_select=True):
+        raise UnsafeSqlQueryError(f"Unsafe filter expression: '{filter_expr}'")
+    return F.expr(filter_expr) if filter_expr else F.lit(True)
+
+
 def safe_json_load(value: str):
     """
     Safely load a JSON string, returning the original value if it fails to parse.

tests/unit/test_manager.py

@@ -1,12 +1,14 @@
 from unittest.mock import create_autospec, PropertyMock
+import pytest
 
 from pyspark.errors import AnalysisException
 from pyspark.sql import DataFrame, SparkSession
 
 from databricks.labs.dqx import check_funcs
 from databricks.labs.dqx.executor import DQCheckResult
 from databricks.labs.dqx.manager import DQRuleManager
-from databricks.labs.dqx.rule import DQRowRule
+from databricks.labs.dqx.rule import DQRowRule, DQDatasetRule
+from databricks.labs.dqx.errors import UnsafeSqlQueryError
 
 
 def _make_manager_with_missing_column(df: DataFrame, spark: SparkSession, *, suppress_skipped: bool) -> DQRuleManager:
@@ -55,3 +57,68 @@ def test_rule_manager_suppress_skipped_false_returns_struct_for_invalid_column()
     assert result.condition is not None
     assert "skipped" in str(result.condition).lower()
     assert result.check_df is df_mock
+
+
+def test_rule_manager_raises_on_unsafe_filter():
+    """A check whose filter contains a forbidden statement (e.g. SELECT) is rejected
+    end-to-end: process() raises UnsafeSqlQueryError via the wired safe_filter_expr guard."""
+    df_mock = create_autospec(DataFrame)
+    spark_mock = create_autospec(SparkSession)
+    manager = DQRuleManager(
+        check=DQRowRule(
+            check_func=check_funcs.is_not_null,
+            column="col1",
+            filter="id IN (SELECT id FROM users)",
+        ),
+        df=df_mock,
+        spark=spark_mock,
+        engine_user_metadata={},
+        run_time_overwrite=None,
+        run_id="test-run",
+        suppress_skipped=False,
+    )
+    with pytest.raises(UnsafeSqlQueryError):
+        manager.process()
+
+
+def test_rule_manager_allows_safe_filter():
+    """A normal predicate filter passes the guard (does not raise)."""
+    df_mock = create_autospec(DataFrame)
+    spark_mock = create_autospec(SparkSession)
+    manager = DQRuleManager(
+        check=DQRowRule(
+            check_func=check_funcs.is_not_null,
+            column="col1",
+            filter="country = 'US'",
+        ),
+        df=df_mock,
+        spark=spark_mock,
+        engine_user_metadata={},
+        run_time_overwrite=None,
+        run_id="test-run",
+        suppress_skipped=False,
+    )
+    # accessing filter_condition compiles the filter; a safe one must not raise
+    assert manager.filter_condition is not None
+
+
+def test_rule_manager_raises_on_unsafe_filter_for_dataset_check():
+    """A dataset check's filter is pushed down to row_filter; an unsafe one is rejected
+    end-to-end through the check_funcs safe_filter_expr wiring."""
+    df_mock = create_autospec(DataFrame)
+    spark_mock = create_autospec(SparkSession)
+    manager = DQRuleManager(
+        check=DQDatasetRule(
+            check_func=check_funcs.is_unique,
+            columns=["col1"],
+            filter="id IN (SELECT id FROM users)",
+        ),
+        df=df_mock,
+        spark=spark_mock,
+        engine_user_metadata={},
+        run_time_overwrite=None,
+        run_id="test-run",
+        suppress_skipped=False,
+    )
+    with pytest.raises(UnsafeSqlQueryError):
+        manager.process()

tests/unit/test_utils.py

@@ -14,6 +14,7 @@
 from databricks.labs.dqx.utils import (
     get_column_name_or_alias,
     is_sql_query_safe,
+    safe_filter_expr,
     normalize_col_str,
     safe_json_load,
     get_columns_as_strings,
@@ -24,7 +25,7 @@
     resolve_variables,
 )
 from databricks.labs.dqx.rule import normalize_bound_args
-from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError
+from databricks.labs.dqx.errors import InvalidParameterError, InvalidConfigError, UnsafeSqlQueryError
 from databricks.labs.dqx.config import InputConfig
 from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig
 
@@ -258,6 +259,32 @@ def test_mixed_content():
     assert not is_sql_query_safe("WITH cte AS (UPDATE users SET x=1) SELECT * FROM cte")
 
 
+def test_select_blocked_when_forbid_select():
+    assert not is_sql_query_safe("id IN (SELECT id FROM users)", forbid_select=True)
+
+
+def test_select_allowed_by_default():
+    assert is_sql_query_safe("SELECT id FROM users WHERE active = true")
+
+
+def test_select_blocked_case_insensitive_when_forbid_select():
+    assert not is_sql_query_safe("id IN (sElEcT id FROM t)", forbid_select=True)
+
+
+def test_select_substring_safe_when_forbid_select():
+    # word-boundary: identifiers that merely contain "select" must NOT be flagged
+    assert is_sql_query_safe("select_flag = true AND selected_count > 0", forbid_select=True)
+
+
+def test_plain_predicate_safe_when_forbid_select():
+    assert is_sql_query_safe("country = 'US' AND amount > 100", forbid_select=True)
+
+
+def test_safe_filter_expr_raises_on_select():
+    with pytest.raises(UnsafeSqlQueryError):
+        safe_filter_expr("id IN (SELECT id FROM users)")
+
+
 def test_safe_json_load_dict():
     value = '{"key": "value"}'
     result = safe_json_load(value)