sqlmesh 0.209.1.dev32__tar.gz → 0.209.1.dev34__tar.gz

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.209.1.dev32
+Version: 0.209.1.dev34
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.209.1.dev32'
-__version_tuple__ = version_tuple = (0, 209, 1, 'dev32')
+__version__ = version = '0.209.1.dev34'
+__version_tuple__ = version_tuple = (0, 209, 1, 'dev34')

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh/core/selector.py

@@ -178,10 +178,10 @@ class Selector:
 
         node = parse(" | ".join(f"({s})" for s in model_selections))
 
-        models = models or self._models
+        all_models = models or self._models
         models_by_tags: t.Dict[str, t.Set[str]] = {}
 
-        for fqn, model in models.items():
+        for fqn, model in all_models.items():
             for tag in model.tags:
                 tag = tag.lower()
                 models_by_tags.setdefault(tag, set())
@@ -193,11 +193,11 @@ class Selector:
                 if "*" in pattern:
                     return {
                         fqn
-                        for fqn, model in models.items()
+                        for fqn, model in all_models.items()
                         if fnmatch.fnmatchcase(model.name, node.this)
                     }
                 fqn = normalize_model_name(pattern, self._default_catalog, self._dialect)
-                return {fqn} if fqn in models else set()
+                return {fqn} if fqn in all_models else set()
             if isinstance(node, exp.And):
                 return evaluate(node.left) & evaluate(node.right)
             if isinstance(node, exp.Or):
@@ -205,7 +205,7 @@ class Selector:
             if isinstance(node, exp.Paren):
                 return evaluate(node.this)
             if isinstance(node, exp.Not):
-                return set(models) - evaluate(node.this)
+                return set(all_models) - evaluate(node.this)
             if isinstance(node, Git):
                 target_branch = node.name
                 git_modified_files = {
@@ -213,7 +213,7 @@ class Selector:
                     *self._git_client.list_uncommitted_changed_files(),
                     *self._git_client.list_committed_changed_files(target_branch=target_branch),
                 }
-                return {m.fqn for m in self._models.values() if m._path in git_modified_files}
+                return {m.fqn for m in all_models.values() if m._path in git_modified_files}
             if isinstance(node, Tag):
                 pattern = node.name.lower()
 
@@ -232,7 +232,7 @@ class Selector:
                     selected.add(model_name)
                     if node.args.get("up"):
                         for u in self._dag.upstream(model_name):
-                            if u in models:
+                            if u in all_models:
                                 selected.add(u)
                     if node.args.get("down"):
                         selected.update(self._dag.downstream(model_name))

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh/core/state_sync/db/facade.py

@@ -256,7 +256,7 @@ class EngineAdapterStateSync(StateSync):
     def unpause_snapshots(
         self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
     ) -> None:
-        self.snapshot_state.unpause_snapshots(snapshots, unpaused_dt, self.interval_state)
+        self.snapshot_state.unpause_snapshots(snapshots, unpaused_dt)
 
     def invalidate_environment(self, name: str, protect_prod: bool = True) -> None:
         self.environment_state.invalidate_environment(name, protect_prod)

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh/core/state_sync/db/migrator.py

@@ -396,7 +396,7 @@ class StateMigrator:
         if updated_prod_environment:
             try:
                 self.snapshot_state.unpause_snapshots(
-                    updated_prod_environment.snapshots, now_timestamp(), self.interval_state
+                    updated_prod_environment.snapshots, now_timestamp()
                 )
             except Exception:
                 logger.warning("Failed to unpause migrated snapshots", exc_info=True)

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh/core/state_sync/db/snapshot.py

@@ -10,6 +10,7 @@ from pydantic import Field
 
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.state_sync.db.utils import (
+    snapshot_name_filter,
     snapshot_name_version_filter,
     snapshot_id_filter,
     fetchone,
@@ -32,15 +33,13 @@ from sqlmesh.core.snapshot import (
     SnapshotChangeCategory,
 )
 from sqlmesh.utils.migration import index_text_type, blob_text_type
-from sqlmesh.utils.date import now_timestamp, TimeLike, now, to_timestamp
+from sqlmesh.utils.date import now_timestamp, TimeLike, to_timestamp
 from sqlmesh.utils.pydantic import PydanticModel
 from sqlmesh.utils import unique
 
 if t.TYPE_CHECKING:
     import pandas as pd
 
-    from sqlmesh.core.state_sync.db.interval import IntervalState
-
 
 logger = logging.getLogger(__name__)
 
@@ -70,6 +69,7 @@ class SnapshotState:
             "unpaused_ts": exp.DataType.build("bigint"),
             "ttl_ms": exp.DataType.build("bigint"),
             "unrestorable": exp.DataType.build("boolean"),
+            "forward_only": exp.DataType.build("boolean"),
         }
 
         self._auto_restatement_columns_to_types = {
@@ -112,84 +112,52 @@ class SnapshotState:
         self,
         snapshots: t.Collection[SnapshotInfoLike],
         unpaused_dt: TimeLike,
-        interval_state: IntervalState,
     ) -> None:
-        """Unpauses given snapshots while pausing all other snapshots that share the same version.
-
-        Args:
-            snapshots: The snapshots to unpause.
-            unpaused_dt: The timestamp to unpause the snapshots at.
-            interval_state: The interval state to use to remove intervals when needed.
-        """
-        current_ts = now()
-
-        target_snapshot_ids = {s.snapshot_id for s in snapshots}
-        same_version_snapshots = self._get_snapshots_with_same_version(
-            snapshots, lock_for_update=True
+        unrestorable_snapshots_by_forward_only: t.Dict[bool, t.List[SnapshotNameVersion]] = (
+            defaultdict(list)
         )
-        target_snapshots_by_version = {
-            (s.name, s.version): s
-            for s in same_version_snapshots
-            if s.snapshot_id in target_snapshot_ids
-        }
-
-        unpaused_snapshots: t.Dict[int, t.List[SnapshotId]] = defaultdict(list)
-        paused_snapshots: t.List[SnapshotId] = []
-        unrestorable_snapshots: t.List[SnapshotId] = []
-
-        for snapshot in same_version_snapshots:
-            is_target_snapshot = snapshot.snapshot_id in target_snapshot_ids
-            if is_target_snapshot and not snapshot.unpaused_ts:
-                logger.info("Unpausing snapshot %s", snapshot.snapshot_id)
-                snapshot.set_unpaused_ts(unpaused_dt)
-                assert snapshot.unpaused_ts is not None
-                unpaused_snapshots[snapshot.unpaused_ts].append(snapshot.snapshot_id)
-            elif not is_target_snapshot:
-                target_snapshot = target_snapshots_by_version[(snapshot.name, snapshot.version)]
-                if (
-                    target_snapshot.normalized_effective_from_ts
-                    and not target_snapshot.disable_restatement
-                ):
-                    # Making sure that there are no overlapping intervals.
-                    effective_from_ts = target_snapshot.normalized_effective_from_ts
-                    logger.info(
-                        "Removing all intervals after '%s' for snapshot %s, superseded by snapshot %s",
-                        target_snapshot.effective_from,
-                        snapshot.snapshot_id,
-                        target_snapshot.snapshot_id,
-                    )
-                    full_snapshot = snapshot.full_snapshot
-                    interval_state.remove_intervals(
-                        [
-                            (
-                                full_snapshot,
-                                full_snapshot.get_removal_interval(effective_from_ts, current_ts),
-                            )
-                        ]
-                    )
 
-                if snapshot.unpaused_ts:
-                    logger.info("Pausing snapshot %s", snapshot.snapshot_id)
-                    snapshot.set_unpaused_ts(None)
-                    paused_snapshots.append(snapshot.snapshot_id)
+        for snapshot in snapshots:
+            # We need to mark all other snapshots that have forward-only opposite to the target snapshot as unrestorable
+            unrestorable_snapshots_by_forward_only[not snapshot.is_forward_only].append(
+                snapshot.name_version
+            )
 
-                if not snapshot.unrestorable and (
-                    (target_snapshot.is_forward_only and not snapshot.is_forward_only)
-                    or (snapshot.is_forward_only and not target_snapshot.is_forward_only)
-                ):
-                    logger.info("Marking snapshot %s as unrestorable", snapshot.snapshot_id)
-                    snapshot.unrestorable = True
-                    unrestorable_snapshots.append(snapshot.snapshot_id)
+        updated_ts = now_timestamp()
+        unpaused_ts = to_timestamp(unpaused_dt)
 
-        if unpaused_snapshots:
-            for unpaused_ts, snapshot_ids in unpaused_snapshots.items():
-                self._update_snapshots(snapshot_ids, unpaused_ts=unpaused_ts)
+        # Pause all snapshots with target names first
+        for where in snapshot_name_filter(
+            [s.name for s in snapshots],
+            batch_size=self.SNAPSHOT_BATCH_SIZE,
+        ):
+            self.engine_adapter.update_table(
+                self.snapshots_table,
+                {"unpaused_ts": None, "updated_ts": updated_ts},
+                where=where,
+            )
 
-        if paused_snapshots:
-            self._update_snapshots(paused_snapshots, unpaused_ts=None)
+        # Now unpause the target snapshots
+        self._update_snapshots(
+            [s.snapshot_id for s in snapshots],
+            unpaused_ts=unpaused_ts,
+            updated_ts=updated_ts,
+        )
 
-        if unrestorable_snapshots:
-            self._update_snapshots(unrestorable_snapshots, unrestorable=True)
+        # Mark unrestorable snapshots
+        for forward_only, snapshot_name_versions in unrestorable_snapshots_by_forward_only.items():
+            forward_only_exp = exp.column("forward_only").is_(exp.convert(forward_only))
+            for where in snapshot_name_version_filter(
+                self.engine_adapter,
+                snapshot_name_versions,
+                batch_size=self.SNAPSHOT_BATCH_SIZE,
+                alias=None,
+            ):
+                self.engine_adapter.update_table(
+                    self.snapshots_table,
+                    {"unrestorable": True, "updated_ts": updated_ts},
+                    where=forward_only_exp.and_(where),
+                )
 
     def get_expired_snapshots(
         self,
@@ -414,7 +382,8 @@ class SnapshotState:
         **kwargs: t.Any,
     ) -> None:
         properties = kwargs
-        properties["updated_ts"] = now_timestamp()
+        if "updated_ts" not in properties:
+            properties["updated_ts"] = now_timestamp()
 
         for where in snapshot_id_filter(
             self.engine_adapter, snapshots, batch_size=self.SNAPSHOT_BATCH_SIZE
@@ -466,6 +435,7 @@ class SnapshotState:
                     updated_ts,
                     unpaused_ts,
                     unrestorable,
+                    forward_only,
                     next_auto_restatement_ts,
                 ) in fetchall(self.engine_adapter, query):
                     snapshot = parse_snapshot(
@@ -473,6 +443,7 @@ class SnapshotState:
                         updated_ts=updated_ts,
                         unpaused_ts=unpaused_ts,
                         unrestorable=unrestorable,
+                        forward_only=forward_only,
                         next_auto_restatement_ts=next_auto_restatement_ts,
                     )
                     snapshot_id = snapshot.snapshot_id
@@ -502,6 +473,7 @@ class SnapshotState:
                         "updated_ts",
                         "unpaused_ts",
                         "unrestorable",
+                        "forward_only",
                         "next_auto_restatement_ts",
                     )
                     .from_(exp.to_table(self.snapshots_table).as_("snapshots"))
@@ -528,6 +500,7 @@ class SnapshotState:
                     updated_ts,
                     unpaused_ts,
                     unrestorable,
+                    forward_only,
                     next_auto_restatement_ts,
                 ) in fetchall(self.engine_adapter, query):
                     snapshot_id = SnapshotId(name=name, identifier=identifier)
@@ -535,6 +508,7 @@ class SnapshotState:
                     snapshot.updated_ts = updated_ts
                     snapshot.unpaused_ts = unpaused_ts
                     snapshot.unrestorable = unrestorable
+                    snapshot.forward_only = forward_only
                     snapshot.next_auto_restatement_ts = next_auto_restatement_ts
                     cached_snapshots_in_state.add(snapshot_id)
 
@@ -568,6 +542,7 @@ class SnapshotState:
                     "snapshots.updated_ts",
                     "snapshots.unpaused_ts",
                     "snapshots.unrestorable",
+                    "snapshots.forward_only",
                     "auto_restatements.next_auto_restatement_ts",
                 )
                 .from_(exp.to_table(self.snapshots_table).as_("snapshots"))
@@ -623,6 +598,7 @@ class SnapshotState:
                     "updated_ts",
                     "unpaused_ts",
                     "unrestorable",
+                    "forward_only",
                 )
                 .from_(exp.to_table(self.snapshots_table).as_("snapshots"))
                 .where(where)
@@ -640,9 +616,10 @@ class SnapshotState:
                 updated_ts=updated_ts,
                 unpaused_ts=unpaused_ts,
                 unrestorable=unrestorable,
+                forward_only=forward_only,
                 snapshot=snapshot,
             )
-            for snapshot, name, identifier, version, updated_ts, unpaused_ts, unrestorable in snapshot_rows
+            for snapshot, name, identifier, version, updated_ts, unpaused_ts, unrestorable, forward_only in snapshot_rows
         ]
 
 
@@ -651,6 +628,7 @@ def parse_snapshot(
     updated_ts: int,
     unpaused_ts: t.Optional[int],
     unrestorable: bool,
+    forward_only: bool,
     next_auto_restatement_ts: t.Optional[int],
 ) -> Snapshot:
     return Snapshot(
@@ -659,6 +637,7 @@ def parse_snapshot(
             "updated_ts": updated_ts,
             "unpaused_ts": unpaused_ts,
             "unrestorable": unrestorable,
+            "forward_only": forward_only,
             "next_auto_restatement_ts": next_auto_restatement_ts,
         }
     )
@@ -673,6 +652,7 @@ def _snapshot_to_json(snapshot: Snapshot) -> str:
             "updated_ts",
             "unpaused_ts",
             "unrestorable",
+            "forward_only",
             "next_auto_restatement_ts",
         }
     )
@@ -693,6 +673,7 @@ def _snapshots_to_df(snapshots: t.Iterable[Snapshot]) -> pd.DataFrame:
                 "unpaused_ts": snapshot.unpaused_ts,
                 "ttl_ms": snapshot.ttl_ms,
                 "unrestorable": snapshot.unrestorable,
+                "forward_only": snapshot.forward_only,
             }
             for snapshot in snapshots
         ]
@@ -762,19 +743,10 @@ class SharedVersionSnapshot(PydanticModel):
                 "updated_ts": self.updated_ts,
                 "unpaused_ts": self.unpaused_ts,
                 "unrestorable": self.unrestorable,
+                "forward_only": self.forward_only,
             }
         )
 
-    def set_unpaused_ts(self, unpaused_dt: t.Optional[TimeLike]) -> None:
-        """Sets the timestamp for when this snapshot was unpaused.
-
-        Args:
-            unpaused_dt: The datetime object of when this snapshot was unpaused.
-        """
-        self.unpaused_ts = (
-            to_timestamp(self.interval_unit.cron_floor(unpaused_dt)) if unpaused_dt else None
-        )
-
     @classmethod
     def from_snapshot_record(
         cls,
@@ -785,6 +757,7 @@ class SharedVersionSnapshot(PydanticModel):
         updated_ts: int,
         unpaused_ts: t.Optional[int],
         unrestorable: bool,
+        forward_only: bool,
         snapshot: str,
     ) -> SharedVersionSnapshot:
         raw_snapshot = json.loads(snapshot)
@@ -803,5 +776,5 @@ class SharedVersionSnapshot(PydanticModel):
             disable_restatement=raw_node.get("kind", {}).get("disable_restatement", False),
             effective_from=raw_snapshot.get("effective_from"),
             raw_snapshot=raw_snapshot,
-            forward_only=raw_snapshot.get("forward_only", False),
+            forward_only=forward_only,
         )

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh/core/state_sync/db/utils.py

@@ -22,6 +22,21 @@ except ImportError:
 T = t.TypeVar("T")
 
 
+def snapshot_name_filter(
+    snapshot_names: t.Iterable[str],
+    batch_size: int,
+    alias: t.Optional[str] = None,
+) -> t.Iterator[exp.Condition]:
+    names = sorted(snapshot_names)
+
+    if not names:
+        yield exp.false()
+    else:
+        batches = create_batches(names, batch_size=batch_size)
+        for names in batches:
+            yield exp.column("name", table=alias).isin(*names)
+
+
 def snapshot_id_filter(
     engine_adapter: EngineAdapter,
     snapshot_ids: t.Iterable[SnapshotIdLike],

sqlmesh-0.209.1.dev34/sqlmesh/migrations/v0090_add_forward_only_column.py

@@ -0,0 +1,100 @@
+"""Add forward_only column to the snapshots table."""
+
+import json
+
+from sqlglot import exp
+
+from sqlmesh.utils.migration import index_text_type, blob_text_type
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    import pandas as pd
+
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    alter_table_exp = exp.Alter(
+        this=exp.to_table(snapshots_table),
+        kind="TABLE",
+        actions=[
+            exp.ColumnDef(
+                this=exp.to_column("forward_only"),
+                kind=exp.DataType.build("boolean"),
+            )
+        ],
+    )
+    engine_adapter.execute(alter_table_exp)
+
+    new_snapshots = []
+
+    for (
+        name,
+        identifier,
+        version,
+        snapshot,
+        kind_name,
+        updated_ts,
+        unpaused_ts,
+        ttl_ms,
+        unrestorable,
+        forward_only,
+    ) in engine_adapter.fetchall(
+        exp.select(
+            "name",
+            "identifier",
+            "version",
+            "snapshot",
+            "kind_name",
+            "updated_ts",
+            "unpaused_ts",
+            "ttl_ms",
+            "unrestorable",
+            "forward_only",
+        ).from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+
+        forward_only = parsed_snapshot.get("forward_only")
+        if forward_only is None:
+            forward_only = parsed_snapshot.get("change_category") == 3
+
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": json.dumps(parsed_snapshot),
+                "kind_name": kind_name,
+                "updated_ts": updated_ts,
+                "unpaused_ts": unpaused_ts,
+                "ttl_ms": ttl_ms,
+                "unrestorable": unrestorable,
+                "forward_only": forward_only,
+            }
+        )
+
+    if new_snapshots:
+        engine_adapter.delete_from(snapshots_table, "TRUE")
+        index_type = index_text_type(engine_adapter.dialect)
+        blob_type = blob_text_type(engine_adapter.dialect)
+
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types={
+                "name": exp.DataType.build(index_type),
+                "identifier": exp.DataType.build(index_type),
+                "version": exp.DataType.build(index_type),
+                "snapshot": exp.DataType.build(blob_type),
+                "kind_name": exp.DataType.build(index_type),
+                "updated_ts": exp.DataType.build("bigint"),
+                "unpaused_ts": exp.DataType.build("bigint"),
+                "ttl_ms": exp.DataType.build("bigint"),
+                "unrestorable": exp.DataType.build("boolean"),
+                "forward_only": exp.DataType.build("boolean"),
+            },
+        )

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.209.1.dev32
+Version: 0.209.1.dev34
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.209.1.dev32 → sqlmesh-0.209.1.dev34}/sqlmesh.egg-info/SOURCES.txt

@@ -469,6 +469,7 @@ sqlmesh/migrations/v0086_check_deterministic_bug.py
 sqlmesh/migrations/v0087_normalize_blueprint_variables.py
 sqlmesh/migrations/v0088_warn_about_variable_python_env_diffs.py
 sqlmesh/migrations/v0089_add_virtual_environment_mode.py
+sqlmesh/migrations/v0090_add_forward_only_column.py
 sqlmesh/schedulers/__init__.py
 sqlmesh/utils/__init__.py
 sqlmesh/utils/aws.py