Skip to content

Commit bfa3eca

Browse files
treysptreysp
committed
List all upstream auto-restated models
1 parent c2994c7 commit bfa3eca

5 files changed

Lines changed: 128 additions & 29 deletions

File tree

‎sqlmesh/core/console.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -428,7 +428,7 @@ def update_snapshot_evaluation_progress(
428428
num_audits_passed: int,
429429
num_audits_failed: int,
430430
audit_only: bool = False,
431-
auto_restatement_trigger: t.Optional[SnapshotId] = None,
431+
auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
432432
) -> None:
433433
"""Updates the snapshot evaluation progress."""
434434

@@ -576,7 +576,7 @@ def update_snapshot_evaluation_progress(
576576
num_audits_passed: int,
577577
num_audits_failed: int,
578578
audit_only: bool = False,
579-
auto_restatement_trigger: t.Optional[SnapshotId] = None,
579+
auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
580580
) -> None:
581581
pass
582582

@@ -1058,7 +1058,7 @@ def update_snapshot_evaluation_progress(
10581058
num_audits_passed: int,
10591059
num_audits_failed: int,
10601060
audit_only: bool = False,
1061-
auto_restatement_trigger: t.Optional[SnapshotId] = None,
1061+
auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
10621062
) -> None:
10631063
"""Update the snapshot evaluation progress."""
10641064
if (
@@ -3642,7 +3642,7 @@ def update_snapshot_evaluation_progress(
36423642
num_audits_passed: int,
36433643
num_audits_failed: int,
36443644
audit_only: bool = False,
3645-
auto_restatement_trigger: t.Optional[SnapshotId] = None,
3645+
auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
36463646
) -> None:
36473647
view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
36483648

@@ -3812,12 +3812,12 @@ def update_snapshot_evaluation_progress(
38123812
num_audits_passed: int,
38133813
num_audits_failed: int,
38143814
audit_only: bool = False,
3815-
auto_restatement_trigger: t.Optional[SnapshotId] = None,
3815+
auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
38163816
) -> None:
38173817
message = f"Evaluating {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
38183818

3819-
if auto_restatement_trigger:
3820-
message += f" | evaluation_triggered_by={auto_restatement_trigger.name}"
3819+
if auto_restatement_triggers:
3820+
message += f" | auto_restatement_triggers={','.join(trigger.name for trigger in auto_restatement_triggers)}"
38213821

38223822
if audit_only:
38233823
message = f"Auditing {snapshot.name} duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"

‎sqlmesh/core/scheduler.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -415,7 +415,7 @@ def run_merged_intervals(
415415
selected_snapshot_ids: t.Optional[t.Set[SnapshotId]] = None,
416416
run_environment_statements: bool = False,
417417
audit_only: bool = False,
418-
auto_restatement_triggers: t.Dict[SnapshotId, SnapshotId] = {},
418+
auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {},
419419
) -> t.Tuple[t.List[NodeExecutionFailedError[SchedulingUnit]], t.List[SchedulingUnit]]:
420420
"""Runs precomputed batches of missing intervals.
421421
@@ -532,7 +532,7 @@ def run_node(node: SchedulingUnit) -> None:
532532
evaluation_duration_ms,
533533
num_audits - num_audits_failed,
534534
num_audits_failed,
535-
auto_restatement_trigger=auto_restatement_triggers.get(snapshot.snapshot_id),
535+
auto_restatement_triggers=auto_restatement_triggers.get(snapshot.snapshot_id),
536536
)
537537
elif isinstance(node, CreateNode):
538538
self.snapshot_evaluator.create_snapshot(
@@ -738,7 +738,7 @@ def _run_or_audit(
738738
for s_id, interval in (remove_intervals or {}).items():
739739
self.snapshots[s_id].remove_interval(interval)
740740

741-
auto_restatement_triggers: t.Dict[SnapshotId, SnapshotId] = {}
741+
auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
742742
if auto_restatement_enabled:
743743
auto_restated_intervals, auto_restatement_triggers = apply_auto_restatements(
744744
self.snapshots, execution_time

‎sqlmesh/core/snapshot/definition.py‎

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2180,7 +2180,7 @@ def snapshots_to_dag(snapshots: t.Collection[Snapshot]) -> DAG[SnapshotId]:
21802180

21812181
def apply_auto_restatements(
21822182
snapshots: t.Dict[SnapshotId, Snapshot], execution_time: TimeLike
2183-
) -> t.Tuple[t.List[SnapshotIntervals], t.Dict[SnapshotId, SnapshotId]]:
2183+
) -> t.Tuple[t.List[SnapshotIntervals], t.Dict[SnapshotId, t.List[SnapshotId]]]:
21842184
"""Applies auto restatements to the snapshots.
21852185
21862186
This operation results in the removal of intervals for snapshots that are ready to be restated based
@@ -2195,8 +2195,7 @@ def apply_auto_restatements(
21952195
A list of SnapshotIntervals with **new** intervals that need to be restated.
21962196
"""
21972197
dag = snapshots_to_dag(snapshots.values())
2198-
snapshots_with_auto_restatements: t.List[SnapshotId] = []
2199-
auto_restatement_triggers: t.Dict[SnapshotId, SnapshotId] = {}
2198+
auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
22002199
auto_restated_intervals_per_snapshot: t.Dict[SnapshotId, Interval] = {}
22012200
for s_id in dag:
22022201
if s_id not in snapshots:
@@ -2211,6 +2210,7 @@ def apply_auto_restatements(
22112210
for parent_s_id in snapshot.parents
22122211
if parent_s_id in auto_restated_intervals_per_snapshot
22132212
]
2213+
upstream_triggers = []
22142214
if next_auto_restated_interval:
22152215
logger.info(
22162216
"Calculated the next auto restated interval (%s, %s) for snapshot %s",
@@ -2221,21 +2221,15 @@ def apply_auto_restatements(
22212221
auto_restated_intervals.append(next_auto_restated_interval)
22222222

22232223
# auto-restated snapshot is its own trigger
2224-
snapshots_with_auto_restatements.append(s_id)
2225-
auto_restatement_triggers[s_id] = s_id
2226-
else:
2227-
for parent_s_id in snapshot.parents:
2228-
# first auto-restated parent is the trigger
2229-
if parent_s_id in snapshots_with_auto_restatements:
2230-
auto_restatement_triggers[s_id] = parent_s_id
2231-
break
2232-
# if no trigger yet and parent has trigger, inherit their trigger
2233-
# - will be overwritten if a different parent is auto-restated
2234-
if (
2235-
parent_s_id in auto_restatement_triggers
2236-
and s_id not in auto_restatement_triggers
2237-
):
2238-
auto_restatement_triggers[s_id] = auto_restatement_triggers[parent_s_id]
2224+
upstream_triggers = [s_id]
2225+
2226+
for parent_s_id in snapshot.parents:
2227+
if parent_s_id in auto_restatement_triggers:
2228+
upstream_triggers.extend(auto_restatement_triggers[parent_s_id])
2229+
2230+
# remove duplicate triggers
2231+
if upstream_triggers:
2232+
auto_restatement_triggers[s_id] = list(dict.fromkeys(upstream_triggers))
22392233

22402234
if auto_restated_intervals:
22412235
auto_restated_interval_start = sys.maxsize

‎tests/core/test_snapshot.py‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3279,6 +3279,111 @@ def test_apply_auto_restatements_disable_restatement_downstream(make_snapshot):
32793279
]
32803280

32813281

3282+
def test_auto_restatement_triggers(make_snapshot):
3283+
model_a = SqlModel(
3284+
name="test_model_a",
3285+
kind=IncrementalByTimeRangeKind(
3286+
time_column=TimeColumn(column="ds"),
3287+
auto_restatement_cron="0 10 * * *",
3288+
auto_restatement_intervals=24,
3289+
),
3290+
start="2020-01-01",
3291+
cron="@daily",
3292+
query=parse_one("SELECT 1 as ds"),
3293+
)
3294+
snapshot_a = make_snapshot(model_a, version="1")
3295+
snapshot_a.add_interval("2020-01-01", "2020-01-05")
3296+
snapshot_a.next_auto_restatement_ts = to_timestamp("2020-01-06 10:00:00")
3297+
3298+
model_b = SqlModel(
3299+
name="test_model_b",
3300+
kind=IncrementalByTimeRangeKind(
3301+
time_column=TimeColumn(column="ds"),
3302+
),
3303+
start="2020-01-01",
3304+
cron="@daily",
3305+
query=parse_one("SELECT ds FROM test_model_a"),
3306+
)
3307+
snapshot_b = make_snapshot(model_b, nodes={model_a.fqn: model_a}, version="1")
3308+
snapshot_b.add_interval("2020-01-01", "2020-01-05")
3309+
3310+
model_c = SqlModel(
3311+
name="test_model_c",
3312+
kind=IncrementalByTimeRangeKind(
3313+
time_column=TimeColumn(column="ds"),
3314+
auto_restatement_cron="0 10 * * *",
3315+
auto_restatement_intervals=24,
3316+
),
3317+
start="2020-01-01",
3318+
cron="@daily",
3319+
query=parse_one("SELECT ds FROM test_model_a"),
3320+
)
3321+
snapshot_c = make_snapshot(model_c, nodes={model_a.fqn: model_a}, version="1")
3322+
snapshot_c.add_interval("2020-01-01", "2020-01-05")
3323+
snapshot_c.next_auto_restatement_ts = to_timestamp("2020-01-06 10:00:00")
3324+
3325+
model_d = SqlModel(
3326+
name="test_model_d",
3327+
kind=IncrementalByTimeRangeKind(
3328+
time_column=TimeColumn(column="ds"),
3329+
auto_restatement_cron="0 10 * * *",
3330+
auto_restatement_intervals=24,
3331+
),
3332+
start="2020-01-01",
3333+
cron="@daily",
3334+
query=parse_one("SELECT 1 as ds"),
3335+
)
3336+
snapshot_d = make_snapshot(model_d, version="1")
3337+
snapshot_d.add_interval("2020-01-01", "2020-01-05")
3338+
snapshot_d.next_auto_restatement_ts = to_timestamp("2020-01-06 10:00:00")
3339+
3340+
model_e = SqlModel(
3341+
name="test_model_e",
3342+
kind=IncrementalByTimeRangeKind(
3343+
time_column=TimeColumn(column="ds"),
3344+
),
3345+
start="2020-01-01",
3346+
cron="@daily",
3347+
query=parse_one(
3348+
"SELECT ds from test_model_b UNION ALL SELECT ds from test_model_c UNION ALL SELECT ds from test_model_d"
3349+
),
3350+
)
3351+
snapshot_e = make_snapshot(
3352+
model_e,
3353+
nodes={
3354+
model_a.fqn: model_a,
3355+
model_b.fqn: model_b,
3356+
model_c.fqn: model_c,
3357+
model_d.fqn: model_d,
3358+
},
3359+
version="1",
3360+
)
3361+
snapshot_e.add_interval("2020-01-01", "2020-01-05")
3362+
3363+
_, auto_restatement_triggers = apply_auto_restatements(
3364+
{
3365+
snapshot_a.snapshot_id: snapshot_a,
3366+
snapshot_b.snapshot_id: snapshot_b,
3367+
snapshot_c.snapshot_id: snapshot_c,
3368+
snapshot_d.snapshot_id: snapshot_d,
3369+
snapshot_e.snapshot_id: snapshot_e,
3370+
},
3371+
"2020-01-06 10:01:00",
3372+
)
3373+
3374+
assert auto_restatement_triggers == {
3375+
snapshot_a.snapshot_id: [snapshot_a.snapshot_id],
3376+
snapshot_d.snapshot_id: [snapshot_d.snapshot_id],
3377+
snapshot_b.snapshot_id: [snapshot_a.snapshot_id],
3378+
snapshot_c.snapshot_id: [snapshot_c.snapshot_id, snapshot_a.snapshot_id],
3379+
snapshot_e.snapshot_id: [
3380+
snapshot_d.snapshot_id,
3381+
snapshot_c.snapshot_id,
3382+
snapshot_a.snapshot_id,
3383+
],
3384+
}
3385+
3386+
32823387
def test_render_signal(make_snapshot, mocker):
32833388
@signal()
32843389
def check_types(batch, env: str, sql: list[SQL], table: exp.Table, default: int = 0):

‎web/server/console.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ def update_snapshot_evaluation_progress(
142142
num_audits_passed: int,
143143
num_audits_failed: int,
144144
audit_only: bool = False,
145-
auto_restatement_trigger: t.Optional[SnapshotId] = None,
145+
auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
146146
) -> None:
147147
if audit_only:
148148
return

0 commit comments

Comments
 (0)