Repository: TobikoData/sqlmesh
Pull request #2200 (2024-03-01 18:17:15)

Fix: map dbt model default dialect sqlserver -> tsql


Patch:
diff --git a/sqlmesh/core/state_sync/common.py b/sqlmesh/core/state_sync/common.py
index a440561b..4bd8cf35 100644
--- a/sqlmesh/core/state_sync/common.py
+++ b/sqlmesh/core/state_sync/common.py
@@ -240,11 +240,23 @@ class CommonStateSyncMixin(StateSync):
                         [(snapshot, snapshot.get_removal_interval(effective_from_ts, current_ts))]
                     )
 
+                update_required = False
+
                 if snapshot.unpaused_ts:
                     logger.info("Pausing snapshot %s", snapshot.snapshot_id)
                     snapshot.set_unpaused_ts(None)
-                    if not snapshot.is_forward_only and target_snapshot.is_forward_only:
-                        snapshot.unrestorable = True
+                    update_required = True
+
+                if (
+                    not snapshot.is_forward_only
+                    and target_snapshot.is_forward_only
+                    and not snapshot.unrestorable
+                ):
+                    logger.info("Marking snapshot %s as unrestorable", snapshot.snapshot_id)
+                    snapshot.unrestorable = True
+                    update_required = True
+
+                if update_required:
                     self._update_snapshot(snapshot)
 
     def _ensure_no_gaps(
diff --git a/sqlmesh/dbt/adapter.py b/sqlmesh/dbt/adapter.py
index 1e070388..3c846074 100644
--- a/sqlmesh/dbt/adapter.py
+++ b/sqlmesh/dbt/adapter.py
@@ -248,9 +248,7 @@ class RuntimeAdapter(BaseAdapter):
 
         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.engine_adapter.dialect)
-            )
+            Column.from_description(name=name, raw_data_type=dtype.sql(dialect=self.dialect))
             for name, dtype in self.engine_adapter.columns(table_name=mapped_table).items()
         ]
 
@@ -290,9 +288,9 @@ class RuntimeAdapter(BaseAdapter):
             self.engine_adapter.fetchdf if fetch else self.engine_adapter.execute  # type: ignore
         )
 
-        expression = parse_one(sql, read=self.engine_adapter.dialect)
+        expression = parse_one(sql, read=self.dialect)
         with _normalize_and_quote(
-            expression, self.engine_adapter.dialect, self.engine_adapter.default_catalog
+            expression, t.cast(str, self.dialect), self.engine_adapter.default_catalog
         ) as expression:
             expression = exp.replace_tables(
                 expression, self.table_mapping, dialect=self.dialect, copy=False
@@ -327,11 +325,11 @@ class RuntimeAdapter(BaseAdapter):
 
         logger.debug("Resolved ref '%s' to snapshot table '%s'", name, physical_table_name)
 
-        return exp.to_table(physical_table_name, dialect=self.engine_adapter.dialect)
+        return exp.to_table(physical_table_name, dialect=self.dialect)
 
     def _relation_to_table(self, relation: BaseRelation) -> exp.Table:
-        table = exp.to_table(relation.render(), dialect=self.engine_adapter.dialect)
-        return exp.to_table(relation.render(), dialect=self.engine_adapter.dialect)
+        table = exp.to_table(relation.render(), dialect=self.dialect)
+        return exp.to_table(relation.render(), dialect=self.dialect)
 
     def _table_to_relation(self, table: exp.Table) -> BaseRelation:
         return self.relation_type.create(
@@ -351,7 +349,7 @@ class RuntimeAdapter(BaseAdapter):
 
     def _normalize(self, input_table: exp.Table) -> exp.Table:
         normalized_name = normalize_model_name(
-            input_table, self.engine_adapter.default_catalog, self.engine_adapter.dialect
+            input_table, self.engine_adapter.default_catalog, self.dialect
         )
         normalized_table = exp.to_table(normalized_name)
         if not input_table.this:
diff --git a/sqlmesh/dbt/builtin.py b/sqlmesh/dbt/builtin.py
index ad9f5713..ed2779b1 100644
--- a/sqlmesh/dbt/builtin.py
+++ b/sqlmesh/dbt/builtin.py
@@ -302,7 +302,7 @@ def create_builtin_globals(
 
     target: t.Optional[AttributeDict] = jinja_globals.get("target", None)
     api = Api(target)
-    dialect = target.type if target else None  # type: ignore
+    dialect = target.dialect if target else None  # type: ignore
 
     builtin_globals["api"] = api
 
diff --git a/sqlmesh/dbt/context.py b/sqlmesh/dbt/context.py
index ceb299ad..71888a07 100644
--- a/sqlmesh/dbt/context.py
+++ b/sqlmesh/dbt/context.py
@@ -62,7 +62,7 @@ class DbtContext:
     def dialect(self) -> str:
         if not self.target:
             raise SQLMeshError("Target must be configured before calling the dialect property.")
-        return TARGET_TYPE_TO_DIALECT_OVERRIDE.get(self.target.type) or self.target.type
+        return self.target.dialect
 
     @property
     def project_name(self) -> t.Optional[str]:
@@ -281,7 +281,3 @@ class DbtContext:
 
 
 SQLMESH_DBT_PACKAGE = "sqlmesh.dbt"
-
-TARGET_TYPE_TO_DIALECT_OVERRIDE: t.Dict[str, str] = {
-    "sqlserver": "tsql",
-}
diff --git a/sqlmesh/dbt/loader.py b/sqlmesh/dbt/loader.py
index 69c62e64..568c783b 100644
--- a/sqlmesh/dbt/loader.py
+++ b/sqlmesh/dbt/loader.py
@@ -42,7 +42,7 @@ def sqlmesh_config(
     context = DbtContext(project_root=project_root)
     profile = Profile.load(context, target_name=dbt_target_name)
     model_defaults = kwargs.get("model_defaults", ModelDefaultsConfig())
-    model_defaults.dialect = profile.target.type
+    model_defaults.dialect = profile.target.dialect
 
     loader_kwargs = kwargs.pop("loader_kwargs", {})
     if variables is not None:
diff --git a/sqlmesh/dbt/model.py b/sqlmesh/dbt/model.py
index fb826f5c..78f48c27 100644
--- a/sqlmesh/dbt/model.py
+++ b/sqlmesh/dbt/model.py
@@ -249,7 +249,7 @@ class ModelConfig(BaseModelConfig):
                 "valid_to_name": "dbt_valid_to",
                 "time_data_type": (
                     exp.DataType.build("TIMESTAMPTZ")
-                    if target.type == "bigquery"
+                    if target.dialect == "bigquery"
                     else exp.DataType.build("TIMESTAMP")
                 ),
             }
@@ -338,7 +338,7 @@ class ModelConfig(BaseModelConfig):
         if self.sql_header:
             model_kwargs["pre_statements"].insert(0, d.jinja_statement(self.sql_header))
 
-        if context.target.type == "bigquery":
+        if context.target.dialect == "bigquery":
             dbt_max_partition_blob = self._dbt_max_partition_blob()
             if dbt_max_partition_blob:
                 model_kwargs["pre_statements"].append(d.jinja_statement(dbt_max_partition_blob))
@@ -352,7 +352,7 @@ class ModelConfig(BaseModelConfig):
             if table_properties:
                 model_kwargs["table_properties"] = table_properties
 
-        if context.target.type == "snowflake" and self.snowflake_warehouse is not None:
+        if context.target.dialect == "snowflake" and self.snowflake_warehouse is not None:
             model_kwargs["session_properties"] = {"warehouse": self.snowflake_warehouse}
 
         return create_sql_model(
diff --git a/sqlmesh/dbt/target.py b/sqlmesh/dbt/target.py
index ef0d0af2..edc1dc19 100644
--- a/sqlmesh/dbt/target.py
+++ b/sqlmesh/dbt/target.py
@@ -134,6 +134,10 @@ class TargetConfig(abc.ABC, DbtConfig):
     def column_class(cls) -> t.Type[Column]:
         return Column
 
+    @property
+    def dialect(self) -> str:
+        return self.type
+
 
 DUCKDB_IN_MEMORY = ":memory:"
 
@@ -651,6 +655,10 @@ class MSSQLConfig(TargetConfig):
 
         return SQLServerColumn
 
+    @property
+    def dialect(self) -> str:
+        return "tsql"
+
     def to_sqlmesh(self) -> ConnectionConfig:
         return MSSQLConnectionConfig(
             host=self.host,


Test patch:
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index b687e5ba..2fbff176 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -1173,6 +1173,70 @@ def test_unpause_snapshots(state_sync: EngineAdapterStateSync, make_snapshot: t.
     assert not actual_snapshots[new_snapshot.snapshot_id].unrestorable
 
 
+def test_unrestorable_snapshot(state_sync: EngineAdapterStateSync, make_snapshot: t.Callable):
+    snapshot = make_snapshot(
+        SqlModel(
+            name="test_snapshot",
+            query=parse_one("select 1, ds"),
+            cron="@daily",
+        ),
+    )
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot.version = "a"
+
+    assert not snapshot.unpaused_ts
+    state_sync.push_snapshots([snapshot])
+
+    unpaused_dt = "2022-01-01"
+    state_sync.unpause_snapshots([snapshot], unpaused_dt)
+
+    actual_snapshot = state_sync.get_snapshots([snapshot])[snapshot.snapshot_id]
+    assert actual_snapshot.unpaused_ts
+    assert actual_snapshot.unpaused_ts == to_timestamp(unpaused_dt)
+
+    new_indirect_non_breaking_snapshot = make_snapshot(
+        SqlModel(name="test_snapshot", query=parse_one("select 2, ds"), cron="@daily")
+    )
+    new_indirect_non_breaking_snapshot.categorize_as(SnapshotChangeCategory.INDIRECT_NON_BREAKING)
+    new_indirect_non_breaking_snapshot.version = "a"
+
+    assert not new_indirect_non_breaking_snapshot.unpaused_ts
+    state_sync.push_snapshots([new_indirect_non_breaking_snapshot])
+    state_sync.unpause_snapshots([new_indirect_non_breaking_snapshot], unpaused_dt)
+
+    actual_snapshots = state_sync.get_snapshots([snapshot, new_indirect_non_breaking_snapshot])
+    assert not actual_snapshots[snapshot.snapshot_id].unpaused_ts
+    assert actual_snapshots[
+        new_indirect_non_breaking_snapshot.snapshot_id
+    ].unpaused_ts == to_timestamp(unpaused_dt)
+
+    assert not actual_snapshots[snapshot.snapshot_id].unrestorable
+    assert not actual_snapshots[new_indirect_non_breaking_snapshot.snapshot_id].unrestorable
+
+    new_forward_only_snapshot = make_snapshot(
+        SqlModel(name="test_snapshot", query=parse_one("select 3, ds"), cron="@daily")
+    )
+    new_forward_only_snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+    new_forward_only_snapshot.version = "a"
+
+    assert not new_forward_only_snapshot.unpaused_ts
+    state_sync.push_snapshots([new_forward_only_snapshot])
+    state_sync.unpause_snapshots([new_forward_only_snapshot], unpaused_dt)
+
+    actual_snapshots = state_sync.get_snapshots(
+        [snapshot, new_indirect_non_breaking_snapshot, new_forward_only_snapshot]
+    )
+    assert not actual_snapshots[snapshot.snapshot_id].unpaused_ts
+    assert not actual_snapshots[new_indirect_non_breaking_snapshot.snapshot_id].unpaused_ts
+    assert actual_snapshots[new_forward_only_snapshot.snapshot_id].unpaused_ts == to_timestamp(
+        unpaused_dt
+    )
+
+    assert actual_snapshots[snapshot.snapshot_id].unrestorable
+    assert actual_snapshots[new_indirect_non_breaking_snapshot.snapshot_id].unrestorable
+    assert not actual_snapshots[new_forward_only_snapshot.snapshot_id].unrestorable
+
+
 def test_unpause_snapshots_remove_intervals(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
 ):