sqlmesh.core.state_sync.db

 75class EngineAdapterStateSync(StateSync):
 76    """Manages state of nodes and snapshot with an existing engine adapter.
 77
 78    This state sync is convenient to use because it requires no additional setup.
 79    You can reuse the same engine/warehouse that your data is stored in.
 80
 81    Args:
 82        engine_adapter: The EngineAdapter to use to store and fetch snapshots.
 83        schema: The schema to store state metadata in. If None or empty string then no schema is defined
 84        console: The console to log information to.
 85        cache_dir: The cache path, used for caching snapshot models.
 86    """
 87
 88    def __init__(
 89        self,
 90        engine_adapter: EngineAdapter,
 91        schema: t.Optional[str],
 92        console: t.Optional[Console] = None,
 93        cache_dir: Path = Path(),
 94    ):
 95        self.interval_state = IntervalState(engine_adapter, schema=schema)
 96        self.environment_state = EnvironmentState(engine_adapter, schema=schema)
 97        self.snapshot_state = SnapshotState(engine_adapter, schema=schema, cache_dir=cache_dir)
 98        self.version_state = VersionState(engine_adapter, schema=schema)
 99        self.migrator = StateMigrator(
100            engine_adapter,
101            version_state=self.version_state,
102            snapshot_state=self.snapshot_state,
103            environment_state=self.environment_state,
104            interval_state=self.interval_state,
105            console=console,
106        )
107        # Make sure that if an empty string is provided that we treat it as None
108        self.schema = schema or None
109        self.engine_adapter = engine_adapter
110        self.console = console or get_console()
111
112    @transactional()
113    def push_snapshots(self, snapshots: t.Iterable[Snapshot]) -> None:
114        """Pushes snapshots to the state store, merging them with existing ones.
115
116        This method first finds all existing snapshots in the store and merges them with
117        the local snapshots. It will then delete all existing snapshots and then
118        insert all the local snapshots. This can be made safer with locks or merge/upsert.
119
120        Args:
121            snapshots: The snapshots to push.
122        """
123        snapshots_by_id = {}
124        for snapshot in snapshots:
125            if not snapshot.version:
126                raise SQLMeshError(
127                    f"Snapshot {snapshot} has not been versioned yet. Create a plan before pushing a snapshot."
128                )
129            snapshots_by_id[snapshot.snapshot_id] = snapshot
130
131        existing = self.snapshots_exist(snapshots_by_id)
132
133        if existing:
134            logger.error(
135                "Snapshots %s already exists. This could be due to a concurrent plan or a hash collision. If this is a hash collision, add a stamp to your model.",
136                str(existing),
137            )
138
139            for sid in tuple(snapshots_by_id):
140                if sid in existing:
141                    snapshots_by_id.pop(sid)
142
143        snapshots = snapshots_by_id.values()
144        if snapshots:
145            self.snapshot_state.push_snapshots(snapshots)
146
147    @transactional()
148    def promote(
149        self,
150        environment: Environment,
151        no_gaps_snapshot_names: t.Optional[t.Set[str]] = None,
152        environment_statements: t.Optional[t.List[EnvironmentStatements]] = None,
153    ) -> PromotionResult:
154        """Update the environment to reflect the current state.
155
156        This method verifies that snapshots have been pushed.
157
158        Args:
159            environment: The environment to promote.
160            no_gaps_snapshot_names: A set of snapshot names to check for data gaps. If None,
161                all snapshots will be checked. The data gap check ensures that models that are already a
162                part of the target environment have no data gaps when compared against previous
163                snapshots for same models.
164
165        Returns:
166           A tuple of (added snapshot table infos, removed snapshot table infos, and environment target suffix for the removed table infos)
167        """
168        logger.info("Promoting environment '%s'", environment.name)
169
170        missing = {s.snapshot_id for s in environment.snapshots} - self.snapshots_exist(
171            environment.snapshots
172        )
173        if missing:
174            raise SQLMeshError(
175                f"Missing snapshots {missing}. Make sure to push and backfill your snapshots."
176            )
177
178        existing_environment = self.environment_state.get_environment(
179            environment.name, lock_for_update=True
180        )
181
182        existing_table_infos = (
183            {table_info.name: table_info for table_info in existing_environment.promoted_snapshots}
184            if existing_environment
185            else {}
186        )
187        table_infos = {table_info.name: table_info for table_info in environment.promoted_snapshots}
188        views_that_changed_location: t.Set[SnapshotTableInfo] = set()
189        if existing_environment:
190            views_that_changed_location = {
191                existing_table_info
192                for name, existing_table_info in existing_table_infos.items()
193                if name in table_infos
194                and existing_table_info.qualified_view_name.for_environment(
195                    existing_environment.naming_info
196                )
197                != table_infos[name].qualified_view_name.for_environment(environment.naming_info)
198            }
199            if not existing_environment.expired:
200                if environment.previous_plan_id != existing_environment.plan_id:
201                    raise ConflictingPlanError(
202                        f"Another plan ({existing_environment.plan_id}) was applied to the target environment '{environment.name}' while your current plan "
203                        f"({environment.plan_id}) was still in progress, interrupting it. Please re-apply your plan to resolve this error."
204                    )
205                if no_gaps_snapshot_names != set():
206                    snapshots = self.get_snapshots(environment.snapshots).values()
207                    self._ensure_no_gaps(
208                        snapshots,
209                        existing_environment,
210                        no_gaps_snapshot_names,
211                    )
212            demoted_snapshots = set(existing_environment.snapshots) - set(environment.snapshots)
213            # Update the updated_at attribute.
214            self.snapshot_state.touch_snapshots(demoted_snapshots)
215
216        missing_models = set(existing_table_infos) - {
217            snapshot.name for snapshot in environment.promoted_snapshots
218        }
219
220        added_table_infos = set(table_infos.values())
221        if existing_environment and environment.can_partially_promote(existing_environment):
222            # Only promote new snapshots.
223            added_table_infos -= set(existing_environment.promoted_snapshots)
224
225        self.environment_state.update_environment(environment)
226
227        # If it is an empty list, we want to update the environment statements
228        # To reflect there are no statements anymore in this environment
229        if environment_statements is not None:
230            self.environment_state.update_environment_statements(
231                environment.name, environment.plan_id, environment_statements
232            )
233
234        removed = {existing_table_infos[name] for name in missing_models}.union(
235            views_that_changed_location
236        )
237
238        return PromotionResult(
239            added=sorted(added_table_infos),
240            removed=list(removed),
241            removed_environment_naming_info=(
242                existing_environment.naming_info if removed and existing_environment else None
243            ),
244        )
245
246    @transactional()
247    def finalize(self, environment: Environment) -> None:
248        """Finalize the target environment, indicating that this environment has been
249        fully promoted and is ready for use.
250
251        Args:
252            environment: The target environment to finalize.
253        """
254        self.environment_state.finalize(environment)
255
256    @transactional()
257    def unpause_snapshots(
258        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
259    ) -> None:
260        self.snapshot_state.unpause_snapshots(snapshots, unpaused_dt)
261
262    def invalidate_environment(self, name: str, protect_prod: bool = True) -> None:
263        self.environment_state.invalidate_environment(name, protect_prod)
264
265    def get_expired_snapshots(
266        self,
267        *,
268        batch_range: ExpiredBatchRange,
269        current_ts: t.Optional[int] = None,
270        ignore_ttl: bool = False,
271    ) -> t.Optional[ExpiredSnapshotBatch]:
272        current_ts = current_ts or now_timestamp()
273        return self.snapshot_state.get_expired_snapshots(
274            environments=self.environment_state.get_environments(),
275            current_ts=current_ts,
276            ignore_ttl=ignore_ttl,
277            batch_range=batch_range,
278        )
279
280    def get_expired_environments(
281        self, current_ts: int, name: t.Optional[str] = None
282    ) -> t.List[EnvironmentSummary]:
283        return self.environment_state.get_expired_environments(current_ts=current_ts, name=name)
284
285    @transactional()
286    def delete_expired_snapshots(
287        self,
288        batch_range: ExpiredBatchRange,
289        ignore_ttl: bool = False,
290        current_ts: t.Optional[int] = None,
291    ) -> None:
292        batch = self.get_expired_snapshots(
293            ignore_ttl=ignore_ttl,
294            current_ts=current_ts,
295            batch_range=batch_range,
296        )
297        if batch and batch.expired_snapshot_ids:
298            self.snapshot_state.delete_snapshots(batch.expired_snapshot_ids)
299            self.interval_state.cleanup_intervals(batch.cleanup_tasks, batch.expired_snapshot_ids)
300
301    @transactional()
302    def delete_expired_environments(
303        self, current_ts: t.Optional[int] = None, name: t.Optional[str] = None
304    ) -> t.List[EnvironmentSummary]:
305        current_ts = current_ts or now_timestamp()
306        return self.environment_state.delete_expired_environments(current_ts=current_ts, name=name)
307
308    def delete_snapshots(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> None:
309        self.snapshot_state.delete_snapshots(snapshot_ids)
310
311    def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[SnapshotId]:
312        return self.snapshot_state.snapshots_exist(snapshot_ids)
313
314    def nodes_exist(self, names: t.Iterable[str], exclude_external: bool = False) -> t.Set[str]:
315        return self.snapshot_state.nodes_exist(names, exclude_external)
316
317    def remove_state(self, including_backup: bool = False) -> None:
318        """Removes the state store objects."""
319        for table in (
320            self.snapshot_state.snapshots_table,
321            self.snapshot_state.auto_restatements_table,
322            self.environment_state.environments_table,
323            self.environment_state.environment_statements_table,
324            self.interval_state.intervals_table,
325            self.version_state.versions_table,
326        ):
327            self.engine_adapter.drop_table(table)
328            if including_backup:
329                self.engine_adapter.drop_table(_backup_table_name(table))
330
331        self.snapshot_state.clear_cache()
332
333    def reset(self, default_catalog: t.Optional[str]) -> None:
334        """Resets the state store to the state when it was first initialized."""
335        self.remove_state()
336        self.migrate(default_catalog)
337
338    @transactional()
339    def update_auto_restatements(
340        self, next_auto_restatement_ts: t.Dict[SnapshotNameVersion, t.Optional[int]]
341    ) -> None:
342        self.snapshot_state.update_auto_restatements(next_auto_restatement_ts)
343
344    def get_environment(self, environment: str) -> t.Optional[Environment]:
345        return self.environment_state.get_environment(environment)
346
347    def get_environment_statements(self, environment: str) -> t.List[EnvironmentStatements]:
348        return self.environment_state.get_environment_statements(environment)
349
350    def get_environments(self) -> t.List[Environment]:
351        """Fetches all environments.
352
353        Returns:
354            A list of all environments.
355        """
356        return self.environment_state.get_environments()
357
358    def get_environments_summary(self) -> t.List[EnvironmentSummary]:
359        """Fetches all environment names along with expiry datetime.
360
361        Returns:
362            A list of all environment summaries.
363        """
364        return self.environment_state.get_environments_summary()
365
366    def get_snapshots(
367        self,
368        snapshot_ids: t.Iterable[SnapshotIdLike],
369    ) -> t.Dict[SnapshotId, Snapshot]:
370        """Fetches snapshots from the state.
371
372        Args:
373            snapshot_ids: The snapshot IDs to fetch.
374
375        Returns:
376            A dict of snapshots.
377        """
378        snapshots = self.snapshot_state.get_snapshots(snapshot_ids)
379        intervals = self.interval_state.get_snapshot_intervals(snapshots.values())
380        Snapshot.hydrate_with_intervals_by_version(snapshots.values(), intervals)
381        return snapshots
382
383    def get_snapshots_by_names(
384        self,
385        snapshot_names: t.Iterable[str],
386        current_ts: t.Optional[int] = None,
387        exclude_expired: bool = True,
388    ) -> t.Set[SnapshotIdAndVersion]:
389        return self.snapshot_state.get_snapshots_by_names(
390            snapshot_names=snapshot_names, current_ts=current_ts, exclude_expired=exclude_expired
391        )
392
393    @transactional()
394    def add_interval(
395        self,
396        snapshot: Snapshot,
397        start: TimeLike,
398        end: TimeLike,
399        is_dev: bool = False,
400        last_altered_ts: t.Optional[int] = None,
401    ) -> None:
402        super().add_interval(snapshot, start, end, is_dev, last_altered_ts)
403
404    @transactional()
405    def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotIntervals]) -> None:
406        intervals_to_insert = []
407        for snapshot_intervals in snapshots_intervals:
408            snapshot_intervals = snapshot_intervals.copy(
409                update={
410                    "intervals": _remove_partial_intervals(
411                        snapshot_intervals.intervals, snapshot_intervals.snapshot_id, is_dev=False
412                    ),
413                    "dev_intervals": _remove_partial_intervals(
414                        snapshot_intervals.dev_intervals,
415                        snapshot_intervals.snapshot_id,
416                        is_dev=True,
417                    ),
418                }
419            )
420            if not snapshot_intervals.is_empty():
421                intervals_to_insert.append(snapshot_intervals)
422        if intervals_to_insert:
423            self.interval_state.add_snapshots_intervals(intervals_to_insert)
424
425    @transactional()
426    def remove_intervals(
427        self,
428        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
429        remove_shared_versions: bool = False,
430    ) -> None:
431        self.interval_state.remove_intervals(snapshot_intervals, remove_shared_versions)
432
433    @transactional()
434    def compact_intervals(self) -> None:
435        self.interval_state.compact_intervals()
436
437    def refresh_snapshot_intervals(self, snapshots: t.Collection[Snapshot]) -> t.List[Snapshot]:
438        return self.interval_state.refresh_snapshot_intervals(snapshots)
439
440    def max_interval_end_per_model(
441        self,
442        environment: str,
443        models: t.Optional[t.Set[str]] = None,
444        ensure_finalized_snapshots: bool = False,
445    ) -> t.Dict[str, int]:
446        env = self.get_environment(environment)
447        if not env:
448            return {}
449
450        snapshots = (
451            env.snapshots if not ensure_finalized_snapshots else env.finalized_or_current_snapshots
452        )
453        if models is not None:
454            snapshots = [s for s in snapshots if s.name in models]
455
456        if not snapshots:
457            return {}
458
459        return self.interval_state.max_interval_end_per_model(snapshots)
460
461    def recycle(self) -> None:
462        self.engine_adapter.recycle()
463
464    def close(self) -> None:
465        self.engine_adapter.close()
466
467    @transactional()
468    def migrate(
469        self,
470        skip_backup: bool = False,
471        promoted_snapshots_only: bool = True,
472    ) -> None:
473        """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
474        self.migrator.migrate(
475            self.schema,
476            skip_backup=skip_backup,
477            promoted_snapshots_only=promoted_snapshots_only,
478        )
479
480    @transactional()
481    def rollback(self) -> None:
482        """Rollback to the previous migration."""
483        self.migrator.rollback()
484
485    @transactional()
486    def export(self, environment_names: t.Optional[t.List[str]] = None) -> StateStream:
487        versions = self.get_versions(
488            validate=True
489        )  # will throw if the state db hasnt been created or there is a version mismatch
490
491        snapshot_ids_to_export: t.Set[SnapshotId] = set()
492        selected_environments: t.List[Environment] = []
493        if environment_names:
494            for env_name in environment_names:
495                environment = self.get_environment(env_name)
496                if not environment:
497                    raise SQLMeshError(f"No such environment: {env_name}")
498                selected_environments.append(environment)
499        else:
500            selected_environments = self.get_environments()
501
502        for env in selected_environments:
503            snapshot_ids_to_export |= set([s.snapshot_id for s in env.snapshots or []])
504
505        def _export_snapshots() -> t.Iterator[Snapshot]:
506            for chunk in chunk_iterable(snapshot_ids_to_export, SnapshotState.SNAPSHOT_BATCH_SIZE):
507                yield from self.get_snapshots(chunk).values()
508
509        def _export_environments() -> t.Iterator[EnvironmentWithStatements]:
510            for env in selected_environments:
511                yield EnvironmentWithStatements(
512                    environment=env, statements=self.get_environment_statements(env.name)
513                )
514
515        return StateStream.from_iterators(
516            versions=versions,
517            snapshots=_export_snapshots(),
518            environments=_export_environments(),
519        )
520
521    @transactional()
522    def import_(self, stream: StateStream, clear: bool = True) -> None:
523        existing_versions = self.get_versions()
524
525        for state_chunk in stream:
526            if isinstance(state_chunk, VersionsChunk):
527                # SQLMesh major/minor version must match so that we can be sure the JSON contained in the state file
528                # is compatible with our Pydantic model definitions. Patch versions dont need to match because the assumption
529                # is that they dont contain any breaking changes
530                incoming_versions = state_chunk.versions
531                if (
532                    incoming_versions.minor_sqlmesh_version
533                    != existing_versions.minor_sqlmesh_version
534                ):
535                    raise SQLMeshError(
536                        f"SQLMesh version mismatch. You are running '{existing_versions.sqlmesh_version}' but the state file was created with '{incoming_versions.sqlmesh_version}'.\n"
537                        "Please upgrade/downgrade your SQLMesh version to match the state file before performing the import."
538                    )
539
540                if clear:
541                    self.reset(default_catalog=None)
542
543            if isinstance(state_chunk, SnapshotsChunk):
544                auto_restatements: t.Dict[SnapshotNameVersion, t.Optional[int]] = {}
545
546                for snapshot_chunk in chunk_iterable(
547                    state_chunk, SnapshotState.SNAPSHOT_BATCH_SIZE
548                ):
549                    snapshot_chunk = list(snapshot_chunk)
550                    overwrite_existing_snapshots = (
551                        not clear
552                    )  # if clear=True, all existing snapshots were dropped anyway
553                    self.snapshot_state.push_snapshots(
554                        snapshot_chunk, overwrite=overwrite_existing_snapshots
555                    )
556                    self.add_snapshots_intervals((s.snapshot_intervals for s in snapshot_chunk))
557
558                    auto_restatements.update(
559                        {
560                            s.name_version: s.next_auto_restatement_ts
561                            for s in snapshot_chunk
562                            if s.next_auto_restatement_ts
563                        }
564                    )
565
566                self.update_auto_restatements(auto_restatements)
567
568            if isinstance(state_chunk, EnvironmentsChunk):
569                for environment_with_statements in state_chunk:
570                    environment = environment_with_statements.environment
571                    self.environment_state.update_environment(environment)
572                    self.environment_state.update_environment_statements(
573                        environment.name,
574                        environment.plan_id,
575                        environment_with_statements.statements,
576                    )
577
578    def state_type(self) -> str:
579        return self.engine_adapter.dialect
580
581    def _get_versions(self) -> Versions:
582        return self.version_state.get_versions()
583
584    def _ensure_no_gaps(
585        self,
586        target_snapshots: t.Iterable[Snapshot],
587        target_environment: Environment,
588        snapshot_names: t.Optional[t.Set[str]],
589    ) -> None:
590        target_snapshots_by_name = {s.name: s for s in target_snapshots}
591
592        changed_version_prev_snapshots_by_name = {
593            s.name: s
594            for s in target_environment.snapshots
595            if s.name in target_snapshots_by_name
596            and target_snapshots_by_name[s.name].version != s.version
597        }
598
599        prev_snapshots = self.get_snapshots(
600            changed_version_prev_snapshots_by_name.values()
601        ).values()
602        cache: t.Dict[str, datetime] = {}
603
604        for prev_snapshot in prev_snapshots:
605            target_snapshot = target_snapshots_by_name[prev_snapshot.name]
606            if (
607                (snapshot_names is None or prev_snapshot.name in snapshot_names)
608                and target_snapshot.is_incremental
609                and prev_snapshot.is_incremental
610                and prev_snapshot.intervals
611            ):
612                start = to_timestamp(
613                    start_date(target_snapshot, target_snapshots_by_name.values(), cache)
614                )
615                end = prev_snapshot.intervals[-1][1]
616
617                if start < end:
618                    missing_intervals = target_snapshot.missing_intervals(
619                        start, end, end_bounded=True
620                    )
621
622                    if missing_intervals:
623                        raise SQLMeshError(
624                            f"Detected missing intervals for model {target_snapshot.name}, interrupting your current plan. "
625                            "Please re-apply your plan to resolve this error."
626                        )
627
628    @contextlib.contextmanager
629    def _transaction(self) -> t.Iterator[None]:
630        with self.engine_adapter.transaction():
631            yield