Feat: enable dbt transpilation (#2806)

Co-authored-by: Iaroslav Zeigerman <zeigerman.ia@gmail.com>

Author: treysp

sqlmesh/dbt/adapter.py

@@ -29,12 +29,12 @@ def __init__(
         self,
         jinja_macros: JinjaMacroRegistry,
         jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
-        dialect: t.Optional[str] = None,
+        project_dialect: t.Optional[str] = None,
     ):
         self.jinja_macros = jinja_macros
         self.jinja_globals = jinja_globals.copy() if jinja_globals else {}
         self.jinja_globals["adapter"] = self
-        self.dialect = dialect
+        self.project_dialect = project_dialect
 
     @abc.abstractmethod
     def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
@@ -93,7 +93,7 @@ def resolve_identifier(self, relation: BaseRelation) -> t.Optional[str]:
 
     def quote(self, identifier: str) -> str:
         """Returns a quoted identifier."""
-        return exp.to_column(identifier).sql(dialect=self.dialect, identify=True)
+        return exp.to_column(identifier).sql(dialect=self.project_dialect, identify=True)
 
     def dispatch(self, name: str, package: t.Optional[str] = None) -> t.Callable:
         """Returns a dialect-specific version of a macro with the given name."""
@@ -189,12 +189,17 @@ def __init__(
         snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
         table_mapping: t.Optional[t.Dict[str, str]] = None,
         deployability_index: t.Optional[DeployabilityIndex] = None,
+        project_dialect: t.Optional[str] = None,
     ):
         from dbt.adapters.base import BaseRelation
         from dbt.adapters.base.column import Column
         from dbt.adapters.base.relation import Policy
 
-        super().__init__(jinja_macros, jinja_globals=jinja_globals, dialect=engine_adapter.dialect)
+        super().__init__(
+            jinja_macros,
+            jinja_globals=jinja_globals,
+            project_dialect=project_dialect or engine_adapter.dialect,
+        )
 
         table_mapping = table_mapping or {}
 
@@ -259,7 +264,9 @@ def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
 
         mapped_table = self._map_table_name(self._normalize(self._relation_to_table(relation)))
         return [
-            Column.from_description(name=name, raw_data_type=dtype.sql(dialect=self.dialect))
+            Column.from_description(
+                name=name, raw_data_type=dtype.sql(dialect=self.project_dialect)
+            )
             for name, dtype in self.engine_adapter.columns(table_name=mapped_table).items()
         ]
 
@@ -298,12 +305,12 @@ def execute(
             self.engine_adapter.fetchdf if fetch else self.engine_adapter.execute  # type: ignore
         )
 
-        expression = parse_one(sql, read=self.dialect)
+        expression = parse_one(sql, read=self.project_dialect)
         with normalize_and_quote(
-            expression, t.cast(str, self.dialect), self.engine_adapter.default_catalog
+            expression, t.cast(str, self.project_dialect), self.engine_adapter.default_catalog
         ) as expression:
             expression = exp.replace_tables(
-                expression, self.table_mapping, dialect=self.dialect, copy=False
+                expression, self.table_mapping, dialect=self.project_dialect, copy=False
             )
 
         if auto_begin:
@@ -328,17 +335,17 @@ def resolve_identifier(self, relation: BaseRelation) -> t.Optional[str]:
         return identifier if identifier else None
 
     def _map_table_name(self, table: exp.Table) -> exp.Table:
-        name = table.sql()
+        name = table.sql(dialect=self.project_dialect)
         physical_table_name = self.table_mapping.get(name)
         if not physical_table_name:
             return table
 
         logger.debug("Resolved ref '%s' to snapshot table '%s'", name, physical_table_name)
 
-        return exp.to_table(physical_table_name, dialect=self.dialect)
+        return exp.to_table(physical_table_name, dialect=self.project_dialect)
 
     def _relation_to_table(self, relation: BaseRelation) -> exp.Table:
-        return exp.to_table(relation.render(), dialect=self.dialect)
+        return exp.to_table(relation.render(), dialect=self.project_dialect)
 
     def _table_to_relation(self, table: exp.Table) -> BaseRelation:
         return self.relation_type.create(
@@ -358,7 +365,7 @@ def _schema(self, schema_relation: BaseRelation) -> exp.Table:
 
     def _normalize(self, input_table: exp.Table) -> exp.Table:
         normalized_name = normalize_model_name(
-            input_table, self.engine_adapter.default_catalog, self.dialect
+            input_table, self.engine_adapter.default_catalog, self.project_dialect
         )
         normalized_table = exp.to_table(normalized_name)
         if not input_table.this:

sqlmesh/dbt/basemodel.py

@@ -105,6 +105,7 @@ class BaseModelConfig(GeneralConfig):
     path: Path = Path()
     dependencies: Dependencies = Dependencies()
     tests: t.List[TestConfig] = []
+    dialect_: t.Optional[str] = Field(None, alias="dialect")
 
     # DBT configuration fields
     name: str = ""
@@ -187,6 +188,9 @@ def config_name(self) -> str:
         """
         return f"{self.package_name}.{self.name}"
 
+    def dialect(self, context: DbtContext) -> str:
+        return self.dialect_ or context.default_dialect
+
     def canonical_name(self, context: DbtContext) -> str:
         """
         Get the sqlmesh model name
@@ -296,7 +300,7 @@ def sqlmesh_model_kwargs(self, context: DbtContext) -> t.Dict[str, t.Any]:
         )
         return {
             "audits": [(test.name, {}) for test in self.tests],
-            "columns": column_types_to_sqlmesh(self.columns, context.dialect) or None,
+            "columns": column_types_to_sqlmesh(self.columns, self.dialect(context)) or None,
             "column_descriptions": column_descriptions_to_sqlmesh(self.columns) or None,
             "depends_on": {
                 model.canonical_name(context) for model in model_context.refs.values()

sqlmesh/dbt/builtin.py

@@ -43,9 +43,9 @@ def warn(self, msg: str) -> str:
 
 
 class Api:
-    def __init__(self, target: t.Optional[AttributeDict] = None) -> None:
-        if target:
-            config_class = TARGET_TYPE_TO_CONFIG_CLASS[target["type"]]
+    def __init__(self, dialect: t.Optional[str]) -> None:
+        if dialect:
+            config_class = TARGET_TYPE_TO_CONFIG_CLASS[dialect]
             self.Relation = config_class.relation_class
             self.Column = config_class.column_class
             self.quote_policy = config_class.quote_policy
@@ -301,8 +301,10 @@ def create_builtin_globals(
     jinja_globals = jinja_globals.copy()
 
     target: t.Optional[AttributeDict] = jinja_globals.get("target", None)
-    api = Api(target)
-    dialect = target.dialect if target else None  # type: ignore
+    project_dialect = jinja_globals.pop("dialect", None) or (
+        target.get("dialect") if target else None
+    )
+    api = Api(project_dialect)
 
     builtin_globals["api"] = api
 
@@ -349,13 +351,14 @@ def create_builtin_globals(
             snapshots=jinja_globals.get("snapshots", {}),
             table_mapping=jinja_globals.get("table_mapping", {}),
             deployability_index=jinja_globals.get("deployability_index"),
+            project_dialect=project_dialect,
         )
     else:
         builtin_globals["flags"] = Flags(which="parse")
         adapter = ParsetimeAdapter(
             jinja_macros,
             jinja_globals={**builtin_globals, **jinja_globals},
-            dialect=dialect,
+            project_dialect=project_dialect,
         )
 
     sql_execution = SQLExecution(adapter)

sqlmesh/dbt/context.py

@@ -59,9 +59,13 @@ class DbtContext:
     _manifest: t.Optional[ManifestHelper] = None
 
     @property
-    def dialect(self) -> str:
+    def default_dialect(self) -> str:
+        if self.sqlmesh_config.dialect:
+            return self.sqlmesh_config.dialect
         if not self.target:
-            raise SQLMeshError("Target must be configured before calling the dialect property.")
+            raise SQLMeshError(
+                "Target must be configured before calling the default_dialect property."
+            )
         return self.target.dialect
 
     @property
@@ -229,6 +233,9 @@ def jinja_globals(self) -> t.Dict[str, JinjaGlobalAttribute]:
             output["project_name"] = self.project_name
         if self._target is not None:
             output["target"] = self._target.attribute_dict()
+        # pass user-specified default dialect if we have already loaded the config
+        if self.sqlmesh_config.dialect:
+            output["dialect"] = self.sqlmesh_config.dialect
         return output
 
     def context_for_dependencies(self, dependencies: Dependencies) -> DbtContext:

sqlmesh/dbt/loader.py

@@ -3,7 +3,6 @@
 import logging
 import typing as t
 from pathlib import Path
-
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import Audit
 from sqlmesh.core.config import (
@@ -43,7 +42,8 @@ def sqlmesh_config(
     context = DbtContext(project_root=project_root)
     profile = Profile.load(context, target_name=dbt_target_name)
     model_defaults = kwargs.pop("model_defaults", ModelDefaultsConfig())
-    model_defaults.dialect = profile.target.dialect
+    if model_defaults.dialect is None:
+        model_defaults.dialect = profile.target.dialect
 
     target_to_sqlmesh_args = {}
     if register_comments is not None:

sqlmesh/dbt/model.py

@@ -201,7 +201,9 @@ def model_kind(self, context: DbtContext) -> ModelKind:
         if materialization == Materialization.VIEW:
             return ViewKind()
         if materialization == Materialization.INCREMENTAL:
-            incremental_materialization_kwargs: t.Dict[str, t.Any] = {"dialect": context.dialect}
+            incremental_materialization_kwargs: t.Dict[str, t.Any] = {
+                "dialect": self.dialect(context)
+            }
             for field in ("batch_size", "lookback", "forward_only"):
                 field_val = getattr(self, field, None) or self.meta.get(field, None)
                 if field_val:
@@ -277,7 +279,7 @@ def model_kind(self, context: DbtContext) -> ModelKind:
                     f"{self.canonical_name(context)}: SQLMesh snapshot strategy is required for snapshot materialization."
                 )
             shared_kwargs = {
-                "dialect": context.dialect,
+                "dialect": self.dialect(context),
                 "unique_key": self.unique_key,
                 "invalidate_hard_deletes": self.invalidate_hard_deletes,
                 "valid_from_name": "dbt_valid_from",
@@ -358,14 +360,14 @@ def sqlmesh_config_fields(self) -> t.Set[str]:
 
     def to_sqlmesh(self, context: DbtContext) -> Model:
         """Converts the dbt model into a SQLMesh model."""
-        dialect = context.dialect
+        model_dialect = self.dialect(context)
         query = d.jinja_query(self.sql_no_config)
 
         optional_kwargs: t.Dict[str, t.Any] = {}
 
         if self.partition_by:
             optional_kwargs["partitioned_by"] = (
-                [exp.to_column(val) for val in self.partition_by]
+                [exp.to_column(val, dialect=model_dialect) for val in self.partition_by]
                 if isinstance(self.partition_by, list)
                 else self._big_query_partition_by_expr(context)
             )
@@ -374,7 +376,7 @@ def to_sqlmesh(self, context: DbtContext) -> Model:
             clustered_by = []
             for c in self.cluster_by:
                 try:
-                    clustered_by.append(d.parse_one(c, dialect=dialect).name)
+                    clustered_by.append(d.parse_one(c, dialect=model_dialect).name)
                 except SqlglotError as e:
                     raise ConfigError(f"Failed to parse cluster_by field '{c}': {e}") from e
             optional_kwargs["clustered_by"] = clustered_by
@@ -403,7 +405,7 @@ def to_sqlmesh(self, context: DbtContext) -> Model:
         return create_sql_model(
             self.canonical_name(context),
             query,
-            dialect=dialect,
+            dialect=model_dialect,
             kind=self.model_kind(context),
             start=self.start,
             **optional_kwargs,

sqlmesh/dbt/seed.py

@@ -50,7 +50,7 @@ def to_sqlmesh(self, context: DbtContext) -> Model:
         return create_seed_model(
             self.canonical_name(context),
             SeedKind(path=seed_path),
-            dialect=context.dialect,
+            dialect=self.dialect(context),
             **kwargs,
         )
 

sqlmesh/dbt/target.py

@@ -852,5 +852,6 @@ def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
     "snowflake": SnowflakeConfig,
     "bigquery": BigQueryConfig,
     "sqlserver": MSSQLConfig,
+    "tsql": MSSQLConfig,
     "trino": TrinoConfig,
 }

sqlmesh/dbt/test.py

@@ -6,7 +6,6 @@
 from pathlib import Path
 
 from pydantic import Field
-
 import sqlmesh.core.dialect as d
 from sqlmesh.core.audit import Audit, ModelAudit, StandaloneAudit
 from sqlmesh.dbt.common import (
@@ -48,6 +47,7 @@ class TestConfig(GeneralConfig):
         interval_unit: The duration of an interval for the audit. By default, it is computed from the cron expression.
         column_name: The name of the column under test.
         dependencies: The macros, refs, and sources the test depends upon.
+        dialect: SQL dialect of the test query.
         package_name: Name of the package that defines the test.
         alias: The alias for the materialized table where failures are stored (Not supported).
         schema: The schema for the materialized table where the failures are stored (Not supported).
@@ -73,6 +73,7 @@ class TestConfig(GeneralConfig):
     interval_unit: t.Optional[str] = None
     column_name: t.Optional[str] = None
     dependencies: Dependencies = Dependencies()
+    dialect_: t.Optional[str] = Field(None, alias="dialect")
 
     # dbt fields
     package_name: str = ""
@@ -110,6 +111,9 @@ def is_standalone(self) -> bool:
     def sqlmesh_config_fields(self) -> t.Set[str]:
         return {"description", "owner", "stamp", "cron", "interval_unit"}
 
+    def dialect(self, context: DbtContext) -> str:
+        return self.dialect_ or context.default_dialect
+
     def to_sqlmesh(self, context: DbtContext) -> Audit:
         """Convert dbt Test to SQLMesh Audit
 
@@ -142,7 +146,7 @@ def to_sqlmesh(self, context: DbtContext) -> Audit:
             jinja_macros.add_globals({"this": self.relation_info})
             audit = StandaloneAudit(
                 name=self.name,
-                dialect=context.dialect,
+                dialect=self.dialect(context),
                 skip=skip,
                 query=query,
                 jinja_macros=jinja_macros,
@@ -158,7 +162,7 @@ def to_sqlmesh(self, context: DbtContext) -> Audit:
         else:
             audit = ModelAudit(
                 name=self.name,
-                dialect=context.dialect,
+                dialect=self.dialect(context),
                 skip=skip,
                 blocking=blocking,
                 query=query,

tests/dbt/test_adapter.py

@@ -73,10 +73,23 @@ def test_normalization(
     context = sushi_test_project.context
     assert context.target
 
+    # bla and bob will be normalized to lowercase since the target is duckdb
     adapter_mock = mocker.MagicMock()
-    adapter_mock.dialect = "snowflake"
     adapter_mock.default_catalog = "test"
+    adapter_mock.dialect = "duckdb"
+
+    duckdb_renderer = runtime_renderer(context, engine_adapter=adapter_mock)
+
+    schema_bla = schema_("bla", "test", quoted=True)
+    relation_bla_bob = exp.table_("bob", db="bla", catalog="test", quoted=True)
 
+    duckdb_renderer("{{ adapter.get_relation(database=None, schema='bla', identifier='bob') }}")
+    adapter_mock.table_exists.assert_has_calls([call(relation_bla_bob)])
+
+    # bla and bob will be normalized to uppercase since the target is Snowflake, even though the default dialect is duckdb
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.default_catalog = "test"
+    adapter_mock.dialect = "snowflake"
     context.target = SnowflakeConfig(
         account="test",
         user="test",
@@ -85,16 +98,19 @@ def test_normalization(
         database="test",
         schema="test",
     )
-
     renderer = runtime_renderer(context, engine_adapter=adapter_mock)
 
-    # bla and bob will be normalized to uppercase since we're dealing with Snowflake
-    schema_bla = schema_("BLA", "TEST", quoted=True)
-    relation_bla_bob = exp.table_("BOB", db="BLA", catalog="TEST", quoted=True)
+    schema_bla = schema_("bla", "test", quoted=True)
+    relation_bla_bob = exp.table_("bob", db="bla", catalog="test", quoted=True)
 
     renderer("{{ adapter.get_relation(database=None, schema='bla', identifier='bob') }}")
     adapter_mock.table_exists.assert_has_calls([call(relation_bla_bob)])
 
+    renderer("{{ adapter.get_relation(database='custom_db', schema='bla', identifier='bob') }}")
+    adapter_mock.table_exists.assert_has_calls(
+        [call(exp.table_("bob", db="bla", catalog="custom_db", quoted=True))]
+    )
+
     renderer(
         "{%- set relation = api.Relation.create(schema='bla') -%}"
         "{{ adapter.create_schema(relation) }}"
@@ -114,7 +130,7 @@ def test_normalization(
     adapter_mock.drop_table.assert_has_calls([call(relation_bla_bob)])
 
     expected_star_query: exp.Select = exp.maybe_parse(
-        'SELECT * FROM "T" as "T"', dialect="snowflake"
+        'SELECT * FROM "t" as "t"', dialect="snowflake"
     )
 
     # The following call to run_query won't return dataframes and so we're expected to