sqlmesh 0.213.1.dev1__py3-none-any.whl → 0.227.2.dev4__py3-none-any.whl

sqlmesh/core/state_sync/db/facade.py

@@ -22,19 +22,19 @@ import typing as t
 from pathlib import Path
 from datetime import datetime
 
-from sqlglot import exp
 
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentSummary
 from sqlmesh.core.snapshot import (
     Snapshot,
+    SnapshotIdAndVersion,
     SnapshotId,
     SnapshotIdLike,
+    SnapshotIdAndVersionLike,
     SnapshotInfoLike,
     SnapshotIntervals,
     SnapshotNameVersion,
-    SnapshotTableCleanupTask,
     SnapshotTableInfo,
     start_date,
 )
@@ -42,7 +42,6 @@ from sqlmesh.core.snapshot.definition import (
     Interval,
 )
 from sqlmesh.core.state_sync.base import (
-    PromotionResult,
     StateSync,
     Versions,
 )
@@ -54,6 +53,9 @@ from sqlmesh.core.state_sync.common import (
     StateStream,
     chunk_iterable,
     EnvironmentWithStatements,
+    ExpiredSnapshotBatch,
+    PromotionResult,
+    ExpiredBatchRange,
 )
 from sqlmesh.core.state_sync.db.interval import IntervalState
 from sqlmesh.core.state_sync.db.environment import EnvironmentState
@@ -89,7 +91,6 @@ class EngineAdapterStateSync(StateSync):
         console: t.Optional[Console] = None,
         cache_dir: Path = Path(),
     ):
-        self.plan_dags_table = exp.table_("_plan_dags", db=schema)
         self.interval_state = IntervalState(engine_adapter, schema=schema)
         self.environment_state = EnvironmentState(engine_adapter, schema=schema)
         self.snapshot_state = SnapshotState(engine_adapter, schema=schema, cache_dir=cache_dir)
@@ -100,7 +101,6 @@ class EngineAdapterStateSync(StateSync):
             snapshot_state=self.snapshot_state,
             environment_state=self.environment_state,
             interval_state=self.interval_state,
-            plan_dags_table=self.plan_dags_table,
             console=console,
         )
         # Make sure that if an empty string is provided that we treat it as None
@@ -262,10 +262,18 @@ class EngineAdapterStateSync(StateSync):
         self.environment_state.invalidate_environment(name, protect_prod)
 
     def get_expired_snapshots(
-        self, current_ts: int, ignore_ttl: bool = False
-    ) -> t.List[SnapshotTableCleanupTask]:
+        self,
+        *,
+        batch_range: ExpiredBatchRange,
+        current_ts: t.Optional[int] = None,
+        ignore_ttl: bool = False,
+    ) -> t.Optional[ExpiredSnapshotBatch]:
+        current_ts = current_ts or now_timestamp()
         return self.snapshot_state.get_expired_snapshots(
-            self.environment_state.get_environments(), current_ts=current_ts, ignore_ttl=ignore_ttl
+            environments=self.environment_state.get_environments(),
+            current_ts=current_ts,
+            ignore_ttl=ignore_ttl,
+            batch_range=batch_range,
         )
 
     def get_expired_environments(self, current_ts: int) -> t.List[EnvironmentSummary]:
@@ -273,17 +281,19 @@ class EngineAdapterStateSync(StateSync):
 
     @transactional()
     def delete_expired_snapshots(
-        self, ignore_ttl: bool = False, current_ts: t.Optional[int] = None
-    ) -> t.List[SnapshotTableCleanupTask]:
-        current_ts = current_ts or now_timestamp()
-        expired_snapshot_ids, cleanup_targets = self.snapshot_state._get_expired_snapshots(
-            self.environment_state.get_environments(), ignore_ttl=ignore_ttl, current_ts=current_ts
+        self,
+        batch_range: ExpiredBatchRange,
+        ignore_ttl: bool = False,
+        current_ts: t.Optional[int] = None,
+    ) -> None:
+        batch = self.get_expired_snapshots(
+            ignore_ttl=ignore_ttl,
+            current_ts=current_ts,
+            batch_range=batch_range,
         )
-
-        self.snapshot_state.delete_snapshots(expired_snapshot_ids)
-        self.interval_state.cleanup_intervals(cleanup_targets, expired_snapshot_ids)
-
-        return cleanup_targets
+        if batch and batch.expired_snapshot_ids:
+            self.snapshot_state.delete_snapshots(batch.expired_snapshot_ids)
+            self.interval_state.cleanup_intervals(batch.cleanup_tasks, batch.expired_snapshot_ids)
 
     @transactional()
     def delete_expired_environments(
@@ -309,7 +319,6 @@ class EngineAdapterStateSync(StateSync):
             self.environment_state.environments_table,
             self.environment_state.environment_statements_table,
             self.interval_state.intervals_table,
-            self.plan_dags_table,
             self.version_state.versions_table,
         ):
             self.engine_adapter.drop_table(table)
@@ -368,6 +377,16 @@ class EngineAdapterStateSync(StateSync):
         Snapshot.hydrate_with_intervals_by_version(snapshots.values(), intervals)
         return snapshots
 
+    def get_snapshots_by_names(
+        self,
+        snapshot_names: t.Iterable[str],
+        current_ts: t.Optional[int] = None,
+        exclude_expired: bool = True,
+    ) -> t.Set[SnapshotIdAndVersion]:
+        return self.snapshot_state.get_snapshots_by_names(
+            snapshot_names=snapshot_names, current_ts=current_ts, exclude_expired=exclude_expired
+        )
+
     @transactional()
     def add_interval(
         self,
@@ -375,8 +394,9 @@ class EngineAdapterStateSync(StateSync):
         start: TimeLike,
         end: TimeLike,
         is_dev: bool = False,
+        last_altered_ts: t.Optional[int] = None,
     ) -> None:
-        super().add_interval(snapshot, start, end, is_dev)
+        super().add_interval(snapshot, start, end, is_dev, last_altered_ts)
 
     @transactional()
     def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
@@ -402,7 +422,7 @@ class EngineAdapterStateSync(StateSync):
     @transactional()
     def remove_intervals(
         self,
-        snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
         remove_shared_versions: bool = False,
     ) -> None:
         self.interval_state.remove_intervals(snapshot_intervals, remove_shared_versions)
@@ -444,14 +464,12 @@ class EngineAdapterStateSync(StateSync):
     @transactional()
     def migrate(
         self,
-        default_catalog: t.Optional[str],
         skip_backup: bool = False,
         promoted_snapshots_only: bool = True,
     ) -> None:
         """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
         self.migrator.migrate(
-            self,
-            default_catalog,
+            self.schema,
             skip_backup=skip_backup,
             promoted_snapshots_only=promoted_snapshots_only,
         )

sqlmesh/core/state_sync/db/interval.py

@@ -15,10 +15,10 @@ from sqlmesh.core.state_sync.db.utils import (
 from sqlmesh.core.snapshot import (
     SnapshotIntervals,
     SnapshotIdLike,
+    SnapshotIdAndVersionLike,
     SnapshotNameVersionLike,
     SnapshotTableCleanupTask,
     SnapshotNameVersion,
-    SnapshotInfoLike,
     Snapshot,
 )
 from sqlmesh.core.snapshot.definition import Interval
@@ -60,6 +60,7 @@ class IntervalState:
             "is_removed": exp.DataType.build("boolean"),
             "is_compacted": exp.DataType.build("boolean"),
             "is_pending_restatement": exp.DataType.build("boolean"),
+            "last_altered_ts": exp.DataType.build("bigint"),
         }
 
     def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
@@ -68,11 +69,11 @@ class IntervalState:
 
     def remove_intervals(
         self,
-        snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
         remove_shared_versions: bool = False,
     ) -> None:
         intervals_to_remove: t.Sequence[
-            t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]
+            t.Tuple[t.Union[SnapshotIdAndVersionLike, SnapshotIntervals], Interval]
         ] = snapshot_intervals
         if remove_shared_versions:
             name_version_mapping = {s.name_version: interval for s, interval in snapshot_intervals}
@@ -215,13 +216,23 @@ class IntervalState:
             for start_ts, end_ts in snapshot.intervals:
                 new_intervals.append(
                     _interval_to_df(
-                        snapshot, start_ts, end_ts, is_dev=False, is_compacted=is_compacted
+                        snapshot,
+                        start_ts,
+                        end_ts,
+                        is_dev=False,
+                        is_compacted=is_compacted,
+                        last_altered_ts=snapshot.last_altered_ts,
                     )
                 )
             for start_ts, end_ts in snapshot.dev_intervals:
                 new_intervals.append(
                     _interval_to_df(
-                        snapshot, start_ts, end_ts, is_dev=True, is_compacted=is_compacted
+                        snapshot,
+                        start_ts,
+                        end_ts,
+                        is_dev=True,
+                        is_compacted=is_compacted,
+                        last_altered_ts=snapshot.dev_last_altered_ts,
                     )
                 )
 
@@ -236,6 +247,7 @@ class IntervalState:
                         is_dev=False,
                         is_compacted=is_compacted,
                         is_pending_restatement=True,
+                        last_altered_ts=snapshot.last_altered_ts,
                     )
                 )
 
@@ -284,6 +296,7 @@ class IntervalState:
                 is_dev,
                 is_removed,
                 is_pending_restatement,
+                last_altered_ts,
             ) in rows:
                 interval_ids.add(interval_id)
                 merge_key = (name, version, dev_version, identifier)
@@ -318,8 +331,10 @@ class IntervalState:
                 else:
                     if is_dev:
                         intervals[merge_key].add_dev_interval(start, end)
+                        intervals[merge_key].update_dev_last_altered_ts(last_altered_ts)
                     else:
                         intervals[merge_key].add_interval(start, end)
+                        intervals[merge_key].update_last_altered_ts(last_altered_ts)
                         # Remove all pending restatement intervals recorded before the current interval has been added
                         intervals[
                             pending_restatement_interval_merge_key
@@ -340,6 +355,7 @@ class IntervalState:
                 "is_dev",
                 "is_removed",
                 "is_pending_restatement",
+                "last_altered_ts",
             )
             .from_(exp.to_table(self.intervals_table).as_("intervals"))
             .order_by(
@@ -431,7 +447,9 @@ class IntervalState:
 
 
 def _intervals_to_df(
-    snapshot_intervals: t.Sequence[t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]],
+    snapshot_intervals: t.Sequence[
+        t.Tuple[t.Union[SnapshotIdAndVersionLike, SnapshotIntervals], Interval]
+    ],
     is_dev: bool,
     is_removed: bool,
 ) -> pd.DataFrame:
@@ -451,13 +469,14 @@ def _intervals_to_df(
 
 
 def _interval_to_df(
-    snapshot: t.Union[SnapshotInfoLike, SnapshotIntervals],
+    snapshot: t.Union[SnapshotIdAndVersionLike, SnapshotIntervals],
     start_ts: int,
     end_ts: int,
     is_dev: bool = False,
     is_removed: bool = False,
     is_compacted: bool = False,
     is_pending_restatement: bool = False,
+    last_altered_ts: t.Optional[int] = None,
 ) -> t.Dict[str, t.Any]:
     return {
         "id": random_id(),
@@ -472,4 +491,5 @@ def _interval_to_df(
         "is_removed": is_removed,
         "is_compacted": is_compacted,
         "is_pending_restatement": is_pending_restatement,
+        "last_altered_ts": last_altered_ts,
     }

sqlmesh/core/state_sync/db/migrator.py

@@ -27,8 +27,9 @@ from sqlmesh.core.snapshot.definition import (
 )
 from sqlmesh.core.state_sync.base import (
     MIGRATIONS,
+    MIN_SCHEMA_VERSION,
+    MIN_SQLMESH_VERSION,
 )
-from sqlmesh.core.state_sync.base import StateSync
 from sqlmesh.core.state_sync.db.environment import EnvironmentState
 from sqlmesh.core.state_sync.db.interval import IntervalState
 from sqlmesh.core.state_sync.db.snapshot import SnapshotState
@@ -41,7 +42,7 @@ from sqlmesh.core.state_sync.db.utils import (
 from sqlmesh.utils import major_minor
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import now_timestamp
-from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.errors import SQLMeshError, StateMigrationError
 
 logger = logging.getLogger(__name__)
 
@@ -61,7 +62,6 @@ class StateMigrator:
         snapshot_state: SnapshotState,
         environment_state: EnvironmentState,
         interval_state: IntervalState,
-        plan_dags_table: TableName,
         console: t.Optional[Console] = None,
     ):
         self.engine_adapter = engine_adapter
@@ -70,7 +70,6 @@ class StateMigrator:
         self.snapshot_state = snapshot_state
         self.environment_state = environment_state
         self.interval_state = interval_state
-        self.plan_dags_table = plan_dags_table
 
         self._state_tables = [
             self.snapshot_state.snapshots_table,
@@ -79,15 +78,13 @@ class StateMigrator:
         ]
         self._optional_state_tables = [
             self.interval_state.intervals_table,
-            self.plan_dags_table,
             self.snapshot_state.auto_restatements_table,
             self.environment_state.environment_statements_table,
         ]
 
     def migrate(
         self,
-        state_sync: StateSync,
-        default_catalog: t.Optional[str],
+        schema: t.Optional[str],
         skip_backup: bool = False,
         promoted_snapshots_only: bool = True,
     ) -> None:
@@ -96,15 +93,13 @@ class StateMigrator:
         migration_start_ts = time.perf_counter()
 
         try:
-            migrate_rows = self._apply_migrations(state_sync, default_catalog, skip_backup)
+            migrate_rows = self._apply_migrations(schema, skip_backup)
 
             if not migrate_rows and major_minor(SQLMESH_VERSION) == versions.minor_sqlmesh_version:
                 return
 
             if migrate_rows:
                 self._migrate_rows(promoted_snapshots_only)
-                # Cleanup plan DAGs since we currently don't migrate snapshot records that are in there.
-                self.engine_adapter.delete_from(self.plan_dags_table, "TRUE")
             self.version_state.update_versions()
 
             analytics.collector.on_migration_end(
@@ -126,6 +121,8 @@ class StateMigrator:
             )
 
             self.console.log_migration_status(success=False)
+            if isinstance(e, StateMigrationError):
+                raise
             raise SQLMeshError("SQLMesh migration failed.") from e
 
         self.console.log_migration_status()
@@ -155,12 +152,21 @@ class StateMigrator:
 
     def _apply_migrations(
         self,
-        state_sync: StateSync,
-        default_catalog: t.Optional[str],
+        schema: t.Optional[str],
         skip_backup: bool,
     ) -> bool:
         versions = self.version_state.get_versions()
-        migrations = MIGRATIONS[versions.schema_version :]
+        first_script_index = 0
+        if versions.schema_version and versions.schema_version < MIN_SCHEMA_VERSION:
+            raise StateMigrationError(
+                "The current state belongs to an old version of SQLMesh that is no longer supported. "
+                f"Please upgrade to {MIN_SQLMESH_VERSION} first before upgrading to {SQLMESH_VERSION}."
+            )
+        elif versions.schema_version > 0:
+            # -1 to skip the baseline migration script
+            first_script_index = versions.schema_version - (MIN_SCHEMA_VERSION - 1)
+
+        migrations = MIGRATIONS[first_script_index:]
         should_backup = any(
             [
                 migrations,
@@ -173,9 +179,14 @@ class StateMigrator:
 
         snapshot_count_before = self.snapshot_state.count() if versions.schema_version else None
 
+        state_table_exist = any(self.engine_adapter.table_exists(t) for t in self._state_tables)
+
         for migration in migrations:
             logger.info(f"Applying migration {migration}")
-            migration.migrate(state_sync, default_catalog=default_catalog)
+            migration.migrate_schemas(engine_adapter=self.engine_adapter, schema=schema)
+            if state_table_exist:
+                # No need to run DML for the initial migration since all tables are empty
+                migration.migrate_rows(engine_adapter=self.engine_adapter, schema=schema)
 
         snapshot_count_after = self.snapshot_state.count()
 
@@ -217,6 +228,7 @@ class StateMigrator:
                 "updated_ts": updated_ts,
                 "unpaused_ts": unpaused_ts,
                 "unrestorable": unrestorable,
+                "forward_only": forward_only,
             }
             for where in (
                 snapshot_id_filter(
@@ -225,10 +237,16 @@ class StateMigrator:
                 if snapshots is not None
                 else [None]
             )
-            for name, identifier, raw_snapshot, updated_ts, unpaused_ts, unrestorable in fetchall(
+            for name, identifier, raw_snapshot, updated_ts, unpaused_ts, unrestorable, forward_only in fetchall(
                 self.engine_adapter,
                 exp.select(
-                    "name", "identifier", "snapshot", "updated_ts", "unpaused_ts", "unrestorable"
+                    "name",
+                    "identifier",
+                    "snapshot",
+                    "updated_ts",
+                    "unpaused_ts",
+                    "unrestorable",
+                    "forward_only",
                 )
                 .from_(self.snapshot_state.snapshots_table)
                 .where(where)