sqlmesh.core.scheduler

  1from __future__ import annotations
  2
  3import logging
  4import traceback
  5import typing as t
  6from datetime import datetime
  7
  8from sqlmesh.core import constants as c
  9from sqlmesh.core.console import Console, get_console
 10from sqlmesh.core.environment import EnvironmentNamingInfo
 11from sqlmesh.core.model import SeedModel
 12from sqlmesh.core.notification_target import (
 13    NotificationEvent,
 14    NotificationTargetManager,
 15)
 16from sqlmesh.core.snapshot import (
 17    DeployabilityIndex,
 18    Snapshot,
 19    SnapshotEvaluator,
 20    earliest_start_date,
 21    missing_intervals,
 22)
 23from sqlmesh.core.snapshot.definition import Interval as SnapshotInterval
 24from sqlmesh.core.snapshot.definition import SnapshotId
 25from sqlmesh.core.state_sync import StateSync
 26from sqlmesh.utils import format_exception
 27from sqlmesh.utils.concurrency import concurrent_apply_to_dag
 28from sqlmesh.utils.dag import DAG
 29from sqlmesh.utils.date import (
 30    TimeLike,
 31    now,
 32    now_timestamp,
 33    to_datetime,
 34    validate_date_range,
 35)
 36from sqlmesh.utils.errors import AuditError, CircuitBreakerError, SQLMeshError
 37
 38logger = logging.getLogger(__name__)
 39Interval = t.Tuple[datetime, datetime]
 40Batch = t.List[Interval]
 41SnapshotToBatches = t.Dict[Snapshot, Batch]
 42# we store snapshot name instead of snapshots/snapshotids because pydantic
 43# is extremely slow to hash. snapshot names should be unique within a dag run
 44SchedulingUnit = t.Tuple[str, t.Tuple[Interval, int]]
 45
 46
 47class Scheduler:
 48    """Schedules and manages the evaluation of snapshots.
 49
 50    The scheduler evaluates multiple snapshots with date intervals in the correct
 51    topological order. It consults the state sync to understand what intervals for each
 52    snapshot needs to be backfilled.
 53
 54    The scheduler comes equipped with a simple ThreadPoolExecutor based evaluation engine.
 55
 56    Args:
 57        snapshots: A collection of snapshots.
 58        snapshot_evaluator: The snapshot evaluator to execute queries.
 59        state_sync: The state sync to pull saved snapshots.
 60        max_workers: The maximum number of parallel queries to run.
 61        console: The rich instance used for printing scheduling information.
 62    """
 63
 64    def __init__(
 65        self,
 66        snapshots: t.Iterable[Snapshot],
 67        snapshot_evaluator: SnapshotEvaluator,
 68        state_sync: StateSync,
 69        default_catalog: t.Optional[str],
 70        max_workers: int = 1,
 71        console: t.Optional[Console] = None,
 72        notification_target_manager: t.Optional[NotificationTargetManager] = None,
 73    ):
 74        self.state_sync = state_sync
 75        self.snapshots = {s.snapshot_id: s for s in snapshots}
 76        self.snapshot_per_version = _resolve_one_snapshot_per_version(self.snapshots.values())
 77        self.default_catalog = default_catalog
 78        self.snapshot_evaluator = snapshot_evaluator
 79        self.max_workers = max_workers
 80        self.console = console or get_console()
 81        self.notification_target_manager = (
 82            notification_target_manager or NotificationTargetManager()
 83        )
 84
 85    def batches(
 86        self,
 87        start: t.Optional[TimeLike] = None,
 88        end: t.Optional[TimeLike] = None,
 89        execution_time: t.Optional[TimeLike] = None,
 90        deployability_index: t.Optional[DeployabilityIndex] = None,
 91        restatements: t.Optional[t.Dict[SnapshotId, SnapshotInterval]] = None,
 92        ignore_cron: bool = False,
 93        end_bounded: bool = False,
 94        selected_snapshots: t.Optional[t.Set[str]] = None,
 95    ) -> SnapshotToBatches:
 96        """Find the optimal date interval paramaters based on what needs processing and maximal batch size.
 97
 98        For each node name, find all dependencies and look for a stored snapshot from the metastore. If a snapshot is found,
 99        calculate the missing intervals that need to be processed given the passed in start and end intervals.
100
101        If a snapshot's node specifies a batch size, consecutive intervals are merged into batches of a size that is less than
102        or equal to the configured one. If no batch size is specified, then it uses the intervals that correspond to the node's cron expression.
103        For example, if a node is supposed to run daily and has 70 days to backfill with a batch size set to 30, there would be 2 jobs
104        with 30 days and 1 job with 10.
105
106        Args:
107            start: The start of the run. Defaults to the min node start date.
108            end: The end of the run. Defaults to now.
109            execution_time: The date/time time reference to use for execution time. Defaults to now.
110            deployability_index: Determines snapshots that are deployable in the context of this evaluation.
111            restatements: A set of snapshot names being restated.
112            ignore_cron: Whether to ignore the node's cron schedule.
113            end_bounded: If set to true, the returned intervals will be bounded by the target end date, disregarding lookback,
114                allow_partials, and other attributes that could cause the intervals to exceed the target end date.
115            selected_snapshots: A set of snapshot names to run. If not provided, all snapshots will be run.
116        """
117        restatements = restatements or {}
118        validate_date_range(start, end)
119
120        snapshots: t.Collection[Snapshot] = self.snapshot_per_version.values()
121        if selected_snapshots is not None:
122            snapshots = [s for s in snapshots if s.name in selected_snapshots]
123
124        self.state_sync.refresh_snapshot_intervals(snapshots)
125
126        return compute_interval_params(
127            snapshots,
128            start=start or earliest_start_date(snapshots),
129            end=end or now(),
130            deployability_index=deployability_index,
131            execution_time=execution_time or now(),
132            restatements=restatements,
133            ignore_cron=ignore_cron,
134            end_bounded=end_bounded,
135        )
136
137    def evaluate(
138        self,
139        snapshot: Snapshot,
140        start: TimeLike,
141        end: TimeLike,
142        execution_time: TimeLike,
143        deployability_index: DeployabilityIndex,
144        **kwargs: t.Any,
145    ) -> None:
146        """Evaluate a snapshot and add the processed interval to the state sync.
147
148        Args:
149            snapshot: Snapshot to evaluate.
150            start: The start datetime to render.
151            end: The end datetime to render.
152            execution_time: The date/time time reference to use for execution time. Defaults to now.
153            deployability_index: Determines snapshots that are deployable in the context of this evaluation.
154            kwargs: Additional kwargs to pass to the renderer.
155        """
156        validate_date_range(start, end)
157
158        snapshots = {
159            self.snapshots[p_sid].name: self.snapshots[p_sid] for p_sid in snapshot.parents
160        }
161        snapshots[snapshot.name] = snapshot
162
163        if isinstance(snapshot.node, SeedModel) and not snapshot.node.is_hydrated:
164            snapshot = self.state_sync.get_snapshots([snapshot], hydrate_seeds=True)[
165                snapshot.snapshot_id
166            ]
167
168        is_deployable = deployability_index.is_deployable(snapshot)
169
170        wap_id = self.snapshot_evaluator.evaluate(
171            snapshot,
172            start=start,
173            end=end,
174            execution_time=execution_time,
175            snapshots=snapshots,
176            deployability_index=deployability_index,
177            **kwargs,
178        )
179        try:
180            self.snapshot_evaluator.audit(
181                snapshot=snapshot,
182                start=start,
183                end=end,
184                execution_time=execution_time,
185                snapshots=snapshots,
186                deployability_index=deployability_index,
187                wap_id=wap_id,
188                **kwargs,
189            )
190        except AuditError as e:
191            self.notification_target_manager.notify(NotificationEvent.AUDIT_FAILURE, e)
192            if is_deployable and snapshot.node.owner:
193                self.notification_target_manager.notify_user(
194                    NotificationEvent.AUDIT_FAILURE, snapshot.node.owner, e
195                )
196            logger.error(f"Audit Failure: {traceback.format_exc()}")
197            raise e
198
199        self.state_sync.add_interval(snapshot, start, end, is_dev=not is_deployable)
200
201    def run(
202        self,
203        environment: str | EnvironmentNamingInfo,
204        start: t.Optional[TimeLike] = None,
205        end: t.Optional[TimeLike] = None,
206        execution_time: t.Optional[TimeLike] = None,
207        restatements: t.Optional[t.Dict[SnapshotId, SnapshotInterval]] = None,
208        ignore_cron: bool = False,
209        end_bounded: bool = False,
210        selected_snapshots: t.Optional[t.Set[str]] = None,
211        circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
212        deployability_index: t.Optional[DeployabilityIndex] = None,
213    ) -> bool:
214        """Concurrently runs all snapshots in topological order.
215
216        Args:
217            environment: The environment naming info the user is targeting when applying their change.
218                Can just be the environment name if the user is targeting a remote environment and wants to get the remote
219                naming info
220            start: The start of the run. Defaults to the min node start date.
221            end: The end of the run. Defaults to now.
222            execution_time: The date/time time reference to use for execution time. Defaults to now.
223            restatements: A dict of snapshots to restate and their intervals.
224            ignore_cron: Whether to ignore the node's cron schedule.
225            end_bounded: If set to true, the evaluated intervals will be bounded by the target end date, disregarding lookback,
226                allow_partials, and other attributes that could cause the intervals to exceed the target end date.
227            selected_snapshots: A set of snapshot names to run. If not provided, all snapshots will be run.
228            circuit_breaker: An optional handler which checks if the run should be aborted.
229            deployability_index: Determines snapshots that are deployable in the context of this render.
230
231        Returns:
232            True if the execution was successful and False otherwise.
233        """
234        restatements = restatements or {}
235        validate_date_range(start, end)
236        if isinstance(environment, str):
237            env = self.state_sync.get_environment(environment)
238            if not env:
239                raise SQLMeshError(
240                    "Was not provided an environment suffix target and the environment doesn't exist."
241                    "Are you running for the first time and need to run plan/apply first?"
242                )
243            environment_naming_info = env.naming_info
244        else:
245            environment_naming_info = environment
246
247        deployability_index = deployability_index or (
248            DeployabilityIndex.create(self.snapshots.values())
249            if environment_naming_info.name != c.PROD
250            else DeployabilityIndex.all_deployable()
251        )
252        execution_time = execution_time or now()
253        batches = self.batches(
254            start,
255            end,
256            execution_time,
257            deployability_index=deployability_index,
258            restatements=restatements,
259            ignore_cron=ignore_cron,
260            end_bounded=end_bounded,
261            selected_snapshots=selected_snapshots,
262        )
263        if not batches:
264            return True
265
266        dag = self._dag(batches)
267
268        self.console.start_evaluation_progress(
269            {snapshot: len(intervals) for snapshot, intervals in batches.items()},
270            environment_naming_info,
271            self.default_catalog,
272        )
273
274        snapshots_by_name = {snapshot.name: snapshot for snapshot in self.snapshots.values()}
275
276        def evaluate_node(node: SchedulingUnit) -> None:
277            if circuit_breaker and circuit_breaker():
278                raise CircuitBreakerError()
279
280            snapshot_name, ((start, end), batch_idx) = node
281            if batch_idx == -1:
282                return
283            snapshot = snapshots_by_name[snapshot_name]
284
285            self.console.start_snapshot_evaluation_progress(snapshot)
286
287            execution_start_ts = now_timestamp()
288            evaluation_duration_ms: t.Optional[int] = None
289
290            try:
291                assert execution_time  # mypy
292                assert deployability_index  # mypy
293                self.evaluate(snapshot, start, end, execution_time, deployability_index)
294                evaluation_duration_ms = now_timestamp() - execution_start_ts
295            finally:
296                self.console.update_snapshot_evaluation_progress(
297                    snapshot, batch_idx, evaluation_duration_ms
298                )
299
300        try:
301            with self.snapshot_evaluator.concurrent_context():
302                errors, skipped_intervals = concurrent_apply_to_dag(
303                    dag,
304                    evaluate_node,
305                    self.max_workers,
306                    raise_on_error=False,
307                )
308        finally:
309            self.state_sync.recycle()
310
311        self.console.stop_evaluation_progress(success=not errors)
312
313        skipped_snapshots = {i[0] for i in skipped_intervals}
314        for skipped in skipped_snapshots:
315            log_message = f"SKIPPED snapshot {skipped}\n"
316            self.console.log_status_update(log_message)
317            logger.info(log_message)
318
319        for error in errors:
320            if isinstance(error.__cause__, CircuitBreakerError):
321                raise error.__cause__
322            sid = error.node[0]
323            formatted_exception = "".join(format_exception(error.__cause__ or error))
324            log_message = f"FAILED processing snapshot {sid}\n{formatted_exception}"
325            self.console.log_error(log_message)
326            # Log with INFO level to prevent duplicate messages in the console.
327            logger.info(log_message)
328
329        return not errors
330
331    def _dag(self, batches: SnapshotToBatches) -> DAG[SchedulingUnit]:
332        """Builds a DAG of snapshot intervals to be evaluated.
333
334        Args:
335            batches: The batches of snapshots and intervals to evaluate.
336
337        Returns:
338            A DAG of snapshot intervals to be evaluated.
339        """
340
341        intervals_per_snapshot = {
342            snapshot.name: intervals for snapshot, intervals in batches.items()
343        }
344
345        dag = DAG[SchedulingUnit]()
346        terminal_node = ((to_datetime(0), to_datetime(0)), -1)
347
348        for snapshot, intervals in batches.items():
349            if not intervals:
350                continue
351
352            upstream_dependencies = []
353
354            for p_sid in snapshot.parents:
355                if p_sid in self.snapshots:
356                    p_intervals = intervals_per_snapshot.get(p_sid.name, [])
357
358                    if len(p_intervals) > 1:
359                        upstream_dependencies.append((p_sid.name, terminal_node))
360                    else:
361                        for i, interval in enumerate(p_intervals):
362                            upstream_dependencies.append((p_sid.name, (interval, i)))
363
364            batch_concurrency = snapshot.node.batch_concurrency
365            if snapshot.depends_on_past:
366                batch_concurrency = 1
367
368            for i, interval in enumerate(intervals):
369                node = (snapshot.name, (interval, i))
370                dag.add(node, upstream_dependencies)
371
372                if len(intervals) > 1:
373                    dag.add((snapshot.name, terminal_node), [node])
374
375                if batch_concurrency and i >= batch_concurrency:
376                    batch_idx_to_wait_for = i - batch_concurrency
377                    dag.add(
378                        node,
379                        [
380                            (
381                                snapshot.name,
382                                (intervals[batch_idx_to_wait_for], batch_idx_to_wait_for),
383                            )
384                        ],
385                    )
386        return dag
387
388
389def compute_interval_params(
390    snapshots: t.Collection[Snapshot],
391    *,
392    start: TimeLike,
393    end: TimeLike,
394    deployability_index: t.Optional[DeployabilityIndex] = None,
395    execution_time: t.Optional[TimeLike] = None,
396    restatements: t.Optional[t.Dict[SnapshotId, SnapshotInterval]] = None,
397    ignore_cron: bool = False,
398    end_bounded: bool = False,
399) -> SnapshotToBatches:
400    """Find the optimal date interval paramaters based on what needs processing and maximal batch size.
401
402    For each node name, find all dependencies and look for a stored snapshot from the metastore. If a snapshot is found,
403    calculate the missing intervals that need to be processed given the passed in start and end intervals.
404
405    If a snapshot's node specifies a batch size, consecutive intervals are merged into batches of a size that is less than
406    or equal to the configured one. If no batch size is specified, then it uses the intervals that correspond to the node's cron expression.
407    For example, if a node is supposed to run daily and has 70 days to backfill with a batch size set to 30, there would be 2 jobs
408    with 30 days and 1 job with 10.
409
410    Args:
411        snapshots: A set of target snapshots for which intervals should be computed.
412        intervals: A list of all snapshot intervals that should be considered.
413        start: Start of the interval.
414        end: End of the interval.
415        deployability_index: Determines snapshots that are deployable in the context of this evaluation.
416        execution_time: The date/time time reference to use for execution time.
417        restatements: A dict of snapshot names being restated and their intervals.
418        ignore_cron: Whether to ignore the node's cron schedule.
419        end_bounded: If set to true, the returned intervals will be bounded by the target end date, disregarding lookback,
420            allow_partials, and other attributes that could cause the intervals to exceed the target end date.
421
422    Returns:
423        A dict containing all snapshots needing to be run with their associated interval params.
424    """
425    snapshot_batches = {}
426
427    for snapshot, intervals in missing_intervals(
428        snapshots,
429        start=start,
430        end=end,
431        execution_time=execution_time,
432        restatements=restatements,
433        deployability_index=deployability_index,
434        ignore_cron=ignore_cron,
435        end_bounded=end_bounded,
436    ).items():
437        batches = []
438        batch_size = snapshot.node.batch_size
439        next_batch: t.List[t.Tuple[int, int]] = []
440
441        for interval in intervals:
442            if (batch_size and len(next_batch) >= batch_size) or (
443                next_batch and interval[0] != next_batch[-1][-1]
444            ):
445                batches.append((next_batch[0][0], next_batch[-1][-1]))
446                next_batch = []
447            next_batch.append(interval)
448        if next_batch:
449            batches.append((next_batch[0][0], next_batch[-1][-1]))
450        snapshot_batches[snapshot] = [(to_datetime(s), to_datetime(e)) for s, e in batches]
451
452    return snapshot_batches
453
454
455def _resolve_one_snapshot_per_version(
456    snapshots: t.Iterable[Snapshot],
457) -> t.Dict[t.Tuple[str, str], Snapshot]:
458    snapshot_per_version: t.Dict[t.Tuple[str, str], Snapshot] = {}
459    for snapshot in snapshots:
460        key = (snapshot.name, snapshot.version_get_or_generate())
461        if key not in snapshot_per_version:
462            snapshot_per_version[key] = snapshot
463        else:
464            prev_snapshot = snapshot_per_version[key]
465            if snapshot.unpaused_ts and (
466                not prev_snapshot.unpaused_ts or snapshot.created_ts > prev_snapshot.created_ts
467            ):
468                snapshot_per_version[key] = snapshot
469
470    return snapshot_per_version