Refactor!: Replace snapshot model with snapshot node (#1137)

Author: crericha

sqlmesh/core/audit/definition.py

@@ -10,7 +10,7 @@
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
 from sqlmesh.core.model.common import bool_validator, expression_validator
-from sqlmesh.core.model.definition import Model, _Model
+from sqlmesh.core.model.definition import _Model
 from sqlmesh.core.renderer import QueryRenderer
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import AuditConfigError, SQLMeshError, raise_config_error
@@ -160,7 +160,7 @@ def load_multiple(
 
     def render_query(
         self,
-        snapshot_or_model: t.Union[Snapshot, Model],
+        snapshot_or_model: t.Union[Snapshot, _Model],
         *,
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
@@ -238,7 +238,7 @@ def macro_definitions(self) -> t.List[d.MacroDef]:
         """All macro definitions from the list of expressions."""
         return [s for s in self.expressions if isinstance(s, d.MacroDef)]
 
-    def _create_query_renderer(self, model: Model) -> QueryRenderer:
+    def _create_query_renderer(self, model: _Model) -> QueryRenderer:
         return QueryRenderer(
             self.query,
             self.dialect or model.dialect,

sqlmesh/core/context.py

@@ -543,7 +543,7 @@ def snapshots(self) -> t.Dict[str, Snapshot]:
 
             snapshot = Snapshot.from_model(
                 model,
-                models=models,
+                nodes=models,
                 audits=audits,
                 cache=fingerprint_cache,
                 ttl=ttl,
@@ -555,7 +555,7 @@ def snapshots(self) -> t.Dict[str, Snapshot]:
 
         for snapshot in stored_snapshots.values():
             # Keep the original model instance to preserve the query cache.
-            snapshot.model = snapshots[snapshot.name].model
+            snapshot.node = snapshots[snapshot.name].node
 
         return {name: stored_snapshots.get(s.snapshot_id, s) for name, s in snapshots.items()}
 

sqlmesh/core/context_diff.py

@@ -127,7 +127,7 @@ def create(
 
             if existing:
                 # Keep the original model instance to preserve the query cache.
-                existing.model = snapshot.model
+                existing.node = snapshot.node
 
                 merged_snapshots[name] = existing.copy()
                 if modified:

sqlmesh/core/engine_adapter/base.py

@@ -39,7 +39,7 @@
         Query,
         QueryOrDF,
     )
-    from sqlmesh.core.model.meta import IntervalUnit
+    from sqlmesh.core.node import IntervalUnit
 
 logger = logging.getLogger(__name__)
 

sqlmesh/core/engine_adapter/bigquery.py

@@ -16,7 +16,7 @@
     DataObjectType,
     TransactionType,
 )
-from sqlmesh.core.model.meta import IntervalUnit
+from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils.date import to_datetime
 from sqlmesh.utils.errors import SQLMeshError

sqlmesh/core/engine_adapter/spark.py

@@ -22,7 +22,7 @@
         PySparkSession,
         QueryOrDF,
     )
-    from sqlmesh.core.model.meta import IntervalUnit
+    from sqlmesh.core.node import IntervalUnit
 
 
 class SparkEngineAdapter(EngineAdapter):

sqlmesh/core/model/definition.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import json
 import logging
 import sys
 import types
@@ -23,13 +24,20 @@
 from sqlmesh.core import dialect as d
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.model.common import expression_validator
-from sqlmesh.core.model.kind import ModelKindName, SeedKind, _Incremental
+from sqlmesh.core.model.kind import (
+    IncrementalByTimeRangeKind,
+    IncrementalByUniqueKeyKind,
+    ModelKindName,
+    SeedKind,
+    _Incremental,
+)
 from sqlmesh.core.model.meta import ModelMeta
 from sqlmesh.core.model.seed import Seed, create_seed
 from sqlmesh.core.renderer import ExpressionRenderer, QueryRenderer
 from sqlmesh.utils import str_to_bool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_datetime
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, raise_config_error
+from sqlmesh.utils.hashing import hash_data
 from sqlmesh.utils.jinja import JinjaMacroRegistry, extract_macro_references
 from sqlmesh.utils.metaprogramming import (
     Executable,
@@ -48,9 +56,9 @@
     from sqlmesh.utils.jinja import MacroReference
 
 if sys.version_info >= (3, 9):
-    from typing import Annotated, Literal
+    from typing import Literal
 else:
-    from typing_extensions import Annotated, Literal
+    from typing_extensions import Literal
 
 logger = logging.getLogger(__name__)
 
@@ -628,6 +636,103 @@ def is_breaking_change(self, previous: Model) -> t.Optional[bool]:
         """
         raise NotImplementedError
 
+    @property
+    def data_hash(self) -> str:
+        """
+        Computes the data hash for the node.
+
+        Returns:
+            The data hash for the node.
+        """
+        return hash_data(self._data_hash_fields)
+
+    @property
+    def _data_hash_fields(self) -> t.List[str]:
+        data = [
+            str(self.sorted_python_env),
+            self.kind.name,
+            self.cron,
+            self.storage_format,
+            str(self.lookback),
+            *(expr.sql() for expr in (self.partitioned_by or [])),
+            *(self.clustered_by or []),
+            self.stamp,
+        ]
+
+        for column_name, column_type in (self.columns_to_types_ or {}).items():
+            data.append(column_name)
+            data.append(column_type.sql())
+
+        if isinstance(self.kind, IncrementalByTimeRangeKind):
+            data.append(self.kind.time_column.column)
+            data.append(self.kind.time_column.format)
+        elif isinstance(self.kind, IncrementalByUniqueKeyKind):
+            data.extend(self.kind.unique_key)
+
+        return data  # type: ignore
+
+    def metadata_hash(self, audits: t.Dict[str, Audit]) -> str:
+        """
+        Computes the metadata hash for the node.
+
+        Args:
+            audits: Available audits by name.
+
+        Returns:
+            The metadata hash for the node.
+        """
+        from sqlmesh.core.audit import BUILT_IN_AUDITS
+
+        metadata = [
+            self.dialect,
+            self.owner,
+            self.description,
+            str(self.start) if self.start else None,
+            str(self.retention) if self.retention else None,
+            str(self.batch_size) if self.batch_size is not None else None,
+            json.dumps(self.mapping_schema, sort_keys=True),
+            *sorted(self.tags),
+            *sorted(self.grain),
+            str(self.forward_only),
+            str(self.disable_restatement),
+        ]
+
+        for audit_name, audit_args in sorted(self.audits, key=lambda a: a[0]):
+            metadata.append(audit_name)
+
+            if audit_name in BUILT_IN_AUDITS:
+                for arg_name, arg_value in audit_args.items():
+                    metadata.append(arg_name)
+                    metadata.append(arg_value.sql(comments=True))
+            elif audit_name in audits:
+                audit = audits[audit_name]
+                query = (
+                    audit.query
+                    if self.hash_raw_query
+                    else audit.render_query(self, **t.cast(t.Dict[str, t.Any], audit_args))
+                    or audit.query
+                )
+                metadata.extend(
+                    [
+                        query.sql(comments=True),
+                        audit.dialect,
+                        str(audit.skip),
+                        str(audit.blocking),
+                    ]
+                )
+            else:
+                raise SQLMeshError(f"Unexpected audit name '{audit_name}'.")
+
+        # Add comments from the query.
+        if self.is_sql:
+            rendered_query = self.render_query()
+            if rendered_query:
+                for e, _, _ in rendered_query.walk():
+                    if e.comments:
+                        metadata.extend(e.comments)
+
+        return hash_data(metadata)
+
 
 class _SqlBasedModel(_Model):
     pre_statements_: t.Optional[t.List[exp.Expression]] = Field(
@@ -728,6 +833,20 @@ def _statement_renderer(self, expression: exp.Expression) -> ExpressionRenderer:
             )
         return self.__statement_renderers[expression_key]
 
+    @property
+    def _data_hash_fields(self) -> t.List[str]:
+        pre_statements = (
+            self.pre_statements if self.hash_raw_query else self.render_pre_statements()
+        )
+        post_statements = (
+            self.post_statements if self.hash_raw_query else self.render_post_statements()
+        )
+        macro_defs = self.macro_definitions if self.hash_raw_query else []
+        return [
+            *super()._data_hash_fields,
+            *[e.sql(comments=False) for e in (*pre_statements, *post_statements, *macro_defs)],
+        ]
+
 
 class SqlModel(_SqlBasedModel):
     """The model definition which relies on a SQL query to fetch the data.
@@ -926,6 +1045,24 @@ def _query_renderer(self) -> QueryRenderer:
             )
         return self.__query_renderer
 
+    @property
+    def _data_hash_fields(self) -> t.List[str]:
+        data = super()._data_hash_fields
+
+        query = self.query if self.hash_raw_query else self.render_query() or self.query
+        data.append(query.sql(comments=False))
+
+        for macro_name, macro in sorted(self.jinja_macros.root_macros.items()):
+            data.append(macro_name)
+            data.append(macro.definition)
+
+        for _, package in sorted(self.jinja_macros.packages.items(), key=lambda x: x[0]):
+            for macro_name, macro in sorted(package.items(), key=lambda x: x[0]):
+                data.append(macro_name)
+                data.append(macro.definition)
+
+        return data
+
     def __repr__(self) -> str:
         return f"Model<name: {self.name}, query: {self.query.sql(dialect=self.dialect)[0:30]}>"
 
@@ -1087,6 +1224,14 @@ def _ensure_hydrated(self) -> None:
         if not self.is_hydrated:
             raise SQLMeshError(f"Seed model '{self.name}' is not hydrated.")
 
+    @property
+    def _data_hash_fields(self) -> t.List[str]:
+        data = super()._data_hash_fields
+        for column_name, column_hash in self.column_hashes.items():
+            data.append(column_name)
+            data.append(column_hash)
+        return data
+
     def __repr__(self) -> str:
         return f"Model<name: {self.name}, seed: {self.kind.path}>"
 
@@ -1139,6 +1284,12 @@ def is_python(self) -> bool:
     def is_breaking_change(self, previous: Model) -> t.Optional[bool]:
         return None
 
+    @property
+    def _data_hash_fields(self) -> t.List[str]:
+        data = super()._data_hash_fields
+        data.append(self.entrypoint)
+        return data
+
     def __repr__(self) -> str:
         return f"Model<name: {self.name}, entrypoint: {self.entrypoint}>"
 
@@ -1156,9 +1307,7 @@ def is_breaking_change(self, previous: Model) -> t.Optional[bool]:
         return None
 
 
-Model = Annotated[
-    t.Union[SqlModel, SeedModel, PythonModel, ExternalModel], Field(discriminator="source_type")
-]
+Model = t.Union[SqlModel, SeedModel, PythonModel, ExternalModel]
 
 
 def load_model(