sqlmesh.core.state_sync.base

  1from __future__ import annotations
  2
  3import abc
  4import importlib
  5import logging
  6import pkgutil
  7import typing as t
  8
  9from sqlglot import __version__ as SQLGLOT_VERSION
 10
 11from sqlmesh import migrations
 12from sqlmesh.core.environment import (
 13    Environment,
 14    EnvironmentStatements,
 15    EnvironmentSummary,
 16)
 17from sqlmesh.core.snapshot import (
 18    Snapshot,
 19    SnapshotId,
 20    SnapshotIdLike,
 21    SnapshotIdAndVersionLike,
 22    SnapshotInfoLike,
 23    SnapshotNameVersion,
 24    SnapshotIdAndVersion,
 25)
 26from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
 27from sqlmesh.utils import major_minor
 28from sqlmesh.utils.date import TimeLike
 29from sqlmesh.utils.errors import SQLMeshError
 30from sqlmesh.utils.pydantic import PydanticModel, field_validator
 31from sqlmesh.core.state_sync.common import (
 32    StateStream,
 33    ExpiredSnapshotBatch,
 34    PromotionResult,
 35    ExpiredBatchRange,
 36)
 37
 38logger = logging.getLogger(__name__)
 39
 40
 41class Versions(PydanticModel):
 42    """Represents the various versions of dependencies in the state sync."""
 43
 44    schema_version: int = 0
 45    sqlglot_version: str = "0.0.0"
 46    sqlmesh_version: str = "0.0.0"
 47
 48    @property
 49    def minor_sqlglot_version(self) -> t.Tuple[int, int]:
 50        return major_minor(self.sqlglot_version)
 51
 52    @property
 53    def minor_sqlmesh_version(self) -> t.Tuple[int, int]:
 54        return major_minor(self.sqlmesh_version)
 55
 56    @field_validator("sqlglot_version", "sqlmesh_version", mode="before")
 57    @classmethod
 58    def _package_version_validator(cls, v: t.Any) -> str:
 59        return "0.0.0" if v is None else str(v)
 60
 61    @field_validator("schema_version", mode="before")
 62    @classmethod
 63    def _schema_version_validator(cls, v: t.Any) -> int:
 64        return 0 if v is None else int(v)
 65
 66
 67MIN_SCHEMA_VERSION = 60
 68MIN_SQLMESH_VERSION = "0.134.0"
 69MIGRATIONS = [
 70    importlib.import_module(f"sqlmesh.migrations.{migration}")
 71    for migration in sorted(info.name for info in pkgutil.iter_modules(migrations.__path__))
 72]
 73# -1 to account for the baseline script
 74SCHEMA_VERSION: int = MIN_SCHEMA_VERSION + len(MIGRATIONS) - 1
 75
 76
 77class StateReader(abc.ABC):
 78    """Abstract base class for read-only operations on snapshot and environment state."""
 79
 80    @abc.abstractmethod
 81    def get_snapshots(
 82        self, snapshot_ids: t.Iterable[SnapshotIdLike]
 83    ) -> t.Dict[SnapshotId, Snapshot]:
 84        """Bulk fetch snapshots given the corresponding snapshot ids.
 85
 86        Args:
 87            snapshot_ids: Iterable of snapshot ids to get.
 88
 89        Returns:
 90            A dictionary of snapshot ids to snapshots for ones that could be found.
 91        """
 92
 93    @abc.abstractmethod
 94    def get_snapshots_by_names(
 95        self,
 96        snapshot_names: t.Iterable[str],
 97        current_ts: t.Optional[int] = None,
 98        exclude_expired: bool = True,
 99    ) -> t.Set[SnapshotIdAndVersion]:
100        """Return the snapshot records for all versions of the specified snapshot names.
101
102        Args:
103            snapshot_names: Iterable of snapshot names to fetch all snapshot records for
104            current_ts: Sets the current time for identifying which snapshots have expired so they can be excluded (only relevant if :exclude_expired=True)
105            exclude_expired: Whether or not to return the snapshot id's of expired snapshots in the result
106
107        Returns:
108            A set containing all the matched snapshot records. To fetch full snapshots, pass it into StateSync.get_snapshots()
109        """
110
111    @abc.abstractmethod
112    def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[SnapshotId]:
113        """Checks if multiple snapshots exist in the state sync.
114
115        Args:
116            snapshot_ids: Iterable of snapshot ids to bulk check.
117
118        Returns:
119            A set of all the existing snapshot ids.
120        """
121
122    @abc.abstractmethod
123    def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.List[Snapshot]:
124        """Updates given snapshots with latest intervals from the state.
125
126        Args:
127            snapshots: The snapshots to refresh.
128
129        Returns:
130            The updated snapshots.
131        """
132
133    @abc.abstractmethod
134    def nodes_exist(self, names: t.Iterable[str], exclude_external: bool = False) -> t.Set[str]:
135        """Returns the node names that exist in the state sync.
136
137        Args:
138            names: Iterable of node names to check.
139            exclude_external: Whether to exclude external models from the output.
140
141        Returns:
142            A set of all the existing node names.
143        """
144
145    @abc.abstractmethod
146    def get_environment(self, environment: str) -> t.Optional[Environment]:
147        """Fetches the environment if it exists.
148
149        Args:
150            environment: The environment
151
152        Returns:
153            The environment object.
154        """
155
156    @abc.abstractmethod
157    def get_environments(self) -> t.List[Environment]:
158        """Fetches all environments.
159
160        Returns:
161            A list of all environments.
162        """
163
164    @abc.abstractmethod
165    def get_environments_summary(self) -> t.List[EnvironmentSummary]:
166        """Fetches all environment names along with expiry datetime.
167
168        Returns:
169            A list of all environment summaries.
170        """
171
172    @abc.abstractmethod
173    def max_interval_end_per_model(
174        self,
175        environment: str,
176        models: t.Optional[t.Set[str]] = None,
177        ensure_finalized_snapshots: bool = False,
178    ) -> t.Dict[str, int]:
179        """Returns the max interval end per model for the given environment.
180
181        Args:
182            environment: The target environment.
183            models: The models to get the max interval end for. If None, all models are considered.
184            ensure_finalized_snapshots: Whether to use snapshots from the latest finalized environment state,
185                or to use whatever snapshots are in the current environment state even if the environment is not finalized.
186
187        Returns:
188            A dictionary of model FQNs to their respective interval ends in milliseconds since epoch.
189        """
190
191    @abc.abstractmethod
192    def recycle(self) -> None:
193        """Closes all open connections and releases all allocated resources associated with any thread
194        except the calling one."""
195
196    @abc.abstractmethod
197    def close(self) -> None:
198        """Closes all open connections and releases all allocated resources."""
199
200    @abc.abstractmethod
201    def state_type(self) -> str:
202        """Returns the type of state sync."""
203
204    @abc.abstractmethod
205    def update_auto_restatements(
206        self, next_auto_restatement_ts: t.Dict[SnapshotNameVersion, t.Optional[int]]
207    ) -> None:
208        """Updates the next auto restatement timestamp for the snapshots.
209
210        Args:
211            next_auto_restatement_ts: A dictionary of snapshot name / version pairs to the next auto restatement timestamp.
212        """
213
214    @abc.abstractmethod
215    def get_environment_statements(self, environment: str) -> t.List[EnvironmentStatements]:
216        """Fetches environment statements from the environment_statements table.
217
218        Returns:
219            A list of the Environment Statements.
220        """
221
222    def get_versions(self, validate: bool = True) -> Versions:
223        """Get the current versions of the SQLMesh schema and libraries.
224
225        Args:
226            validate: Whether or not to raise error if the running version is different from what's in state.
227
228        Returns:
229            The versions object.
230        """
231        from sqlmesh._version import __version__ as SQLMESH_VERSION
232
233        versions = self._get_versions()
234
235        if validate:
236
237            def raise_error(
238                lib: str,
239                local: str | int,
240                remote: str | int,
241                remote_package_version: t.Optional[str] = None,
242                ahead: bool = False,
243            ) -> None:
244                if ahead:
245                    raise SQLMeshError(
246                        f"{lib} (local) is using version '{local}' which is ahead of '{remote}' (remote). "
247                        "Please run a migration ('sqlmesh migrate' command)."
248                    )
249
250                if remote_package_version:
251                    upgrade_suggestion = f" Please upgrade {lib} ('pip install --upgrade \"{lib.lower()}=={remote_package_version}\"' command)."
252                else:
253                    upgrade_suggestion = ""
254
255                raise SQLMeshError(
256                    f"{lib} (local) is using version '{local}' which is behind '{remote}' (remote).{upgrade_suggestion}"
257                )
258
259            if major_minor(SQLMESH_VERSION) != major_minor(versions.sqlmesh_version):
260                raise_error(
261                    "SQLMesh",
262                    SQLMESH_VERSION,
263                    versions.sqlmesh_version,
264                    remote_package_version=versions.sqlmesh_version,
265                    ahead=major_minor(SQLMESH_VERSION) > major_minor(versions.sqlmesh_version),
266                )
267
268            if SCHEMA_VERSION != versions.schema_version:
269                raise_error(
270                    "SQLMesh",
271                    SCHEMA_VERSION,
272                    versions.schema_version,
273                    remote_package_version=versions.sqlmesh_version,
274                    ahead=SCHEMA_VERSION > versions.schema_version,
275                )
276
277            if major_minor(SQLGLOT_VERSION) != major_minor(versions.sqlglot_version):
278                raise_error(
279                    "SQLGlot",
280                    SQLGLOT_VERSION,
281                    versions.sqlglot_version,
282                    remote_package_version=versions.sqlglot_version,
283                    ahead=major_minor(SQLGLOT_VERSION) > major_minor(versions.sqlglot_version),
284                )
285
286        return versions
287
288    @abc.abstractmethod
289    def _get_versions(self) -> Versions:
290        """Queries the store to get the current versions of SQLMesh and deps.
291
292        Returns:
293            The versions object.
294        """
295
296    @abc.abstractmethod
297    def export(self, environment_names: t.Optional[t.List[str]] = None) -> StateStream:
298        """Export the contents of this StateSync as a StateStream
299
300        Args:
301            environment_names: An optional list of environment names to export. If not specified, all environments will be exported.
302        """
303
304    @abc.abstractmethod
305    def get_expired_snapshots(
306        self,
307        *,
308        batch_range: ExpiredBatchRange,
309        current_ts: t.Optional[int] = None,
310        ignore_ttl: bool = False,
311    ) -> t.Optional[ExpiredSnapshotBatch]:
312        """Returns a single batch of expired snapshots ordered by (updated_ts, name, identifier).
313
314        Args:
315            current_ts: Timestamp used to evaluate expiration.
316            ignore_ttl: If True, include snapshots regardless of TTL (only checks if unreferenced).
317            batch_range: The range of the batch to fetch.
318
319        Returns:
320            A batch describing expired snapshots or None if no snapshots are pending cleanup.
321        """
322
323    @abc.abstractmethod
324    def get_expired_environments(self, current_ts: int) -> t.List[EnvironmentSummary]:
325        """Returns the expired environments.
326
327        Expired environments are environments that have exceeded their time-to-live value.
328        Returns:
329            The list of environment summaries to remove.
330        """
331
332
333class StateSync(StateReader, abc.ABC):
334    """Abstract base class for snapshot and environment state management."""
335
336    @abc.abstractmethod
337    def push_snapshots(self, snapshots: t.Iterable[Snapshot]) -> None:
338        """Push snapshots into the state sync.
339
340        This method only allows for pushing new snapshots. If existing snapshots are found,
341        this method should raise an error.
342
343        Raises:
344            SQLMeshError when existing snapshots are pushed.
345
346        Args:
347            snapshots: A list of snapshots to save in the state sync.
348        """
349
350    @abc.abstractmethod
351    def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
352        """Delete snapshots from the state sync.
353
354        Args:
355            snapshot_ids: A list of snapshot like objects to delete.
356        """
357
358    @abc.abstractmethod
359    def delete_expired_snapshots(
360        self,
361        batch_range: ExpiredBatchRange,
362        ignore_ttl: bool = False,
363        current_ts: t.Optional[int] = None,
364    ) -> None:
365        """Removes expired snapshots.
366
367        Expired snapshots are snapshots that have exceeded their time-to-live
368        and are no longer in use within an environment.
369
370        Args:
371            batch_range: The range of snapshots to delete in this batch.
372            ignore_ttl: Ignore the TTL on the snapshot when considering it expired. This has the effect of deleting
373                all snapshots that are not referenced in any environment
374            current_ts: Timestamp used to evaluate expiration.
375        """
376
377    @abc.abstractmethod
378    def invalidate_environment(self, name: str, protect_prod: bool = True) -> None:
379        """Invalidates the target environment by setting its expiration timestamp to now.
380
381        Args:
382            name: The name of the environment to invalidate.
383            protect_prod: If True, prevents invalidation of the production environment.
384        """
385
386    @abc.abstractmethod
387    def remove_state(self, including_backup: bool = False) -> None:
388        """Removes the state store objects."""
389
390    @abc.abstractmethod
391    def remove_intervals(
392        self,
393        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
394        remove_shared_versions: bool = False,
395    ) -> None:
396        """Remove an interval from a list of snapshots and sync it to the store.
397
398        Because multiple snapshots can be pointing to the same version or physical table, this method
399        can also grab all snapshots tied to the passed in version.
400
401        Args:
402            snapshot_intervals: The snapshot intervals to remove.
403            remove_shared_versions: Whether to remove intervals for snapshots that share the same version with the target snapshots.
404        """
405
406    @abc.abstractmethod
407    def promote(
408        self,
409        environment: Environment,
410        no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
411        environment_statements: t.Optional[t.List[EnvironmentStatements]] = None,
412    ) -> PromotionResult:
413        """Update the environment to reflect the current state.
414
415        This method verifies that snapshots have been pushed.
416
417        Args:
418            environment: The environment to promote.
419            no_gaps_snapshot_names: A set of snapshot names to check for data gaps. If None,
420                all snapshots will be checked. The data gap check ensures that models that are already a
421                part of the target environment have no data gaps when compared against previous
422                snapshots for same models.
423
424        Returns:
425           A tuple of (added snapshot table infos, removed snapshot table infos)
426        """
427
428    @abc.abstractmethod
429    def finalize(self, environment: Environment) -> None:
430        """Finalize the target environment, indicating that this environment has been
431        fully promoted and is ready for use.
432
433        Args:
434            environment: The target environment to finalize.
435        """
436
437    @abc.abstractmethod
438    def delete_expired_environments(
439        self, current_ts: t.Optional[int] = None
440    ) -> t.List[EnvironmentSummary]:
441        """Removes expired environments.
442
443        Expired environments are environments that have exceeded their time-to-live value.
444
445        Returns:
446            The list of removed environments.
447        """
448
449    @abc.abstractmethod
450    def unpause_snapshots(
451        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
452    ) -> None:
453        """Unpauses target snapshots.
454
455        Unpaused snapshots are scheduled for evaluation on a recurring basis.
456        Once unpaused a snapshot can't be paused again.
457
458        Args:
459            snapshots: Target snapshots.
460            unpaused_dt: The datetime object which indicates when target snapshots
461                were unpaused.
462        """
463
464    @abc.abstractmethod
465    def compact_intervals(self) -> None:
466        """Compacts intervals for all snapshots.
467
468        Compaction process involves merging of existing interval records into new records and
469        then deleting the old ones.
470        """
471
472    @abc.abstractmethod
473    def migrate(
474        self,
475        skip_backup: bool = False,
476        promoted_snapshots_only: bool = True,
477    ) -> None:
478        """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
479
480    @abc.abstractmethod
481    def rollback(self) -> None:
482        """Rollback to previous backed up state."""
483
484    @abc.abstractmethod
485    def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
486        """Add snapshot intervals to state
487
488        Args:
489            snapshots_intervals: The intervals to add.
490        """
491
492    def add_interval(
493        self,
494        snapshot: Snapshot,
495        start: TimeLike,
496        end: TimeLike,
497        is_dev: bool = False,
498        last_altered_ts: t.Optional[int] = None,
499    ) -> None:
500        """Add an interval to a snapshot and sync it to the store.
501
502        Args:
503            snapshot: The snapshot like object to add an interval to.
504            start: The start of the interval to add.
505            end: The end of the interval to add.
506            is_dev: Indicates whether the given interval is being added while in development mode
507            last_altered_ts: The timestamp of the last modification of the physical table
508        """
509        start_ts, end_ts = snapshot.inclusive_exclusive(start, end, strict=False, expand=False)
510        if not snapshot.version:
511            raise SQLMeshError("Snapshot version must be set to add an interval.")
512        intervals = [(start_ts, end_ts)]
513        snapshot_intervals = SnapshotIntervals(
514            name=snapshot.name,
515            identifier=snapshot.identifier,
516            version=snapshot.version,
517            dev_version=snapshot.dev_version,
518            intervals=intervals if not is_dev else [],
519            dev_intervals=intervals if is_dev else [],
520            last_altered_ts=last_altered_ts if not is_dev else None,
521            dev_last_altered_ts=last_altered_ts if is_dev else None,
522        )
523        self.add_snapshots_intervals([snapshot_intervals])
524
525    @abc.abstractmethod
526    def import_(self, stream: StateStream, clear: bool = True) -> None:
527        """
528        Replace the existing state with the state contained in the StateStream
529
530        Args:
531            stream: The stream of new state
532            clear: Whether or not to clear existing state before inserting state from the stream
533        """
534
535
536class DelegatingStateSync(StateSync):
537    def __init__(self, state_sync: StateSync) -> None:
538        self.state_sync = state_sync
539
540
541def _create_delegate_method(name: str) -> t.Callable:
542    def delegate(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
543        return getattr(self.state_sync, name)(*args, **kwargs)
544
545    return delegate
546
547
548DelegatingStateSync.__abstractmethods__ = frozenset()
549for name in StateSync.__abstractmethods__:
550    setattr(DelegatingStateSync, name, _create_delegate_method(name))