Fix: ensure we keep identifiers in unique keys (#1537)

Author: tobymao

sqlmesh/core/engine_adapter/base.py

@@ -1130,10 +1130,7 @@ def merge(
             source=False,
             then=exp.Update(
                 expressions=[
-                    exp.EQ(
-                        this=exp.column(col, MERGE_TARGET_ALIAS),
-                        expression=exp.column(col, MERGE_SOURCE_ALIAS),
-                    )
+                    exp.column(col, MERGE_TARGET_ALIAS).eq(exp.column(col, MERGE_SOURCE_ALIAS))
                     for col in columns_to_types
                 ],
             ),

sqlmesh/core/model/common.py

@@ -63,9 +63,7 @@ def parse_expressions(cls: t.Type, v: t.Any, values: t.Dict[str, t.Any]) -> t.Li
     results = []
 
     for expr in expressions:
-        expr = normalize_identifiers(
-            exp.to_column(expr.name) if isinstance(expr, exp.Identifier) else expr
-        )
+        expr = normalize_identifiers(exp.column(expr) if isinstance(expr, exp.Identifier) else expr)
         expr.meta["dialect"] = dialect
         results.append(expr)
 

sqlmesh/core/model/definition.py

@@ -245,7 +245,7 @@ def render_query(
         """
         return exp.select(
             *(
-                exp.cast(exp.Null(), column_type, copy=False).as_(name, copy=False)
+                exp.cast(exp.Null(), column_type, copy=False).as_(name, copy=False, quoted=True)
                 for name, column_type in (self.columns_to_types or {}).items()
             ),
             copy=False,

tests/core/engine_adapter/test_base.py

@@ -722,23 +722,37 @@ def table_columns(table_name: str) -> t.Dict[str, exp.DataType]:
     adapter.cursor.execute.assert_has_calls([call(x) for x in expected])
 
 
-def test_merge_upsert(make_mocked_engine_adapter: t.Callable):
+def test_merge_upsert(make_mocked_engine_adapter: t.Callable, assert_exp_eq):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
     adapter.merge(
         target_table="target",
-        source_table=t.cast(exp.Select, parse_one("SELECT id, ts, val FROM source")),
+        source_table=t.cast(exp.Select, parse_one('SELECT "ID", ts, val FROM source')),
         columns_to_types={
-            "id": exp.DataType.Type.INT,
+            "ID": exp.DataType.Type.INT,
             "ts": exp.DataType.Type.TIMESTAMP,
             "val": exp.DataType.Type.INT,
         },
-        unique_key=[exp.to_identifier("id")],
+        unique_key=[exp.to_identifier("ID", quoted=True)],
     )
-    adapter.cursor.execute.assert_called_once_with(
-        'MERGE INTO "target" AS "__MERGE_TARGET__" USING (SELECT "id", "ts", "val" FROM "source") AS "__MERGE_SOURCE__" ON "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id" '
-        'WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id", "__MERGE_TARGET__"."ts" = "__MERGE_SOURCE__"."ts", "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val" '
-        'WHEN NOT MATCHED THEN INSERT ("id", "ts", "val") VALUES ("__MERGE_SOURCE__"."id", "__MERGE_SOURCE__"."ts", "__MERGE_SOURCE__"."val")'
+
+    assert_exp_eq(
+        adapter.cursor.execute.call_args[0][0],
+        """
+MERGE INTO "target" AS "__MERGE_TARGET__" USING (
+  SELECT
+    "ID",
+    "ts",
+    "val"
+  FROM "source"
+) AS "__MERGE_SOURCE__"
+  ON "__MERGE_TARGET__"."ID" = "__MERGE_SOURCE__"."ID"
+  WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."ID" = "__MERGE_SOURCE__"."ID",
+    "__MERGE_TARGET__"."ts" = "__MERGE_SOURCE__"."ts",
+    "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val"
+  WHEN NOT MATCHED THEN INSERT ("ID", "ts", "val")
+    VALUES ("__MERGE_SOURCE__"."ID", "__MERGE_SOURCE__"."ts", "__MERGE_SOURCE__"."val")
+""",
     )
 
     adapter.cursor.reset_mock()

tests/core/test_model.py

@@ -1248,16 +1248,30 @@ def test_parse(assert_exp_eq):
 CONST = "bar"
 
 
-def test_python_model_deps() -> None:
-    @model(name="my_model", kind="full", columns={"foo": "int"})
+def test_python_model(assert_exp_eq) -> None:
+    @model(name="my_model", kind="full", columns={'"COL"': "int"})
     def my_model(context, **kwargs):
         context.table("foo")
         context.table(model_name=CONST + ".baz")
 
-    assert model.get_registry()["my_model"].model(
+    m = model.get_registry()["my_model"].model(
         module_path=Path("."),
         path=Path("."),
-    ).depends_on == {"foo", "bar.baz"}
+    )
+
+    assert m.depends_on == {"foo", "bar.baz"}
+    assert m.columns_to_types == {"COL": exp.DataType.build("int")}
+    assert_exp_eq(
+        m.ctas_query(),
+        """
+SELECT
+  CAST(NULL AS INT) AS "COL"
+FROM (VALUES
+  (1)) AS t(dummy)
+WHERE
+  FALSE
+	""",
+    )
 
 
 def test_python_models_returning_sql(assert_exp_eq) -> None:
@@ -2163,11 +2177,11 @@ def test_scd_type_2_defaults():
         MODEL (
             name db.table,
             kind SCD_TYPE_2 (
-                unique_key id,
+                unique_key "ID",
             ),
         );
         SELECT
-            1 as id,
+            1 as "ID",
             '2020-01-01' as ds,
             '2020-01-01' as test_updated_at,
             '2020-01-01' as test_valid_from,
@@ -2176,7 +2190,16 @@ def test_scd_type_2_defaults():
         """
     )
     scd_type_2_model = load_sql_based_model(view_model_expressions)
-    assert scd_type_2_model.unique_key == [exp.to_column("id")]
+    assert scd_type_2_model.unique_key == [exp.to_column("ID", quoted=True)]
+    assert scd_type_2_model.columns_to_types == {
+        "ID": exp.DataType.build("int"),
+        "ds": exp.DataType.build("varchar"),
+        "test_updated_at": exp.DataType.build("varchar"),
+        "test_valid_from": exp.DataType.build("varchar"),
+        "test_valid_to": exp.DataType.build("varchar"),
+        "valid_from": exp.DataType.build("TIMESTAMP"),
+        "valid_to": exp.DataType.build("TIMESTAMP"),
+    }
     assert scd_type_2_model.managed_columns == {
         "valid_from": exp.DataType.build("TIMESTAMP"),
         "valid_to": exp.DataType.build("TIMESTAMP"),