Fix!: Make sure internal models are not accidentally converted to external ones (#926)

Author: izeigerman

sqlmesh/core/context.py

@@ -853,6 +853,7 @@ def create_external_models(self) -> None:
                     if self.config_for_model(model) is config
                 },
                 adapter=self._engine_adapter,
+                state_reader=self.state_reader,
                 dialect=config.model_defaults.dialect,
                 max_workers=self.concurrent_tasks,
             )

sqlmesh/core/model/definition.py

@@ -348,8 +348,12 @@ def text_diff(self, other: Model) -> str:
         Returns:
             A unified text diff showing additions and deletions.
         """
-        meta_a, *statements_a, query_a = self.render_definition()
-        meta_b, *statements_b, query_b = other.render_definition()
+        meta_a, *statements_a = self.render_definition()
+        meta_b, *statements_b = other.render_definition()
+
+        query_a = statements_a.pop() if statements_a else None
+        query_b = statements_b.pop() if statements_b else None
+
         return "\n".join(
             (
                 d.text_diff(meta_a, meta_b, self.dialect),

sqlmesh/core/plan/evaluator.py

@@ -159,17 +159,20 @@ def _restate(self, plan: Plan) -> None:
         if not plan.restatements:
             return
 
-        all_snapshots = (
-            [s for s in plan.snapshots if s.name in plan.restatements]
-            if plan.is_dev
-            else self.state_sync.get_snapshots_by_models(*plan.restatements)
-        )
-        self.state_sync.remove_interval(
-            [],
-            start=plan.start,
-            end=plan.end,
-            all_snapshots=all_snapshots,
-        )
+        target_snapshots = [s for s in plan.snapshots if s.name in plan.restatements]
+        if plan.is_dev:
+            self.state_sync.remove_interval(
+                [],
+                start=plan.start,
+                end=plan.end,
+                all_snapshots=target_snapshots,
+            )
+        else:
+            self.state_sync.remove_interval(
+                target_snapshots,
+                start=plan.start,
+                end=plan.end,
+            )
 
 
 class AirflowPlanEvaluator(PlanEvaluator):

sqlmesh/core/schema_loader.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import typing as t
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
@@ -9,12 +10,16 @@
 
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.model import Model
+from sqlmesh.core.state_sync import StateReader
+
+logger = logging.getLogger(__name__)
 
 
 def create_schema_file(
     path: Path,
     models: t.Dict[str, Model],
     adapter: EngineAdapter,
+    state_reader: StateReader,
     dialect: DialectType,
     max_workers: int = 1,
 ) -> None:
@@ -24,13 +29,24 @@ def create_schema_file(
         path: The path to store the YAML file.
         models: A dictionary of models to fetch columns from the db.
         adapter: The engine adapter.
+        state_reader: The state reader.
         dialect: The dialect to serialize the schema as.
         max_workers: The max concurrent workers to fetch columns.
     """
     external_tables = {
         dep for model in models.values() for dep in model.depends_on if dep not in models
     }
 
+    # Make sure we don't convert internal models into external ones.
+    existing_models = state_reader.models_exist(external_tables, exclude_external=True)
+    if existing_models:
+        logger.warning(
+            "The following models already exist and can't be converted to external: %s."
+            "Perhaps these models have been removed, while downstream models that reference them weren't updated accordingly",
+            ", ".join(existing_models),
+        )
+        external_tables -= existing_models
+
     with ThreadPoolExecutor(max_workers=max_workers) as pool:
         schemas = [
             {

sqlmesh/core/state_sync/base.py

@@ -76,6 +76,18 @@ def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[Sna
             A set of all the existing snapshot ids.
         """
 
+    @abc.abstractmethod
+    def models_exist(self, names: t.Iterable[str], exclude_external: bool = False) -> t.Set[str]:
+        """Returns the model names that exist in the state sync.
+
+        Args:
+            names: Iterable of model names to check.
+            exclude_external: Whether to exclude external models from the output.
+
+        Returns:
+            A set of all the existing model names.
+        """
+
     @abc.abstractmethod
     def get_environment(self, environment: str) -> t.Optional[Environment]:
         """Fetches the environment if it exists.
@@ -95,14 +107,6 @@ def get_environments(self) -> t.List[Environment]:
             A list of all environments.
         """
 
-    @abc.abstractmethod
-    def get_snapshots_by_models(self, *names: str) -> t.List[Snapshot]:
-        """Get all snapshots by model name.
-
-        Returns:
-            The list of snapshots.
-        """
-
     @abc.abstractmethod
     def get_snapshot_intervals(
         self, snapshots: t.Optional[t.Iterable[SnapshotNameVersionLike]]

sqlmesh/core/state_sync/common.py

@@ -49,21 +49,6 @@ def get_snapshots(
     def get_environment(self, environment: str) -> t.Optional[Environment]:
         return self._get_environment(environment)
 
-    def get_snapshots_by_models(
-        self, *names: str, lock_for_update: bool = False
-    ) -> t.List[Snapshot]:
-        """
-        Get all snapshots by model name.
-
-        Returns:
-            The list of snapshots.
-        """
-        return [
-            snapshot
-            for snapshot in self._get_snapshots(lock_for_update=lock_for_update).values()
-            if snapshot.name in names
-        ]
-
     @transactional()
     def promote(
         self, environment: Environment, no_gaps: bool = False

sqlmesh/core/state_sync/engine_adapter.py

@@ -30,7 +30,7 @@
 from sqlmesh.core.audit import Audit
 from sqlmesh.core.engine_adapter import EngineAdapter, TransactionType
 from sqlmesh.core.environment import Environment
-from sqlmesh.core.model import Model, SeedModel
+from sqlmesh.core.model import Model, ModelKindName, SeedModel
 from sqlmesh.core.snapshot import (
     Intervals,
     Snapshot,
@@ -87,6 +87,7 @@ def __init__(
             "identifier": exp.DataType.build("text"),
             "version": exp.DataType.build("text"),
             "snapshot": exp.DataType.build("text"),
+            "kind_name": exp.DataType.build("text"),
         }
 
         self._environment_columns_to_types = {
@@ -239,6 +240,17 @@ def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[Sna
             )
         }
 
+    def models_exist(self, names: t.Iterable[str], exclude_external: bool = False) -> t.Set[str]:
+        query = (
+            exp.select("name")
+            .from_(self.snapshots_table)
+            .where(exp.column("name").isin(*names))
+            .distinct()
+        )
+        if exclude_external:
+            query = query.where(exp.column("kind_name").neq(ModelKindName.EXTERNAL.value))
+        return {name for name, in self.engine_adapter.fetchall(query)}
+
     def reset(self) -> None:
         """Resets the state store to the state when it was first initialized."""
         self.engine_adapter.drop_table(self.snapshots_table)
@@ -821,6 +833,7 @@ def _snapshots_to_df(snapshots: t.Iterable[Snapshot]) -> pd.DataFrame:
                 "identifier": snapshot.identifier,
                 "version": snapshot.version,
                 "snapshot": snapshot.json(exclude={"intervals", "dev_intervals"}),
+                "kind_name": snapshot.model_kind_name.value,
             }
             for snapshot in snapshots
         ]

sqlmesh/migrations/v0011_add_model_kind_name.py

@@ -0,0 +1,54 @@
+"""Add the kind_name column to the snapshots table."""
+import json
+
+import pandas as pd
+from sqlglot import exp
+
+
+def migrate(state_sync):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = f"{schema}._snapshots"
+
+    alter_table_exp = exp.AlterTable(
+        this=exp.to_table(snapshots_table),
+        actions=[
+            exp.ColumnDef(
+                this=exp.to_column("kind_name"),
+                kind=exp.DataType.build("text"),
+            )
+        ],
+    )
+    engine_adapter.execute(alter_table_exp)
+
+    new_snapshots = []
+
+    for name, identifier, version, snapshot in engine_adapter.fetchall(
+        exp.select("name", "identifier", "version", "snapshot").from_(snapshots_table)
+    ):
+        parsed_snapshot = json.loads(snapshot)
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": snapshot,
+                "kind_name": parsed_snapshot["model"]["kind"]["name"],
+            }
+        )
+
+    if new_snapshots:
+        engine_adapter.delete_from(snapshots_table, "TRUE")
+
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types={
+                "name": exp.DataType.build("text"),
+                "identifier": exp.DataType.build("text"),
+                "version": exp.DataType.build("text"),
+                "snapshot": exp.DataType.build("text"),
+                "kind_name": exp.DataType.build("text"),
+            },
+            contains_json=True,
+        )

sqlmesh/schedulers/airflow/api.py

@@ -101,6 +101,17 @@ def get_snapshot_intervals() -> Response:
         )
 
 
+@sqlmesh_api_v1.get("/models")
+@csrf.exempt
+@check_authentication
+def models_exist() -> Response:
+    with util.scoped_state_sync() as state_sync:
+        names = _csv_arg("names")
+        exclude_external = "exclude_external" in request.args
+        existing_models = state_sync.models_exist(names, exclude_external=exclude_external)
+        return _success(common.ExistingModelsResponse(names=list(existing_models)))
+
+
 @sqlmesh_api_v1.get("/versions")
 @csrf.exempt
 @check_authentication
@@ -138,3 +149,9 @@ def _snapshot_name_versions_from_request() -> t.Optional[t.List[SnapshotNameVers
 
     raw_versions = json.loads(request.args["versions"])
     return [SnapshotNameVersion.parse_obj(v) for v in raw_versions]
+
+
+def _csv_arg(arg: str) -> t.List[str]:
+    if arg not in request.args:
+        return []
+    return [v.strip() for v in request.args[arg].split(",")]

sqlmesh/schedulers/airflow/client.py

@@ -35,6 +35,7 @@
 SNAPSHOTS_PATH = f"{common.SQLMESH_API_BASE_PATH}/snapshots"
 SEEDS_PATH = f"{common.SQLMESH_API_BASE_PATH}/seeds"
 INTERVALS_PATH = f"{common.SQLMESH_API_BASE_PATH}/intervals"
+MODELS_PATH = f"{common.SQLMESH_API_BASE_PATH}/models"
 VERSIONS_PATH = f"{common.SQLMESH_API_BASE_PATH}/versions"
 
 
@@ -105,6 +106,14 @@ def snapshots_exist(self, snapshot_ids: t.List[SnapshotId]) -> t.Set[SnapshotId]
             ).snapshot_ids
         )
 
+    def models_exist(self, names: t.Iterable[str], exclude_external: bool = False) -> t.Set[str]:
+        flags = ["exclude_external"] if exclude_external else []
+        return set(
+            common.ExistingModelsResponse.parse_obj(
+                self._get(MODELS_PATH, *flags, names=",".join(names))
+            ).names
+        )
+
     def get_snapshot_intervals(
         self, snapshot_name_versions: t.List[SnapshotNameVersion]
     ) -> t.List[SnapshotIntervals]: