fix(core): qualify check_filter columns in valid_reference_data JOIN (DTL-1517)

soda-core/src/soda_core/common/sql_dialect.py

@@ -100,7 +100,10 @@
     SqlExpression,
     SqlExpressionStr,
 )
-from soda_core.common.sql_utils import apply_sampling_to_sql
+from soda_core.common.sql_utils import (
+    apply_sampling_to_sql,
+    qualify_unqualified_columns_with_alias,
+)
 from soda_core.common.statements.table_types import FullyQualifiedObjectName, TableType
 from typing_extensions import deprecated
 
@@ -1452,6 +1455,14 @@ def apply_sampling(
             write_dialect=self.SQLGLOT_DIALECT,
         )
 
+    def qualify_unqualified_columns_with_alias(self, sql_expression: str, alias: str) -> str:
+        return qualify_unqualified_columns_with_alias(
+            sql_expression=sql_expression,
+            alias=alias,
+            read_dialect=self.SQLGLOT_DIALECT,
+            write_dialect=self.SQLGLOT_DIALECT,
+        )
+
     ########################################################
     # Metadata columns query
     ########################################################

soda-core/src/soda_core/common/sql_utils.py

@@ -37,6 +37,52 @@ def attach_sample_to_relation(rel: exp.Expression, sampler_limit: Number, sample
         rel.set("sample", build_sample_clause(sampler_limit, sampler_type))
 
 
+def qualify_unqualified_columns_with_alias(
+    sql_expression: str,
+    alias: str,
+    read_dialect: str | None = None,
+    write_dialect: str | None = None,
+) -> str:
+    """
+    Qualify every unqualified top-level column reference in a SQL expression with the given
+    table alias.
+
+    Already-qualified columns are left untouched. Columns inside subqueries are also left
+    untouched — they belong to the subquery's own scope, not the outer alias. If sqlglot
+    cannot parse or print the expression for any reason (parse error, unknown dialect),
+    the input is returned unchanged so any database error surfaces unmodified to the user.
+
+    Note: a successful parse round-trips through sqlglot's printer, which canonicalises some
+    syntax (e.g. ``col::int`` becomes ``CAST(col AS INT)``, function names get uppercased).
+    On the same read/write dialect the rewrite is semantically equivalent in standard cases,
+    but recognised functions may be rewritten into canonical equivalents.
+
+    If ``read_dialect`` / ``write_dialect`` are not provided (or are empty), sqlglot's
+    default parser/printer is used.
+    """
+    if not sql_expression or not sql_expression.strip():
+        return sql_expression
+
+    try:
+        tree = (
+            sqlglot.parse_one(sql_expression, read=read_dialect) if read_dialect else sqlglot.parse_one(sql_expression)
+        )
+
+        for column in tree.find_all(exp.Column):
+            if column.table:
+                continue
+            # Skip columns inside any nested SELECT scope (subqueries via exp.Subquery,
+            # EXISTS bodies via exp.Exists, scalar subqueries, CTEs). These resolve in
+            # their own scope first and shouldn't be silently aliased to the outer one.
+            if column.find_ancestor(exp.Select):
+                continue
+            column.set("table", exp.to_identifier(alias, quoted=True))
+
+        return tree.sql(dialect=write_dialect) if write_dialect else tree.sql()
+    except Exception:
+        return sql_expression
+
+
 def apply_sampling_to_sql(
     sql: str,
     sampler_limit: Number,

soda-core/src/soda_core/contracts/impl/check_types/invalidity_check.py

@@ -244,21 +244,36 @@ def __init__(
 
         self._cte = cte
 
-        sql_ast = self.build_query(cte=self._cte, select_clause=SELECT(COUNT(STAR())))
+        sql_ast = self.build_query(cte=self._cte, select_clause=SELECT(COUNT(STAR())), check_filter=self.check_filter)
         self.sql = self.data_source_impl.sql_dialect.build_select_sql(sql_ast)
 
-    def build_query(self, cte: CTE, select_clause: SqlExpression) -> list[SqlExpression]:
+    def build_query(
+        self, cte: CTE, select_clause: SqlExpression, check_filter: Optional[str] = None
+    ) -> list[SqlExpression]:
         query = [
             WITH([cte, self.referenced_cte()]),
             select_clause,
             FROM(cte.alias).AS(self.referencing_alias),
-            WHERE.optional(SqlExpressionStr.optional(self.check_filter)),
+            WHERE.optional(SqlExpressionStr.optional(check_filter)),
         ]
 
         query.extend(self.query_join())
 
         return query
 
+    def _build_qualified_filter_sql(self) -> str:
+        """Re-render the JOIN query with check_filter columns qualified against the
+        source alias. Used as a one-shot retry when the original query errors."""
+        qualified_filter: str = self.data_source_impl.sql_dialect.qualify_unqualified_columns_with_alias(
+            self.check_filter, self.referencing_alias
+        )
+        sql_ast = self.build_query(
+            cte=self._cte,
+            select_clause=SELECT(COUNT(STAR())),
+            check_filter=qualified_filter,
+        )
+        return self.data_source_impl.sql_dialect.build_select_sql(sql_ast)
+
     def referenced_cte(self) -> CTE:
         valid_reference_data: ValidReferenceData = self.metric_impl.missing_and_validity.valid_reference_data
         referenced_dataset_name: str = valid_reference_data.dataset_name
@@ -332,9 +347,39 @@ def query_join(self) -> SqlExpression:
     def execute(self) -> list[Measurement]:
         try:
             query_result: QueryResult = self.data_source_impl.execute_query(self.sql)
-        except Exception as e:
-            logger.error(msg=f"Could not execute invalid reference query {self.sql}: {e}", exc_info=True)
-            return []
+        except Exception as initial_error:
+            # The most common cause of failure here is an ambiguous-column error when the
+            # check_filter references a name that exists in both the source and reference
+            # datasets (the JOIN exposes both). Engine error wording varies, so instead of
+            # sniffing the message we retry once with the filter columns qualified against
+            # the source alias. If the retry also fails (or there is no filter to retry
+            # with), surface the original error and, where applicable, the retry error too.
+            if not self.check_filter:
+                logger.error(
+                    msg=f"Could not execute invalid reference query {self.sql}: {initial_error}",
+                    exc_info=True,
+                )
+                return []
+
+            retry_sql: str = self._build_qualified_filter_sql()
+            try:
+                query_result = self.data_source_impl.execute_query(retry_sql)
+            except Exception as retry_error:
+                logger.error(
+                    msg=(
+                        f"Could not execute invalid reference query.\n"
+                        f"Initial query failed: {self.sql}\n"
+                        f"  Initial error: {initial_error}\n"
+                        f"Retried with the check_filter columns qualified against source alias "
+                        f"{self.referencing_alias!r} (common fix for ambiguous-column errors when "
+                        f"the filter references a column that exists in both the source and "
+                        f"reference datasets); the retry also failed.\n"
+                        f"Retry query: {retry_sql}\n"
+                        f"  Retry error: {retry_error}"
+                    ),
+                    exc_info=True,
+                )
+                return []
 
         metric_value = query_result.rows[0][0]
         metric_impl: MetricImpl = self.metrics[0]

soda-tests/tests/integration/test_invalid_reference_check.py

@@ -43,6 +43,25 @@
 )
 
 
+# Reference table that intentionally shares an `id` column with the source
+# (`referencing_table_specification`). Used to exercise the case where a
+# check-level filter on `id` becomes ambiguous in the JOIN-style invalid
+# reference query.
+referenced_table_with_shared_id_specification = (
+    TestTableSpecification.builder()
+    .table_purpose("invalid_referenced_shared_id")
+    .column_integer("id")
+    .column_varchar("country_code", 2)
+    .rows(
+        rows=[
+            (100, "NL"),
+            (101, "BE"),
+        ]
+    )
+    .build()
+)
+
+
 def test_invalid_count(data_source_test_helper: DataSourceTestHelper):
     # https://dev.sodadata.io/o/f35cb402-ad17-4aca-9166-02c9eb75c979/datasets/2945ba9d-b1ff-4cfd-b277-d5e4edfa2bd5/checks
 
@@ -114,6 +133,31 @@ def test_invalid_count_with_check_filter(data_source_test_helper: DataSourceTest
     assert get_diagnostic_value(check_result, "check_rows_tested") == 3
 
 
+def test_invalid_count_with_check_filter_on_shared_column(data_source_test_helper: DataSourceTestHelper):
+    referencing_table = data_source_test_helper.ensure_test_table(referencing_table_specification)
+    referenced_table = data_source_test_helper.ensure_test_table(referenced_table_with_shared_id_specification)
+
+    id_quoted: str = data_source_test_helper.quote_column("id")
+
+    contract_verification_result: ContractVerificationResult = data_source_test_helper.assert_contract_fail(
+        test_table=referencing_table,
+        contract_yaml_str=f"""
+            columns:
+              - name: country
+                valid_reference_data:
+                  dataset: {data_source_test_helper.build_dqn(referenced_table)}
+                  column: country_code
+                checks:
+                  - invalid:
+                      filter: |
+                        {id_quoted} > 1
+        """,
+    )
+    check_result: CheckResult = contract_verification_result.check_results[0]
+    assert get_diagnostic_value(check_result, "invalid_count") == 3
+    assert get_diagnostic_value(check_result, "check_rows_tested") == 7
+
+
 def test_invalid_count_with_check_and_dataset_filter(data_source_test_helper: DataSourceTestHelper):
     referencing_table = data_source_test_helper.ensure_test_table(referencing_table_specification)
     referenced_table = data_source_test_helper.ensure_test_table(referenced_table_specification)

soda-tests/tests/unit/test_sql_utils.py

@@ -7,7 +7,10 @@
 import pytest
 import sqlglot
 from soda_core.common.metadata_types import SamplerType
-from soda_core.common.sql_utils import apply_sampling_to_sql
+from soda_core.common.sql_utils import (
+    apply_sampling_to_sql,
+    qualify_unqualified_columns_with_alias,
+)
 from sqlglot import exp
 
 
@@ -259,3 +262,107 @@ def test_apply_sampling_complex_query() -> None:
     subquery = _get_first_subquery_with_alias(tree, "subquery")
     subquery_sample = subquery.args.get("sample")
     assert subquery_sample is None
+
+
+# ------------------------------
+# qualify_unqualified_columns_with_alias
+# ------------------------------
+
+
+def test_qualify_basic_unqualified_column() -> None:
+    out = qualify_unqualified_columns_with_alias("id > 1", "C", read_dialect="postgres", write_dialect="postgres")
+    assert out == '"C".id > 1'
+
+
+def test_qualify_leaves_already_qualified_columns_alone() -> None:
+    out = qualify_unqualified_columns_with_alias(
+        '"R".country_code = country', "C", read_dialect="postgres", write_dialect="postgres"
+    )
+    assert out == '"R".country_code = "C".country'
+
+
+def test_qualify_does_not_touch_string_literals() -> None:
+    # 'foo' is a string literal — only the column `name` should be qualified
+    out = qualify_unqualified_columns_with_alias(
+        "lower(name) = 'foo'", "C", read_dialect="postgres", write_dialect="postgres"
+    )
+    assert "'foo'" in out
+    assert '"C".name' in out
+
+
+@pytest.mark.parametrize("expr", ["", "   ", "\n\t"])
+def test_qualify_empty_input_returned_unchanged(expr: str) -> None:
+    out = qualify_unqualified_columns_with_alias(expr, "C", read_dialect="postgres", write_dialect="postgres")
+    assert out == expr
+
+
+def test_qualify_unparseable_input_returned_unchanged() -> None:
+    expr = "(((unbalanced"
+    out = qualify_unqualified_columns_with_alias(expr, "C", read_dialect="postgres", write_dialect="postgres")
+    assert out == expr
+
+
+def test_qualify_unknown_dialect_returned_unchanged() -> None:
+    # Unknown dialects raise ValueError inside sqlglot; the helper must not propagate it.
+    expr = "id > 1"
+    out = qualify_unqualified_columns_with_alias(
+        expr, "C", read_dialect="totally_fake_dialect", write_dialect="totally_fake_dialect"
+    )
+    assert out == expr
+
+
+def test_qualify_skips_columns_inside_in_subquery() -> None:
+    # Outer `id` belongs to the source and should be qualified.
+    # Inner `user_id` belongs to `other` and must NOT be aliased to "C".
+    out = qualify_unqualified_columns_with_alias(
+        "id IN (SELECT user_id FROM other)", "C", read_dialect="postgres", write_dialect="postgres"
+    )
+    assert '"C".id' in out
+    assert '"C".user_id' not in out
+
+
+def test_qualify_skips_columns_inside_exists_subquery() -> None:
+    # EXISTS bodies are wrapped in exp.Exists > exp.Select (not exp.Subquery).
+    # Inner unqualified `id` must NOT be aliased — it resolves in the subquery's own scope.
+    out = qualify_unqualified_columns_with_alias(
+        "EXISTS (SELECT 1 FROM other WHERE id = 1)", "C", read_dialect="postgres", write_dialect="postgres"
+    )
+    assert '"C".id' not in out
+
+
+def test_qualify_skips_columns_inside_scalar_subquery() -> None:
+    # Outer `id` qualified; inner aggregate's `x` left alone.
+    out = qualify_unqualified_columns_with_alias(
+        "(SELECT MAX(x) FROM other) > id", "C", read_dialect="postgres", write_dialect="postgres"
+    )
+    assert '"C".id' in out
+    assert '"C".x' not in out
+
+
+@pytest.mark.parametrize(
+    "dialect,expected_substring",
+    [
+        ("postgres", '"C".id'),
+        ("snowflake", '"C".id'),
+        ("duckdb", '"C".id'),
+        ("trino", '"C".id'),
+        ("athena", '"C".id'),
+        ("bigquery", "`C`.id"),
+        ("databricks", "`C`.id"),
+        ("tsql", "[C].id"),
+    ],
+)
+def test_qualify_uses_dialect_specific_quoting(dialect: str, expected_substring: str) -> None:
+    out = qualify_unqualified_columns_with_alias("id > 1", "C", read_dialect=dialect, write_dialect=dialect)
+    assert expected_substring in out
+
+
+def test_qualify_handles_window_function() -> None:
+    out = qualify_unqualified_columns_with_alias(
+        "ROW_NUMBER() OVER (PARTITION BY col1 ORDER BY col2) > 5",
+        "C",
+        read_dialect="postgres",
+        write_dialect="postgres",
+    )
+    assert '"C".col1' in out
+    assert '"C".col2' in out