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(
325        self, current_ts: int, name: t.Optional[str] = None
326    ) -> t.List[EnvironmentSummary]:
327        """Returns the expired environments.
328
329        Expired environments are environments that have exceeded their time-to-live value.
330
331        Args:
332            current_ts: The current timestamp in milliseconds used to determine expiration.
333            name: If provided, only the environment with this name is considered.
334
335        Returns:
336            The list of environment summaries to remove.
337        """
338
339
340class StateSync(StateReader, abc.ABC):
341    """Abstract base class for snapshot and environment state management."""
342
343    @abc.abstractmethod
344    def push_snapshots(self, snapshots: t.Iterable[Snapshot]) -> None:
345        """Push snapshots into the state sync.
346
347        This method only allows for pushing new snapshots. If existing snapshots are found,
348        this method should raise an error.
349
350        Raises:
351            SQLMeshError when existing snapshots are pushed.
352
353        Args:
354            snapshots: A list of snapshots to save in the state sync.
355        """
356
357    @abc.abstractmethod
358    def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
359        """Delete snapshots from the state sync.
360
361        Args:
362            snapshot_ids: A list of snapshot like objects to delete.
363        """
364
365    @abc.abstractmethod
366    def delete_expired_snapshots(
367        self,
368        batch_range: ExpiredBatchRange,
369        ignore_ttl: bool = False,
370        current_ts: t.Optional[int] = None,
371    ) -> None:
372        """Removes expired snapshots.
373
374        Expired snapshots are snapshots that have exceeded their time-to-live
375        and are no longer in use within an environment.
376
377        Args:
378            batch_range: The range of snapshots to delete in this batch.
379            ignore_ttl: Ignore the TTL on the snapshot when considering it expired. This has the effect of deleting
380                all snapshots that are not referenced in any environment
381            current_ts: Timestamp used to evaluate expiration.
382        """
383
384    @abc.abstractmethod
385    def invalidate_environment(self, name: str, protect_prod: bool = True) -> None:
386        """Invalidates the target environment by setting its expiration timestamp to now.
387
388        Args:
389            name: The name of the environment to invalidate.
390            protect_prod: If True, prevents invalidation of the production environment.
391        """
392
393    @abc.abstractmethod
394    def remove_state(self, including_backup: bool = False) -> None:
395        """Removes the state store objects."""
396
397    @abc.abstractmethod
398    def remove_intervals(
399        self,
400        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
401        remove_shared_versions: bool = False,
402    ) -> None:
403        """Remove an interval from a list of snapshots and sync it to the store.
404
405        Because multiple snapshots can be pointing to the same version or physical table, this method
406        can also grab all snapshots tied to the passed in version.
407
408        Args:
409            snapshot_intervals: The snapshot intervals to remove.
410            remove_shared_versions: Whether to remove intervals for snapshots that share the same version with the target snapshots.
411        """
412
413    @abc.abstractmethod
414    def promote(
415        self,
416        environment: Environment,
417        no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
418        environment_statements: t.Optional[t.List[EnvironmentStatements]] = None,
419    ) -> PromotionResult:
420        """Update the environment to reflect the current state.
421
422        This method verifies that snapshots have been pushed.
423
424        Args:
425            environment: The environment to promote.
426            no_gaps_snapshot_names: A set of snapshot names to check for data gaps. If None,
427                all snapshots will be checked. The data gap check ensures that models that are already a
428                part of the target environment have no data gaps when compared against previous
429                snapshots for same models.
430
431        Returns:
432           A tuple of (added snapshot table infos, removed snapshot table infos)
433        """
434
435    @abc.abstractmethod
436    def finalize(self, environment: Environment) -> None:
437        """Finalize the target environment, indicating that this environment has been
438        fully promoted and is ready for use.
439
440        Args:
441            environment: The target environment to finalize.
442        """
443
444    @abc.abstractmethod
445    def delete_expired_environments(
446        self, current_ts: t.Optional[int] = None, name: t.Optional[str] = None
447    ) -> t.List[EnvironmentSummary]:
448        """Removes expired environments.
449
450        Expired environments are environments that have exceeded their time-to-live value.
451
452        Args:
453            current_ts: The current timestamp in milliseconds. Defaults to now.
454            name: If provided, only the environment with this name is deleted.
455
456        Returns:
457            The list of removed environments.
458        """
459
460    @abc.abstractmethod
461    def unpause_snapshots(
462        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
463    ) -> None:
464        """Unpauses target snapshots.
465
466        Unpaused snapshots are scheduled for evaluation on a recurring basis.
467        Once unpaused a snapshot can't be paused again.
468
469        Args:
470            snapshots: Target snapshots.
471            unpaused_dt: The datetime object which indicates when target snapshots
472                were unpaused.
473        """
474
475    @abc.abstractmethod
476    def compact_intervals(self) -> None:
477        """Compacts intervals for all snapshots.
478
479        Compaction process involves merging of existing interval records into new records and
480        then deleting the old ones.
481        """
482
483    @abc.abstractmethod
484    def migrate(
485        self,
486        skip_backup: bool = False,
487        promoted_snapshots_only: bool = True,
488    ) -> None:
489        """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
490
491    @abc.abstractmethod
492    def rollback(self) -> None:
493        """Rollback to previous backed up state."""
494
495    @abc.abstractmethod
496    def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
497        """Add snapshot intervals to state
498
499        Args:
500            snapshots_intervals: The intervals to add.
501        """
502
503    def add_interval(
504        self,
505        snapshot: Snapshot,
506        start: TimeLike,
507        end: TimeLike,
508        is_dev: bool = False,
509        last_altered_ts: t.Optional[int] = None,
510    ) -> None:
511        """Add an interval to a snapshot and sync it to the store.
512
513        Args:
514            snapshot: The snapshot like object to add an interval to.
515            start: The start of the interval to add.
516            end: The end of the interval to add.
517            is_dev: Indicates whether the given interval is being added while in development mode
518            last_altered_ts: The timestamp of the last modification of the physical table
519        """
520        start_ts, end_ts = snapshot.inclusive_exclusive(start, end, strict=False, expand=False)
521        if not snapshot.version:
522            raise SQLMeshError("Snapshot version must be set to add an interval.")
523        intervals = [(start_ts, end_ts)]
524        snapshot_intervals = SnapshotIntervals(
525            name=snapshot.name,
526            identifier=snapshot.identifier,
527            version=snapshot.version,
528            dev_version=snapshot.dev_version,
529            intervals=intervals if not is_dev else [],
530            dev_intervals=intervals if is_dev else [],
531            last_altered_ts=last_altered_ts if not is_dev else None,
532            dev_last_altered_ts=last_altered_ts if is_dev else None,
533        )
534        self.add_snapshots_intervals([snapshot_intervals])
535
536    @abc.abstractmethod
537    def import_(self, stream: StateStream, clear: bool = True) -> None:
538        """
539        Replace the existing state with the state contained in the StateStream
540
541        Args:
542            stream: The stream of new state
543            clear: Whether or not to clear existing state before inserting state from the stream
544        """
545
546
547class DelegatingStateSync(StateSync):
548    def __init__(self, state_sync: StateSync) -> None:
549        self.state_sync = state_sync
550
551
552def _create_delegate_method(name: str) -> t.Callable:
553    def delegate(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
554        return getattr(self.state_sync, name)(*args, **kwargs)
555
556    return delegate
557
558
559DelegatingStateSync.__abstractmethods__ = frozenset()
560for name in StateSync.__abstractmethods__:
561    setattr(DelegatingStateSync, name, _create_delegate_method(name))