Repository: TobikoData/sqlmesh
Pull request #2482 (2024-04-22 12:45:54)

Fix: use MacroEvaluator dialect in certain builtin macros
Context: https://tobiko-data.slack.com/archives/C044BRE5W4S/p1713786694499789


Patch:
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 19f1b8aa..c0ba8da4 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -680,11 +680,14 @@ def star(
 
     exclude = {e.name for e in except_.expressions}
     quoted = quote_identifiers.this
+    table_identifier = alias.name or relation.name
 
     return [
-        exp.cast(exp.column(column, table=alias.name, quoted=quoted), type_).as_(
-            f"{prefix.this}{column}{suffix.this}", quoted=quoted
-        )
+        exp.cast(
+            exp.column(column, table=table_identifier, quoted=quoted),
+            type_,
+            dialect=evaluator.dialect,
+        ).as_(f"{prefix.this}{column}{suffix.this}", quoted=quoted)
         for column, type_ in evaluator.columns_to_types(relation).items()
         if column not in exclude
     ]
@@ -794,7 +797,7 @@ def union(
     }
 
     projections = [
-        exp.cast(column, type_).as_(column)
+        exp.cast(column, type_, dialect=evaluator.dialect).as_(column)
         for column, type_ in evaluator.columns_to_types(tables[0]).items()
         if column in columns
     ]


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index 94dec960..b70e5905 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -70,8 +70,8 @@ class ModelTest(unittest.TestCase):
 
         self._validate_and_normalize_test()
 
-        # This ID is appended to each input fixture name to avoid concurrency issues
-        self._test_id = random_id(short=True)
+        # The test schema name is randomized to avoid concurrency issues
+        self._test_schema = f"sqlmesh_test_{random_id(short=True)}"
 
         self._engine_adapter_dialect = Dialect.get_or_raise(self.engine_adapter.dialect)
         self._transforms = self._engine_adapter_dialect.generator_class.TRANSFORMS
@@ -131,10 +131,9 @@ class ModelTest(unittest.TestCase):
                         known_columns_to_types[col] = v_type
 
             test_fixture_table = self._test_fixture_table(name)
-            if test_fixture_table.db:
-                self.engine_adapter.create_schema(
-                    schema_(test_fixture_table.args["db"], test_fixture_table.args.get("catalog"))
-                )
+            self.engine_adapter.create_schema(
+                schema_(test_fixture_table.args["db"], test_fixture_table.args.get("catalog"))
+            )
 
             df = _create_df(rows, columns=known_columns_to_types)
             self.engine_adapter.create_view(test_fixture_table, df, known_columns_to_types)
@@ -142,8 +141,7 @@ class ModelTest(unittest.TestCase):
     def tearDown(self) -> None:
         """Drop all fixture tables."""
         if not self.preserve_fixtures:
-            for name in self.body.get("inputs", {}):
-                self.engine_adapter.drop_view(self._fixture_table_cache[name])
+            self.engine_adapter.drop_schema(self._test_schema, cascade=True)
 
     def assert_equal(
         self,
@@ -366,7 +364,12 @@ class ModelTest(unittest.TestCase):
         table = self._fixture_table_cache.get(name)
         if not table:
             table = exp.to_table(name, dialect=self.dialect)
-            table.this.set("this", f"{table.this.this}__fixture__{self._test_id}")
+
+            # We change both the schema and the catalog, so we need to ensure there are no name clashes
+            table.this.set("this", "__".join(part.name for part in table.parts))
+            table.set("db", exp.to_identifier(self._test_schema))
+            table.set("catalog", exp.to_identifier(self.default_catalog))
+
             self._fixture_table_cache[name] = table
 
         return table
diff --git a/tests/conftest.py b/tests/conftest.py
index 2f43d758..4eb5aa9c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -15,6 +15,7 @@ import pandas as pd
 import pytest
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp, maybe_parse, parse_one
+from sqlglot.dialects.dialect import DialectType
 from sqlglot.helper import ensure_list
 
 from sqlmesh.core.config import DuckDBConnectionConfig
@@ -281,9 +282,11 @@ def init_and_plan_context(copy_to_temp_path, mocker) -> t.Callable:
 
 @pytest.fixture
 def assert_exp_eq() -> t.Callable:
-    def _assert_exp_eq(source: exp.Expression | str, expected: exp.Expression | str) -> None:
-        source_exp = maybe_parse(source)
-        expected_exp = maybe_parse(expected)
+    def _assert_exp_eq(
+        source: exp.Expression | str, expected: exp.Expression | str, dialect: DialectType = None
+    ) -> None:
+        source_exp = maybe_parse(source, dialect=dialect)
+        expected_exp = maybe_parse(expected, dialect=dialect)
 
         if not source_exp:
             raise ValueError(f"Could not parse {source}")
@@ -291,7 +294,9 @@ def assert_exp_eq() -> t.Callable:
             raise ValueError(f"Could not parse {expected}")
 
         if source_exp != expected_exp:
-            assert source_exp.sql(pretty=True) == expected_exp.sql(pretty=True)
+            assert source_exp.sql(dialect=dialect, pretty=True) == expected_exp.sql(
+                dialect=dialect, pretty=True
+            )
             assert source_exp == expected_exp
 
     return _assert_exp_eq
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 588854c7..32b7a31f 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -3,7 +3,7 @@ import typing as t
 from unittest import mock
 
 import pytest
-from sqlglot import exp, parse_one
+from sqlglot import MappingSchema, exp, parse_one
 
 from sqlmesh.core.dialect import StagedFilePath
 from sqlmesh.core.macros import MacroEvaluator, macro
@@ -33,7 +33,15 @@ def macro_evaluator() -> MacroEvaluator:
     )
 
 
-def test_case(macro_evaluator):
+def test_star(assert_exp_eq) -> None:
+    sql = """SELECT @STAR(foo) FROM foo"""
+    expected_sql = """SELECT CAST([foo].[a] AS DATETIMEOFFSET) AS [a] FROM foo"""
+    schema = MappingSchema({"foo": {"a": "datetimeoffset"}}, dialect="tsql")
+    evaluator = MacroEvaluator(schema=schema, dialect="tsql")
+    assert_exp_eq(evaluator.transform(parse_one(sql, read="tsql")), expected_sql, dialect="tsql")
+
+
+def test_case(macro_evaluator: MacroEvaluator) -> None:
     assert macro.get_registry()["upper"]
 
 
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 76197eb6..c94a7440 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import datetime
 import typing as t
 from pathlib import Path
+from unittest.mock import call
 
 import pandas as pd
 import pytest
@@ -104,9 +105,10 @@ def full_model_with_two_ctes(request) -> SqlModel:
 
 
 def test_ctes(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
-    test = _create_test(
-        body=load_yaml(
-            """
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
 test_foo:
   model: sushi.foo
   inputs:
@@ -123,21 +125,14 @@ test_foo:
   vars:
     start: 2022-01-01
     end: 2022-01-01
-            """
-        ),
-        test_name="test_foo",
-        model=sushi_context.upsert_model(full_model_with_two_ctes),
-        context=sushi_context,
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_with_two_ctes),
+            context=sushi_context,
+        ).run()
     )
 
-    random_id = "jzngz56a"
-    test._test_id = random_id
-    _check_successful_or_raise(test.run())
-
-    assert len(test._fixture_table_cache) == len(sushi_context.models) + 1
-    for table in test._fixture_table_cache.values():
-        assert table.name.endswith(f"__fixture__{random_id}")
-
 
 def test_ctes_only(sushi_context: Context, full_model_with_two_ctes: SqlModel) -> None:
     _check_successful_or_raise(
@@ -581,15 +576,13 @@ test_child:
         """
     )
     test = _create_test(body, "test_child", child, sushi_context)
-
-    random_id = "jzngz56a"
-    test._test_id = random_id
+    test._test_schema = "sqlmesh_test_jzngz56a"
 
     spy_execute = mocker.spy(test.engine_adapter, "_execute")
     _check_successful_or_raise(test.run())
 
     spy_execute.assert_any_call(
-        f'CREATE OR REPLACE VIEW "memory"."sushi"."parent__fixture__{random_id}" ("s", "a", "b") AS '
+        f'CREATE OR REPLACE VIEW "memory"."sqlmesh_test_jzngz56a"."memory__sushi__parent" ("s", "a", "b") AS '
         "SELECT "
         'CAST("s" AS STRUCT("d" DATE)) AS "s", '
         'CAST("a" AS INT) AS "a", '
@@ -922,15 +915,21 @@ test_foo:
         ),
         context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
     )
+    test._test_schema = f"sqlmesh_test_jzngz56a"
 
     spy_execute = mocker.spy(test.engine_adapter, "_execute")
     _check_successful_or_raise(test.run())
 
-    spy_execute.assert_called_with(
-        "SELECT "
-        """CAST('2023-01-01 12:05:03+00:00' AS DATE) AS "cur_date", """
-        """CAST('2023-01-01 12:05:03+00:00' AS TIME) AS "cur_time", """
-        '''CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMP) AS "cur_timestamp"''',
+    spy_execute.assert_has_calls(
+        [
+            call(
+                "SELECT "
+                """CAST('2023-01-01 12:05:03+00:00' AS DATE) AS "cur_date", """
+                """CAST('2023-01-01 12:05:03+00:00' AS TIME) AS "cur_time", """
+                '''CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMP) AS "cur_timestamp"'''
+            ),
+            call('DROP SCHEMA IF EXISTS "sqlmesh_test_jzngz56a" CASCADE'),
+        ]
     )
 
     @model("py_model", columns={"ts1": "timestamptz", "ts2": "timestamptz"})
@@ -971,6 +970,33 @@ def test_successes(sushi_context: Context) -> None:
     assert "test_customer_revenue_by_day" in successful_tests
 
 
+def test_create_external_model_fixture(sushi_context: Context) -> None:
+    test = _create_test(
+        body=load_yaml(
+            """
+test_foo:
+  model: sushi.foo
+  inputs:
+    c.db.external:
+      - x: 1
+  outputs:
+    query:
+      - x: 1
+            """
+        ),
+        test_name="test_foo",
+        model=_create_model("SELECT x FROM c.db.external"),
+        context=sushi_context,
+    )
+    test._test_schema = f"sqlmesh_test_jzngz56a"
+    _check_successful_or_raise(test.run())
+
+    assert len(test._fixture_table_cache) == len(sushi_context.models) + 1
+    for table in test._fixture_table_cache.values():
+        assert table.catalog == "memory"
+        assert table.db == f"sqlmesh_test_jzngz56a"
+
+
 def test_test_generation(tmp_path: Path) -> None:
     init_example_project(tmp_path, dialect="duckdb")