Repository: TobikoData/sqlmesh
Pull request #2489 (2024-04-22 21:37:51)

Fix: When removing intervals also remove them for snapshots that have been deleted by the janitor process
When the janitor deletes obsolete snapshots, intervals associated with them remain, since the deleted snapshot may have been sharing its version with other snapshots.

When we delete intervals for restatement purposes, we have to delete them for **all** unique (name, identifier) pairs that are associated with the same version, because the interval hydration happens on the version basis.

Before this update, we used to source snapshots with the shared version from the `_snapshots` table. This led to incomplete removal of intervals since only intervals for existing snapshots were considered, not for the garbage collected ones.

With this update, snapshots with shared versions are sourced directly from the `_intervals` table which is the source of truth in this case.


Patch:
diff --git a/sqlmesh/core/engine_adapter/base_postgres.py b/sqlmesh/core/engine_adapter/base_postgres.py
index 4ca5627f..9ad49c4f 100644
--- a/sqlmesh/core/engine_adapter/base_postgres.py
+++ b/sqlmesh/core/engine_adapter/base_postgres.py
@@ -22,25 +22,37 @@ if t.TYPE_CHECKING:
 
 class BasePostgresEngineAdapter(EngineAdapter):
     DEFAULT_BATCH_SIZE = 400
-    COLUMNS_TABLE = "information_schema.columns"
     CATALOG_SUPPORT = CatalogSupport.SINGLE_CATALOG_ONLY
     COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
     COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 
+    def _columns_query(self, table: exp.Table) -> exp.Select:
+        sql = (
+            exp.select(
+                "attname AS column_name",
+                "pg_catalog.format_type(atttypid, atttypmod) AS data_type",
+            )
+            .from_("pg_catalog.pg_attribute")
+            .join("pg_catalog.pg_class", on="oid = attrelid")
+            .join("pg_catalog.pg_namespace", on="oid = relnamespace")
+            .where(
+                exp.and_(
+                    "attnum > 0",
+                    "NOT attisdropped",
+                    exp.column("nspname").eq(table.alias_or_name),
+                )
+            )
+        )
+        if table.args.get("db"):
+            sql = sql.where(exp.column("relname").eq(table.args["db"].name))
+        return sql
+
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
     ) -> t.Dict[str, exp.DataType]:
         """Fetches column names and types for the target table."""
         table = exp.to_table(table_name)
-        sql = (
-            exp.select("column_name", "data_type")
-            .from_(self.COLUMNS_TABLE)
-            .where(f"table_name = '{table.alias_or_name}'")
-        )
-        if table.args.get("db"):
-            sql = sql.where(f"table_schema = '{table.args['db'].name}'")
-
-        self.execute(sql)
+        self.execute(self._columns_query(table))
         resp = self.cursor.fetchall()
         if not resp:
             raise SQLMeshError("Could not get columns for table '%s'. Table not found.", table_name)
diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py
index 981f9c78..d8bb59c5 100644
--- a/sqlmesh/core/engine_adapter/redshift.py
+++ b/sqlmesh/core/engine_adapter/redshift.py
@@ -34,12 +34,21 @@ class RedshiftEngineAdapter(
     NonTransactionalTruncateMixin,
 ):
     DIALECT = "redshift"
-    COLUMNS_TABLE = "svv_columns"  # Includes late-binding views
     CURRENT_CATALOG_EXPRESSION = exp.func("current_database")
     # Redshift doesn't support comments for VIEWs WITH NO SCHEMA BINDING (which we always use)
     COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
     SUPPORTS_REPLACE_TABLE = False
 
+    def _columns_query(self, table: exp.Table) -> exp.Select:
+        sql = (
+            exp.select("column_name", "data_type")
+            .from_("svv_columns")  # Includes late-binding views
+            .where(exp.column("table_name").eq(table.alias_or_name))
+        )
+        if table.args.get("db"):
+            sql = sql.where(exp.column("table_schema").eq(table.args["db"].name))
+        return sql
+
     @property
     def cursor(self) -> t.Any:
         # Redshift by default uses a `format` paramstyle that has issues when we try to write our snapshot
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 1ef846cf..e88577b0 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -162,6 +162,10 @@ class SnapshotIntervals(PydanticModel, frozen=True):
     def snapshot_id(self) -> SnapshotId:
         return SnapshotId(name=self.name, identifier=self.identifier)
 
+    @property
+    def name_version(self) -> SnapshotNameVersion:
+        return SnapshotNameVersion(name=self.name, version=self.version)
+
 
 class SnapshotDataVersion(PydanticModel, frozen=True):
     fingerprint: SnapshotFingerprint
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index bd785e96..b9163634 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -645,26 +645,38 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
         execution_time: t.Optional[TimeLike] = None,
         remove_shared_versions: bool = False,
     ) -> None:
+        intervals_to_remove: t.Sequence[
+            t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]
+        ] = snapshot_intervals
         if remove_shared_versions:
-            name_version_mapping = {
-                s.name_version: (s, interval) for s, interval in snapshot_intervals
-            }
-            all_snapshots = self._get_snapshots_with_same_version(
-                [s[0] for s in snapshot_intervals]
-            )
-            snapshot_intervals = [
-                (snapshot, name_version_mapping[snapshot.name_version][1])
+            name_version_mapping = {s.name_version: interval for s, interval in snapshot_intervals}
+            all_snapshots = []
+            for where in self._snapshot_name_version_filter(name_version_mapping, alias=None):
+                all_snapshots.extend(
+                    [
+                        SnapshotIntervals(
+                            name=r[0], identifier=r[1], version=r[2], intervals=[], dev_intervals=[]
+                        )
+                        for r in self._fetchall(
+                            exp.select("name", "identifier", "version")
+                            .from_(self.intervals_table)
+                            .where(where)
+                        )
+                    ]
+                )
+            intervals_to_remove = [
+                (snapshot, name_version_mapping[snapshot.name_version])
                 for snapshot in all_snapshots
             ]
 
         if logger.isEnabledFor(logging.INFO):
-            snapshot_ids = ", ".join(str(s.snapshot_id) for s, _ in snapshot_intervals)
+            snapshot_ids = ", ".join(str(s.snapshot_id) for s, _ in intervals_to_remove)
             logger.info("Removing interval for snapshots: %s", snapshot_ids)
 
         for is_dev in (True, False):
             self.engine_adapter.insert_append(
                 self.intervals_table,
-                _intervals_to_df(snapshot_intervals, is_dev=is_dev, is_removed=True),
+                _intervals_to_df(intervals_to_remove, is_dev=is_dev, is_removed=True),
                 columns_to_types=self._interval_columns_to_types,
             )
 
@@ -917,6 +929,8 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                 migration.migrate(self, default_catalog=default_catalog)
             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._update_versions()
         except Exception as e:
             if skip_backup:
@@ -1260,7 +1274,9 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                 )
 
     def _snapshot_name_version_filter(
-        self, snapshot_name_versions: t.Iterable[SnapshotNameVersionLike], alias: str = "snapshots"
+        self,
+        snapshot_name_versions: t.Iterable[SnapshotNameVersionLike],
+        alias: t.Optional[str] = "snapshots",
     ) -> t.Iterator[exp.Condition]:
         name_versions = sorted({(s.name, s.version) for s in snapshot_name_versions})
         batches = self._snapshot_batches(name_versions)
@@ -1303,7 +1319,7 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
 
 
 def _intervals_to_df(
-    snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+    snapshot_intervals: t.Sequence[t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]],
     is_dev: bool,
     is_removed: bool,
 ) -> pd.DataFrame:
diff --git a/web/server/main.py b/web/server/main.py
index d2c11566..69ce174e 100644
--- a/web/server/main.py
+++ b/web/server/main.py
@@ -2,6 +2,8 @@ import asyncio
 import mimetypes
 import pathlib
 import threading
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse, JSONResponse
@@ -14,20 +16,9 @@ from web.server.exceptions import ApiException
 from web.server.settings import get_context_or_raise, get_settings
 from web.server.watcher import watch_project
 
-app = FastAPI()
 
-app.include_router(api_router, prefix="/api")
-WEB_DIRECTORY = pathlib.Path(__file__).parent.parent
-
-# Starlette uses mimetypes.guess_type to determine a file response's content type. Since this method
-# is not consistent across different computers, operating systems, etc., we enumerate some of the
-# more common ones we use here.
-mimetypes.add_type("application/javascript", ".js")
-mimetypes.add_type("text/css", ".css")
-
-
-@app.on_event("startup")
-async def startup_event() -> None:
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncIterator:
     async def dispatch() -> None:
         while True:
             item = await api_console.queue.get()
@@ -40,15 +31,26 @@ async def startup_event() -> None:
     app.state.watch_task = asyncio.create_task(watch_project())
     app.state.circuit_breaker = threading.Event()
 
+    yield
 
-@app.on_event("shutdown")
-async def shutdown_event() -> None:
     app.state.dispatch_task.cancel()
     app.state.watch_task.cancel()
     context = await get_context_or_raise(settings=get_settings())
     context.close()
 
 
+app = FastAPI(lifespan=lifespan)
+
+app.include_router(api_router, prefix="/api")
+WEB_DIRECTORY = pathlib.Path(__file__).parent.parent
+
+# Starlette uses mimetypes.guess_type to determine a file response's content type. Since this method
+# is not consistent across different computers, operating systems, etc., we enumerate some of the
+# more common ones we use here.
+mimetypes.add_type("application/javascript", ".js")
+mimetypes.add_type("text/css", ".css")
+
+
 @app.exception_handler(ApiException)
 async def handle_api_exception(_: Request, e: ApiException) -> JSONResponse:
     return JSONResponse(status_code=e.status_code, content=e.to_dict())


Test patch:
diff --git a/tests/core/engine_adapter/test_base_postgres.py b/tests/core/engine_adapter/test_base_postgres.py
index a7a9369f..2e8809b1 100644
--- a/tests/core/engine_adapter/test_base_postgres.py
+++ b/tests/core/engine_adapter/test_base_postgres.py
@@ -16,7 +16,12 @@ def test_columns(make_mocked_engine_adapter: t.Callable):
 
     resp = adapter.columns("db.table")
     adapter.cursor.execute.assert_called_once_with(
-        """SELECT "column_name", "data_type" FROM "information_schema"."columns" WHERE "table_name" = 'table' AND "table_schema" = 'db'"""
+        'SELECT "attname" AS "column_name", '
+        '"pg_catalog".FORMAT_TYPE("atttypid", "atttypmod") AS "data_type" '
+        'FROM "pg_catalog"."pg_attribute" '
+        'JOIN "pg_catalog"."pg_class" ON "oid" = "attrelid" '
+        'JOIN "pg_catalog"."pg_namespace" ON "oid" = "relnamespace" '
+        """WHERE ("attnum" > 0 AND NOT "attisdropped" AND "nspname" = 'table') AND "relname" = 'db'"""
     )
     assert resp == {"col": exp.DataType.build("INT")}
 
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index 8187b910..8de2b07d 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -330,6 +330,49 @@ def test_remove_interval(state_sync: EngineAdapterStateSync, make_snapshot: t.Ca
     ]
 
 
+def test_remove_interval_missing_snapshot(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
+) -> None:
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            cron="@daily",
+            query=parse_one("select 1, ds"),
+        ),
+        version="a",
+    )
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="a",
+            cron="@daily",
+            query=parse_one("select 2::INT, '2022-01-01'::TEXT AS ds"),
+        ),
+        version="a",
+    )
+    # Only add snapshot_a to simulate that snapshot_b is missing
+    state_sync.push_snapshots([snapshot_a])
+    state_sync.add_interval(snapshot_a, "2020-01-01", "2020-01-10")
+    state_sync.add_interval(snapshot_b, "2020-01-11", "2020-01-30")
+
+    snapshots = state_sync.get_snapshots([snapshot_a, snapshot_b])
+    assert len(snapshots) == 1
+    assert snapshots[snapshot_a.snapshot_id].intervals == [
+        (to_timestamp("2020-01-01"), to_timestamp("2020-01-31")),
+    ]
+
+    state_sync.remove_interval(
+        [(snapshot_a, snapshot_a.inclusive_exclusive("2020-01-15", "2020-01-17"))],
+        remove_shared_versions=True,
+    )
+
+    snapshots = state_sync.get_snapshots([snapshot_a, snapshot_b])
+    assert len(snapshots) == 1
+    assert snapshots[snapshot_a.snapshot_id].intervals == [
+        (to_timestamp("2020-01-01"), to_timestamp("2020-01-15")),
+        (to_timestamp("2020-01-18"), to_timestamp("2020-01-31")),
+    ]
+
+
 def test_refresh_snapshot_intervals(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
 ) -> None: