sqlmesh.core.janitor

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6
  7from sqlmesh.core.engine_adapter import EngineAdapter
  8from sqlmesh.core.console import Console
  9from sqlmesh.core.dialect import schema_
 10from sqlmesh.core.environment import Environment
 11from sqlmesh.core.snapshot import SnapshotEvaluator
 12from sqlmesh.core.state_sync import StateSync
 13from sqlmesh.core.state_sync.common import (
 14    logger,
 15    iter_expired_snapshot_batches,
 16    RowBoundary,
 17    ExpiredBatchRange,
 18)
 19
 20
 21def cleanup_expired_views(
 22    default_adapter: EngineAdapter,
 23    engine_adapters: t.Dict[str, EngineAdapter],
 24    environments: t.List[Environment],
 25    console: t.Optional[Console] = None,
 26) -> t.List[str]:
 27    failures: t.List[str] = []
 28
 29    expired_schema_or_catalog_environments = [
 30        environment
 31        for environment in environments
 32        if environment.suffix_target.is_schema or environment.suffix_target.is_catalog
 33    ]
 34    expired_table_environments = [
 35        environment for environment in environments if environment.suffix_target.is_table
 36    ]
 37
 38    # We have to use the corresponding adapter if the virtual layer is gateway managed
 39    def get_adapter(gateway_managed: bool, gateway: t.Optional[str] = None) -> EngineAdapter:
 40        if gateway_managed and gateway:
 41            return engine_adapters.get(gateway, default_adapter)
 42        return default_adapter
 43
 44    catalogs_to_drop: t.Set[t.Tuple[EngineAdapter, str]] = set()
 45    schemas_to_drop: t.Set[t.Tuple[EngineAdapter, exp.Table]] = set()
 46
 47    # Collect schemas and catalogs to drop
 48    for engine_adapter, expired_catalog, expired_schema, suffix_target in {
 49        (
 50            (engine_adapter := get_adapter(environment.gateway_managed, snapshot.model_gateway)),
 51            snapshot.qualified_view_name.catalog_for_environment(
 52                environment.naming_info, dialect=engine_adapter.dialect
 53            ),
 54            snapshot.qualified_view_name.schema_for_environment(
 55                environment.naming_info, dialect=engine_adapter.dialect
 56            ),
 57            environment.suffix_target,
 58        )
 59        for environment in expired_schema_or_catalog_environments
 60        for snapshot in environment.snapshots
 61        if snapshot.is_model and not snapshot.is_symbolic
 62    }:
 63        if suffix_target.is_catalog:
 64            if expired_catalog:
 65                catalogs_to_drop.add((engine_adapter, expired_catalog))
 66        else:
 67            schema = schema_(expired_schema, expired_catalog)
 68            schemas_to_drop.add((engine_adapter, schema))
 69
 70    # Drop the views for the expired environments
 71    for engine_adapter, expired_view in {
 72        (
 73            (engine_adapter := get_adapter(environment.gateway_managed, snapshot.model_gateway)),
 74            snapshot.qualified_view_name.for_environment(
 75                environment.naming_info, dialect=engine_adapter.dialect
 76            ),
 77        )
 78        for environment in expired_table_environments
 79        for snapshot in environment.snapshots
 80        if snapshot.is_model and not snapshot.is_symbolic
 81    }:
 82        try:
 83            engine_adapter.drop_view(expired_view, ignore_if_not_exists=True)
 84            if console:
 85                console.update_cleanup_progress(expired_view)
 86        except Exception as e:
 87            message = f"Failed to drop the expired environment view '{expired_view}': {e}"
 88            logger.warning(message)
 89            failures.append(message)
 90
 91    # Drop the schemas for the expired environments
 92    for engine_adapter, schema in schemas_to_drop:
 93        try:
 94            engine_adapter.drop_schema(
 95                schema,
 96                ignore_if_not_exists=True,
 97                cascade=True,
 98            )
 99            if console:
100                console.update_cleanup_progress(schema.sql(dialect=engine_adapter.dialect))
101        except Exception as e:
102            message = f"Failed to drop the expired environment schema '{schema}': {e}"
103            logger.warning(message)
104            failures.append(message)
105
106    # Drop any catalogs that were associated with a snapshot where the engine adapter supports dropping catalogs
107    # catalogs_to_drop is only populated when environment_suffix_target is set to 'catalog'
108    for engine_adapter, catalog in catalogs_to_drop:
109        if engine_adapter.SUPPORTS_CREATE_DROP_CATALOG:
110            try:
111                engine_adapter.drop_catalog(catalog)
112                if console:
113                    console.update_cleanup_progress(catalog)
114            except Exception as e:
115                message = f"Failed to drop the expired environment catalog '{catalog}': {e}"
116                logger.warning(message)
117                failures.append(message)
118
119    return failures
120
121
122def delete_expired_snapshots(
123    state_sync: StateSync,
124    snapshot_evaluator: SnapshotEvaluator,
125    *,
126    current_ts: int,
127    ignore_ttl: bool = False,
128    force_delete: bool = False,
129    batch_size: t.Optional[int] = None,
130    console: t.Optional[Console] = None,
131) -> t.List[str]:
132    """Delete all expired snapshots in batches.
133
134    This helper function encapsulates the logic for deleting expired snapshots in batches,
135    eliminating code duplication across different use cases.
136
137    Args:
138        state_sync: StateSync instance to query and delete expired snapshots from.
139        snapshot_evaluator: SnapshotEvaluator instance to clean up tables associated with snapshots.
140        current_ts: Timestamp used to evaluate expiration.
141        ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
142        force_delete: If True, delete snapshot state records even when physical table cleanup fails.
143        batch_size: Maximum number of snapshots to fetch per batch.
144        console: Optional console for reporting progress.
145
146    Returns:
147        List of failure messages so callers can surface them at the end of the janitor run.
148    """
149    failures: t.List[str] = []
150    num_expired_snapshots = 0
151    for batch in iter_expired_snapshot_batches(
152        state_reader=state_sync,
153        current_ts=current_ts,
154        ignore_ttl=ignore_ttl,
155        batch_size=batch_size,
156    ):
157        end_info = (
158            f"updated_ts={batch.batch_range.end.updated_ts}"
159            if isinstance(batch.batch_range.end, RowBoundary)
160            else f"limit={batch.batch_range.end.batch_size}"
161        )
162        logger.info(
163            "Processing batch of size %s with end %s",
164            len(batch.expired_snapshot_ids),
165            end_info,
166        )
167        cleanup_succeeded = True
168        try:
169            snapshot_evaluator.cleanup(
170                target_snapshots=batch.cleanup_tasks,
171                on_complete=console.update_cleanup_progress if console else None,
172            )
173        except Exception as failed_drops:
174            message = f"Failed to clean up: {failed_drops}"
175            logger.warning(message)
176            failures.append(message)
177            cleanup_succeeded = False
178
179        if cleanup_succeeded or force_delete:
180            try:
181                state_sync.delete_expired_snapshots(
182                    batch_range=ExpiredBatchRange(
183                        start=RowBoundary.lowest_boundary(),
184                        end=batch.batch_range.end,
185                    ),
186                    ignore_ttl=ignore_ttl,
187                )
188                logger.info("Cleaned up expired snapshots batch")
189                num_expired_snapshots += len(batch.expired_snapshot_ids)
190            except Exception as e:
191                message = f"Failed to delete expired snapshot state records: {e}"
192                logger.warning(message)
193                failures.append(message)
194    logger.info("Cleaned up %s expired snapshots", num_expired_snapshots)
195    return failures