Repository: TobikoData/sqlmesh
Pull request #2975 (2024-07-31 22:28:28)

Fix!: Improve handling of the `--where` option in table_diff
Prior to this change, running `sqlmesh table_diff` with something like `--where "foo = 'bar'"` would throw an error like:

```
duckdb.duckdb.BinderException: Binder Error: Ambiguous reference to column name "foo" (use: "t.foo" or "s.foo")
```

This was because the where clause was injected verbatim without automatically qualifying the columns. The user had to know that table_diff used a subselect internally and supply `--where "s.foo = 'bar' and t.foo = 'bar'"` in order to get this to work.

This PR refactors the generated table_diff query to use CTE's to build the source / target / stats datasets so the `--where` clause can still be injected verbatim but it's now scoped to each dataset so there are no more ambiguous column references


Patch:
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 2a686fbb..c9e63837 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -272,7 +272,22 @@ class TableDiff:
             def name(e: exp.Expression) -> str:
                 return e.args["alias"].sql(identify=True)
 
-            query = (
+            source_query = (
+                exp.select(*(exp.column(c) for c in source_schema))
+                .from_(self.source)
+                .where(self.where)
+            )
+            target_query = (
+                exp.select(*(exp.column(c) for c in target_schema))
+                .from_(self.target)
+                .where(self.where)
+            )
+
+            source_table = exp.table_("__source")
+            target_table = exp.table_("__target")
+            stats_table = exp.table_("__stats")
+
+            stats_query = (
                 exp.select(
                     *s_selects.values(),
                     *t_selects.values(),
@@ -313,31 +328,32 @@ class TableDiff:
                     ).as_("null_grain"),
                     *comparisons,
                 )
-                .from_(exp.alias_(self.source, "s"))
-                .join(
-                    self.target,
-                    on=self.on,
-                    join_type="FULL",
-                    join_alias="t",
-                )
-                .where(self.where)
+                .from_(source_table.as_("s"))
+                .join(target_table.as_("t"), on=self.on, join_type="FULL")
             )
 
-            query = exp.select(
-                "*",
-                exp.Case()
-                .when(
-                    exp.and_(
-                        *[
-                            exp.column(f"{c}_matches").eq(exp.Literal.number(1))
-                            for c in matched_columns
-                        ]
-                    ),
-                    exp.Literal.number(1),
+            query = (
+                exp.Select()
+                .with_(source_table, source_query)
+                .with_(target_table, target_query)
+                .with_(stats_table, stats_query)
+                .select(
+                    "*",
+                    exp.Case()
+                    .when(
+                        exp.and_(
+                            *[
+                                exp.column(f"{c}_matches").eq(exp.Literal.number(1))
+                                for c in matched_columns
+                            ]
+                        ),
+                        exp.Literal.number(1),
+                    )
+                    .else_(exp.Literal.number(0))
+                    .as_("row_full_match"),
                 )
-                .else_(exp.Literal.number(0))
-                .as_("row_full_match"),
-            ).from_(query.subquery("stats"))
+                .from_(stats_table)
+            )
 
             query = quote_identifiers(query, dialect=self.model_dialect or self.dialect)
             temp_table = exp.table_("diff", db="sqlmesh_temp", quoted=True)
diff --git a/sqlmesh/dbt/basemodel.py b/sqlmesh/dbt/basemodel.py
index b88ca295..a95b25b9 100644
--- a/sqlmesh/dbt/basemodel.py
+++ b/sqlmesh/dbt/basemodel.py
@@ -26,7 +26,6 @@ from sqlmesh.dbt.common import (
 from sqlmesh.dbt.relation import Policy, RelationType
 from sqlmesh.dbt.test import TestConfig
 from sqlmesh.utils import AttributeDict
-from sqlmesh.utils.conversions import ensure_bool
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import field_validator
 
@@ -144,11 +143,6 @@ class BaseModelConfig(GeneralConfig):
 
         return hooks
 
-    @field_validator("full_refresh", mode="before")
-    @classmethod
-    def _validate_bool(cls, v: str) -> bool:
-        return ensure_bool(v)
-
     @field_validator("grants", mode="before")
     @classmethod
     def _validate_grants(cls, v: t.Dict[str, str]) -> t.Dict[str, t.List[str]]:


Test patch:
diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py
index fa4175be..383f16a6 100644
--- a/tests/core/test_table_diff.py
+++ b/tests/core/test_table_diff.py
@@ -246,7 +246,7 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
         ),
     )
 
-    query_sql = 'CREATE TABLE IF NOT EXISTS "sqlmesh_temp"."__temp_diff_abcdefgh" AS SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "t"."key" AS "t__key", "t"."value" AS "t__value", CASE WHEN NOT "s"."key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."key" = "t"."key" AND NOT "s"."key" IS NULL AND NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "table_diff_source" AS "s" FULL JOIN "table_diff_target" AS "t" ON ("s"."key" = "t"."key") OR (("s"."key" IS NULL) AND ("t"."key" IS NULL))) AS "stats"'
+    query_sql = 'CREATE TABLE IF NOT EXISTS "sqlmesh_temp"."__temp_diff_abcdefgh" AS WITH "__source" AS (SELECT "key", "value" FROM "table_diff_source"), "__target" AS (SELECT "key", "value" FROM "table_diff_target"), "__stats" AS (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "t"."key" AS "t__key", "t"."value" AS "t__value", CASE WHEN NOT "s"."key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."key" = "t"."key" AND NOT "s"."key" IS NULL AND NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "__source" AS "s" FULL JOIN "__target" AS "t" ON ("s"."key" = "t"."key") OR (("s"."key" IS NULL) AND ("t"."key" IS NULL))) SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM "__stats"'
     summary_query_sql = 'SELECT SUM("s_exists") AS "s_count", SUM("t_exists") AS "t_count", SUM("row_joined") AS "join_count", SUM("null_grain") AS "null_grain_count", SUM("row_full_match") AS "full_match_count", SUM("key_matches") AS "key_matches", SUM("value_matches") AS "value_matches", COUNT(DISTINCT ("s__key")) AS "distinct_count_s", COUNT(DISTINCT ("t__key")) AS "distinct_count_t" FROM "sqlmesh_temp"."__temp_diff_abcdefgh"'
     sample_query_sql = 'SELECT "s_exists", "t_exists", "row_joined", "row_full_match", "s__key", "s__value", "t__key", "t__value" FROM "sqlmesh_temp"."__temp_diff_abcdefgh" WHERE "key_matches" = 0 OR "value_matches" = 0 ORDER BY "s__key" NULLS FIRST, "t__key" NULLS FIRST LIMIT 20'
 
@@ -263,3 +263,17 @@ def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture)
     spy_execute.assert_any_call(query_sql)
     spy_execute.assert_any_call(summary_query_sql)
     spy_execute.assert_any_call(sample_query_sql)
+
+    spy_execute.reset_mock()
+
+    # Also check WHERE clause is propagated correctly
+    sushi_context_fixed_date.table_diff(
+        source="table_diff_source",
+        target="table_diff_target",
+        on=["key"],
+        skip_columns=["ignored"],
+        where="key = 2",
+    )
+
+    query_sql_where = 'CREATE TABLE IF NOT EXISTS "sqlmesh_temp"."__temp_diff_abcdefgh" AS WITH "__source" AS (SELECT "key", "value" FROM "table_diff_source" WHERE "key" = 2), "__target" AS (SELECT "key", "value" FROM "table_diff_target" WHERE "key" = 2), "__stats" AS (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "t"."key" AS "t__key", "t"."value" AS "t__value", CASE WHEN NOT "s"."key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."key" = "t"."key" AND NOT "s"."key" IS NULL AND NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "__source" AS "s" FULL JOIN "__target" AS "t" ON ("s"."key" = "t"."key") OR (("s"."key" IS NULL) AND ("t"."key" IS NULL))) SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM "__stats"'
+    spy_execute.assert_any_call(query_sql_where)
diff --git a/tests/dbt/test_transformation.py b/tests/dbt/test_transformation.py
index 4dfb6805..6b5b7762 100644
--- a/tests/dbt/test_transformation.py
+++ b/tests/dbt/test_transformation.py
@@ -229,6 +229,10 @@ def test_model_kind():
         materialized=Materialization.INCREMENTAL, incremental_strategy="append"
     ).model_kind(context) == IncrementalUnmanagedKind(disable_restatement=False)
 
+    assert ModelConfig(
+        materialized=Materialization.INCREMENTAL, incremental_strategy="append", full_refresh=None
+    ).model_kind(context) == IncrementalUnmanagedKind(disable_restatement=False)
+
     assert ModelConfig(
         materialized=Materialization.INCREMENTAL,
         incremental_strategy="insert_overwrite",