Refactor!: move project from snapshot to model/node (#1345)

Author: tobymao

sqlmesh/core/context.py

@@ -1159,7 +1159,7 @@ def _snapshots(
         projects = {config.project for config in self.configs.values()}
 
         for name, snapshot in remote_snapshots.items():
-            if name not in models and snapshot.project not in projects:
+            if name not in models and snapshot.model.project not in projects:
                 models[name] = snapshot.model
 
                 for audit in snapshot.audits:
@@ -1172,19 +1172,16 @@ def _snapshots(
             if model.name not in self._models and model.name in remote_snapshots:
                 snapshot = remote_snapshots[model.name]
                 ttl = snapshot.ttl
-                project = snapshot.project
             else:
                 config = self.config_for_model(model)
                 ttl = config.snapshot_ttl
-                project = config.project
 
             snapshot = Snapshot.from_model(
                 model,
                 nodes=models,
                 audits=audits,
                 cache=fingerprint_cache,
                 ttl=ttl,
-                project=project,
             )
             snapshots[model.name] = snapshot
 

sqlmesh/core/loader.py

@@ -174,6 +174,7 @@ def _load_external_models(self) -> UniqueKeyDict[str, Model]:
                             **row,
                             dialect=config.model_defaults.dialect,
                             path=path,
+                            project=config.project,
                         )
                         models[model.name] = model
         return models
@@ -275,6 +276,7 @@ def _load() -> Model:
                         dialect=config.model_defaults.dialect,
                         time_column_format=config.time_column_format,
                         physical_schema_override=config.physical_schema_override,
+                        project=config.project,
                     )
 
                 model = cache.get_or_load_model(path, _load)
@@ -310,6 +312,7 @@ def _load_python_models(self) -> UniqueKeyDict[str, Model]:
                         dialect=config.model_defaults.dialect,
                         time_column_format=config.time_column_format,
                         physical_schema_override=config.physical_schema_override,
+                        project=config.project,
                     )
                     models[model.name] = model
 

sqlmesh/core/model/decorator.py

@@ -58,6 +58,7 @@ def model(
         dialect: t.Optional[str] = None,
         time_column_format: str = c.DEFAULT_TIME_COLUMN_FORMAT,
         physical_schema_override: t.Optional[t.Dict[str, str]] = None,
+        project: str = "",
     ) -> Model:
         """Get the model registered by this function."""
         env: t.Dict[str, t.Any] = {}
@@ -71,6 +72,7 @@ def model(
             time_column_format=time_column_format,
             python_env=serialize_env(env, path=module_path),
             physical_schema_override=physical_schema_override,
+            project=project,
             **self.kwargs,
         )
 

sqlmesh/core/model/definition.py

@@ -722,6 +722,7 @@ def metadata_hash(self, audits: t.Dict[str, Audit]) -> str:
             str(self.forward_only),
             str(self.disable_restatement),
             str(self.interval_unit_) if self.interval_unit_ is not None else None,
+            self.project,
         ]
 
         for audit_name, audit_args in sorted(self.audits, key=lambda a: a[0]):

sqlmesh/core/node.py

@@ -141,6 +141,7 @@ class Node(PydanticModel):
 
     Args:
         name: The name of the node.
+        description: The name of the project this node belongs to, used in multi-repo deployments.
         description: The optional node description.
         owner: The owner of the node.
         start: The earliest date that the node will be executed for. If this is None,
@@ -154,6 +155,7 @@ class Node(PydanticModel):
     """
 
     name: str
+    project: str = ""
     description: t.Optional[str] = None
     owner: t.Optional[str] = None
     start: t.Optional[TimeLike] = None

sqlmesh/core/snapshot/definition.py

@@ -361,7 +361,6 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         audits: The list of audits used by the model.
         intervals: List of [start, end) intervals showing which time ranges a snapshot has data for.
         dev_intervals: List of [start, end) intervals showing development intervals (forward-only).
-        project: The name of the project this snapshot is associated with.
         created_ts: Epoch millis timestamp when a snapshot was first created.
         updated_ts: Epoch millis timestamp when a snapshot was last updated.
         ttl: The time-to-live of a snapshot determines when it should be deleted after it's no longer referenced
@@ -385,7 +384,6 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
     audits: t.Tuple[Audit, ...]
     intervals: Intervals = []
     dev_intervals: Intervals = []
-    project: str = ""
     created_ts: int
     updated_ts: int
     ttl: str
@@ -473,7 +471,6 @@ def from_model(
         *,
         nodes: t.Dict[str, SnapshotNode],
         ttl: str = c.DEFAULT_SNAPSHOT_TTL,
-        project: str = "",
         version: t.Optional[str] = None,
         audits: t.Optional[t.Dict[str, Audit]] = None,
         cache: t.Optional[t.Dict[str, SnapshotFingerprint]] = None,
@@ -519,7 +516,6 @@ def from_model(
             audits=tuple(model.referenced_audits(audits)),
             intervals=[],
             dev_intervals=[],
-            project=project,
             created_ts=created_ts,
             updated_ts=created_ts,
             ttl=ttl,

sqlmesh/migrations/v0022_move_project_to_model.py

@@ -0,0 +1,53 @@
+"""Move project attr from snapshot to model."""
+import json
+
+import pandas as pd
+from sqlglot import exp
+
+from sqlmesh.utils.migration import index_text_type
+
+
+def migrate(state_sync):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    new_snapshots = []
+
+    for name, identifier, version, snapshot, kind_name in engine_adapter.fetchall(
+        exp.select("name", "identifier", "version", "snapshot", "kind_name").from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+
+        parsed_snapshot["node"]["project"] = parsed_snapshot.pop("project", "")
+
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": json.dumps(parsed_snapshot),
+                "kind_name": kind_name,
+            }
+        )
+
+    engine_adapter.delete_from(snapshots_table, "TRUE")
+
+    text_type = index_text_type(engine_adapter.dialect)
+
+    if new_snapshots:
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types={
+                "name": exp.DataType.build(text_type),
+                "identifier": exp.DataType.build(text_type),
+                "version": exp.DataType.build(text_type),
+                "snapshot": exp.DataType.build("text"),
+                "kind_name": exp.DataType.build("text"),
+            },
+            contains_json=True,
+        )

tests/core/test_integration.py

@@ -645,6 +645,7 @@ def test_multi(mocker):
         paths=["examples/multi/repo_1"], engine_adapter=context.engine_adapter, gateway="memory"
     )
     model = context.models["bronze.a"]
+    assert model.project == "repo_1"
     context.upsert_model(model.copy(update={"query": model.query.select("'c' AS c")}))
     plan = context.plan()
     assert set(snapshot.name for snapshot in plan.directly_modified) == {"bronze.a", "bronze.b"}

tests/core/test_snapshot.py

@@ -100,6 +100,7 @@ def test_json(snapshot: Snapshot):
             "dialect": "spark",
             "name": "name",
             "partitioned_by": [],
+            "project": "",
             "owner": "owner",
             "query": "SELECT @EACH([1, 2], x -> x), ds FROM parent.tbl",
             "jinja_macros": {
@@ -118,7 +119,6 @@ def test_json(snapshot: Snapshot):
         "name": "name",
         "parents": [{"name": "parent.tbl", "identifier": snapshot.parents[0].identifier}],
         "previous_versions": [],
-        "project": "",
         "indirect_versions": {},
         "updated_ts": 1663891973000,
         "version": snapshot.fingerprint.to_version(),
@@ -393,7 +393,7 @@ def test_fingerprint(model: Model, parent_model: Model):
 
     original_fingerprint = SnapshotFingerprint(
         data_hash="1116890341",
-        metadata_hash="1237394431",
+        metadata_hash="1312958471",
     )
 
     assert fingerprint == original_fingerprint
@@ -440,7 +440,7 @@ def test_fingerprint_seed_model():
 
     expected_fingerprint = SnapshotFingerprint(
         data_hash="1421766360",
-        metadata_hash="3585221762",
+        metadata_hash="1617581697",
     )
 
     model = load_sql_based_model(expressions, path=Path("./examples/sushi/models/test_model.sql"))
@@ -480,7 +480,7 @@ def test_fingerprint_jinja_macros(model: Model):
     )
     original_fingerprint = SnapshotFingerprint(
         data_hash="4053778362",
-        metadata_hash="1237394431",
+        metadata_hash="1312958471",
     )
 
     fingerprint = fingerprint_from_node(model, nodes={})

tests/schedulers/airflow/test_client.py

@@ -88,6 +88,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
                     "partitioned_by": ["a"],
                     "query": "SELECT a, ds FROM tbl",
                     "references": [],
+                    "project": "",
                     "storage_format": "parquet",
                     "jinja_macros": {
                         "global_objs": {},
@@ -104,7 +105,6 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
                 "name": "test_model",
                 "parents": [],
                 "previous_versions": [],
-                "project": "",
                 "updated_ts": 1665014400000,
                 "version": snapshot.version,
                 "change_category": snapshot.change_category,
@@ -131,7 +131,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
             "promoted_snapshot_ids": [
                 {
                     "name": "test_model",
-                    "identifier": "3192766394",
+                    "identifier": "3474928511",
                 }
             ],
             "suffix_target": "schema",