Feat: Add external models and loader to improve optimizer. (#871)

* Feat!: Add external models and loader to improve optimizer.

* Update sqlmesh/core/model/definition.py

Co-authored-by: Vincent Chan <vchan@users.noreply.github.com>

* Update sqlmesh/core/schema_loader.py

Co-authored-by: Vincent Chan <vchan@users.noreply.github.com>

* Chore: bump sqlglot

* Chore: cleanup

* pr feedback

* fix remaining issues

* Update setup.py

Co-authored-by: Jo <46752250+GeorgeSittas@users.noreply.github.com>

---------

Co-authored-by: Vincent Chan <vchan@users.noreply.github.com>
Co-authored-by: Jo <46752250+GeorgeSittas@users.noreply.github.com>

Author: 3 people

setup.py

@@ -44,7 +44,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot~=13.0.1",
+        "sqlglot~=13.2.1",
         "fsspec",
     ],
     extras_require={

sqlmesh/cli/main.py

@@ -379,6 +379,14 @@ def rollback(obj: Context) -> None:
     obj.rollback()
 
 
+@cli.command("create_external_models")
+@click.pass_obj
+@error_handler
+def create_external_models(obj: Context) -> None:
+    """Create a schema file containing external model schemas."""
+    obj.create_external_models()
+
+
 @cli.command("prompt")
 @click.argument("prompt")
 @click.option(

sqlmesh/core/console.py

@@ -470,12 +470,12 @@ def _snapshot_change_choices(
         indirect = "indirectly modified children"
         if use_rich_formatting:
             indirect = f"[indirect]{indirect}[/indirect]"
-        if snapshot.is_view_kind:
+        if snapshot.is_view:
             choices = {
                 SnapshotChangeCategory.BREAKING: f"Update {direct} and backfill {indirect}",
                 SnapshotChangeCategory.NON_BREAKING: f"Update {direct} but don't backfill {indirect}",
             }
-        elif snapshot.is_embedded_kind:
+        elif snapshot.is_symbolic:
             choices = {
                 SnapshotChangeCategory.BREAKING: f"Backfill {indirect}",
                 SnapshotChangeCategory.NON_BREAKING: f"Don't backfill {indirect}",

sqlmesh/core/constants.py

@@ -37,3 +37,4 @@
 MODELS = "models"
 TESTS = "tests"
 CACHE = ".cache"
+SCHEMA_YAML = "schema.yaml"

sqlmesh/core/context.py

@@ -59,6 +59,7 @@
 from sqlmesh.core.model.definition import _Model
 from sqlmesh.core.plan import Plan
 from sqlmesh.core.scheduler import Scheduler
+from sqlmesh.core.schema_loader import create_schema_file
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotEvaluator,
@@ -841,6 +842,25 @@ def rollback(self) -> None:
         """
         self._new_state_sync().rollback()
 
+    def create_external_models(self) -> None:
+        """Create a schema file with all external models.
+
+        The schema file contains all columns and types of external models, allowing for more robust
+        lineage, validation, and optimizations.
+        """
+        for path, config in self.configs.items():
+            create_schema_file(
+                path=path / c.SCHEMA_YAML,
+                models={
+                    name: model
+                    for name, model in self._models.items()
+                    if self.config_for_model(model) is config
+                },
+                adapter=self._engine_adapter,
+                dialect=config.model_defaults.dialect,
+                max_workers=self.concurrent_tasks,
+            )
+
     def print_info(self) -> None:
         """Prints information about connections, models, macros, etc. to the console."""
         self.console.log_status_update(f"Models: {len(self.models)}")

sqlmesh/core/engine_adapter/base_spark.py

@@ -82,7 +82,7 @@ def _create_table_properties(
         format_property = None
         partition_columns_property = None
         if storage_format:
-            format_property = exp.TableFormatProperty(this=exp.Var(this=storage_format))
+            format_property = exp.FileFormatProperty(this=exp.Var(this=storage_format))
         if partitioned_by:
             partition_columns_property = exp.PartitionedByProperty(
                 this=exp.Schema(

sqlmesh/core/loader.py

@@ -12,16 +12,23 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from ruamel.yaml import YAML
 from sqlglot.errors import SqlglotError
 from sqlglot.optimizer.qualify_columns import validate_qualify_columns
-from sqlglot.schema import MappingSchema, _nested_set
+from sqlglot.schema import MappingSchema, nested_set
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import Audit
 from sqlmesh.core.dialect import parse
 from sqlmesh.core.hooks import HookRegistry, hook
 from sqlmesh.core.macros import MacroRegistry, macro
-from sqlmesh.core.model import Model, ModelCache, SeedModel, load_model
+from sqlmesh.core.model import (
+    Model,
+    ModelCache,
+    SeedModel,
+    create_external_model,
+    load_model,
+)
 from sqlmesh.core.model import model as model_registry
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.dag import DAG
@@ -51,7 +58,7 @@ def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> No
             mapping_schema = schema.find(table)
 
             if mapping_schema:
-                _nested_set(
+                nested_set(
                     model.mapping_schema,
                     tuple(str(part) for part in table.parts),
                     {k: str(v) for k, v in mapping_schema.items()},
@@ -228,10 +235,29 @@ def _load_models(
         audits into a Dict and creates the dag
         """
         models = self._load_sql_models(macros, hooks, jinja_macros)
+        models.update(self._load_external_models())
         models.update(self._load_python_models())
 
         return models
 
+    def _load_external_models(self) -> UniqueKeyDict[str, Model]:
+        models: UniqueKeyDict = UniqueKeyDict("models")
+        for context_path, config in self._context.configs.items():
+            path = Path(context_path / c.SCHEMA_YAML)
+
+            if path.exists():
+                self._track_file(path)
+
+                with open(path, "r", encoding="utf-8") as file:
+                    for row in YAML().load(file.read()):
+                        model = create_external_model(
+                            **row,
+                            dialect=config.model_defaults.dialect,
+                            path=path,
+                        )
+                        models[model.name] = model
+        return models
+
     def _load_sql_models(
         self, macros: MacroRegistry, hooks: HookRegistry, jinja_macros: JinjaMacroRegistry
     ) -> UniqueKeyDict[str, Model]:

sqlmesh/core/model/__init__.py

@@ -6,6 +6,7 @@
     PythonModel,
     SeedModel,
     SqlModel,
+    create_external_model,
     create_python_model,
     create_seed_model,
     create_sql_model,
@@ -15,6 +16,7 @@
     IncrementalByTimeRangeKind,
     IncrementalByUniqueKeyKind,
     ModelKind,
+    ModelKindMixin,
     ModelKindName,
     SeedKind,
     TimeColumn,

sqlmesh/core/model/definition.py

@@ -881,7 +881,15 @@ def __repr__(self) -> str:
         return f"Model<name: {self.name}, entrypoint: {self.entrypoint}>"
 
 
-Model = Annotated[t.Union[SqlModel, SeedModel, PythonModel], Field(discriminator="source_type")]
+class ExternalModel(_Model):
+    """The model definition which represents an external source/table."""
+
+    source_type: Literal["external"] = "external"
+
+
+Model = Annotated[
+    t.Union[SqlModel, SeedModel, PythonModel, ExternalModel], Field(discriminator="source_type")
+]
 
 
 def load_model(
@@ -1129,6 +1137,30 @@ def create_python_model(
     )
 
 
+def create_external_model(
+    name: str,
+    *,
+    dialect: t.Optional[str] = None,
+    path: Path = Path(),
+    **kwargs: t.Any,
+) -> Model:
+    """Creates an external model.
+
+    Args:
+        name: The name of the model, which is of the form [catalog].[db].table.
+            The catalog and db are optional.
+        dialect: The dialect to serialize.
+        path: An optional path to the model definition file.
+    """
+    return _create_model(
+        ExternalModel,
+        name,
+        dialect=dialect,
+        path=path,
+        **kwargs,
+    )
+
+
 def _create_model(
     klass: t.Type[_Model],
     name: str,

sqlmesh/core/model/kind.py

@@ -13,68 +13,81 @@
 from sqlmesh.utils.pydantic import PydanticModel
 
 
-# TODO: switch to autoname when sqlglot is typed
-class ModelKindName(str, Enum):
-    """The kind of model, determining how this data is computed and stored in the warehouse."""
-
-    INCREMENTAL_BY_TIME_RANGE = "INCREMENTAL_BY_TIME_RANGE"
-    INCREMENTAL_BY_UNIQUE_KEY = "INCREMENTAL_BY_UNIQUE_KEY"
-    FULL = "FULL"
-    VIEW = "VIEW"
-    EMBEDDED = "EMBEDDED"
-    SEED = "SEED"
-    # TODO: Add support for snapshots
-    # SNAPSHOT = "SNAPSHOT"
-
-
-class ModelKind(PydanticModel):
-    name: ModelKindName
+class ModelKindMixin:
+    @property
+    def model_kind_name(self) -> ModelKindName:
+        """Returns the model kind name."""
+        raise NotImplementedError
 
     @property
     def is_incremental_by_time_range(self) -> bool:
-        return self.name == ModelKindName.INCREMENTAL_BY_TIME_RANGE
+        return self.model_kind_name == ModelKindName.INCREMENTAL_BY_TIME_RANGE
 
     @property
     def is_incremental_by_unique_key(self) -> bool:
-        return self.name == ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
-
-    @property
-    def is_incremental(self) -> bool:
-        """Whether or not this model is incremental."""
-        return isinstance(self, _Incremental)
+        return self.model_kind_name == ModelKindName.INCREMENTAL_BY_UNIQUE_KEY
 
     @property
     def is_full(self) -> bool:
-        return self.name == ModelKindName.FULL
-
-    # @property
-    # def is_snapshot(self) -> bool:
-    #     return self.name == ModelKindName.SNAPSHOT
+        return self.model_kind_name == ModelKindName.FULL
 
     @property
     def is_view(self) -> bool:
-        return self.name == ModelKindName.VIEW
+        return self.model_kind_name == ModelKindName.VIEW
 
     @property
     def is_embedded(self) -> bool:
-        return self.name == ModelKindName.EMBEDDED
+        return self.model_kind_name == ModelKindName.EMBEDDED
 
     @property
     def is_seed(self) -> bool:
-        return self.name == ModelKindName.SEED
+        return self.model_kind_name == ModelKindName.SEED
+
+    @property
+    def is_external(self) -> bool:
+        return self.model_kind_name == ModelKindName.EXTERNAL
+
+    @property
+    def is_symbolic(self) -> bool:
+        """A symbolic model is one that doesn't execute at all."""
+        return self.model_kind_name in (ModelKindName.EMBEDDED, ModelKindName.EXTERNAL)
 
     @property
     def is_materialized(self) -> bool:
-        return self.name not in (ModelKindName.VIEW, ModelKindName.EMBEDDED)
+        return not (self.is_symbolic or self.is_view)
 
     @property
     def only_latest(self) -> bool:
         """Whether or not this model only cares about latest date to render."""
-        return self.name in (ModelKindName.VIEW, ModelKindName.FULL)
+        return self.model_kind_name in (ModelKindName.VIEW, ModelKindName.FULL)
+
+
+class ModelKindName(str, ModelKindMixin, Enum):
+    """The kind of model, determining how this data is computed and stored in the warehouse."""
+
+    INCREMENTAL_BY_TIME_RANGE = "INCREMENTAL_BY_TIME_RANGE"
+    INCREMENTAL_BY_UNIQUE_KEY = "INCREMENTAL_BY_UNIQUE_KEY"
+    FULL = "FULL"
+    VIEW = "VIEW"
+    EMBEDDED = "EMBEDDED"
+    SEED = "SEED"
+    EXTERNAL = "EXTERNAL"
+
+    @property
+    def model_kind_name(self) -> ModelKindName:
+        return self
+
+
+class ModelKind(PydanticModel, ModelKindMixin):
+    name: ModelKindName
 
     def to_expression(self, **kwargs: t.Any) -> d.ModelKind:
         return d.ModelKind(this=self.name.value.upper(), **kwargs)
 
+    @property
+    def model_kind_name(self) -> ModelKindName:
+        return self.name
+
 
 class TimeColumn(PydanticModel):
     column: str