Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/events/triggers.py: 83%

1""" 

2The triggers consumer watches events streaming in from the event bus and decides whether 

3to act on them based on the automations that users have set up. 

4""" 

5 

6import asyncio 

7import fnmatch 

8from contextlib import AsyncExitStack, asynccontextmanager 

9from datetime import timedelta 

10from typing import ( 

11 TYPE_CHECKING, 

12 AsyncGenerator, 

13 Collection, 

14 Dict, 

15 List, 

16 Optional, 

17 Set, 

18 Tuple, 

19) 

20from uuid import UUID 

21 

22import orjson 

23import sqlalchemy as sa 

24from sqlalchemy.ext.asyncio import AsyncSession 

25from typing_extensions import Literal, TypeAlias 

26 

27import prefect.types._datetime 

28from prefect._internal.retries import retry_async_fn 

29from prefect.logging import get_logger 

30from prefect.server.database import PrefectDBInterface, db_injector 

31from prefect.server.events import messaging 

32from prefect.server.events.actions import ServerActionTypes 

33from prefect.server.events.models.automations import ( 

34 AUTOMATION_CHANGES_CHANNEL, 

35 AutomationChangeEvent, 

36 automations_session, 

37 read_automation, 

38) 

39from prefect.server.events.models.composite_trigger_child_firing import ( 

40 acquire_composite_trigger_lock, 

41 clear_child_firings, 

42 clear_old_child_firings, 

43 get_child_firings, 

44 upsert_child_firing, 

45) 

46from prefect.server.events.ordering import ( 

47 PRECEDING_EVENT_LOOKBACK, 

48 EventArrivedEarly, 

49 get_triggers_causal_ordering, 

50) 

51from prefect.server.events.schemas.automations import ( 

52 Automation, 

53 CompositeTrigger, 

54 EventTrigger, 

55 Firing, 

56 Posture, 

57 Trigger, 

58 TriggeredAction, 

59 TriggerState, 

60) 

61from prefect.server.events.schemas.events import ReceivedEvent 

62from prefect.server.utilities.messaging import Message, MessageHandler 

63from prefect.server.utilities.postgres_listener import ( 

64 get_pg_notify_connection, 

65 pg_listen, 

66) 

67from prefect.settings import PREFECT_EVENTS_EXPIRED_BUCKET_BUFFER 

68from prefect.settings.context import get_current_settings 

69 

70if TYPE_CHECKING: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 import logging 

72 

73 from prefect.server.database.orm_models import ORMAutomationBucket 

74 

75 

76logger: "logging.Logger" = get_logger(__name__) 

77 

78AutomationID: TypeAlias = UUID 

79TriggerID: TypeAlias = UUID 

80AutomationStateSnapshot: TypeAlias = Tuple[ 

81 Optional[prefect.types._datetime.DateTime], int 

82] 

83 

84 

85AUTOMATION_BUCKET_BATCH_SIZE = 500 

86 

87 

88async def evaluate( 

89 session: AsyncSession, 

90 trigger: EventTrigger, 

91 bucket: "ORMAutomationBucket", 

92 now: prefect.types._datetime.DateTime, 

93 triggering_event: Optional[ReceivedEvent], 

94) -> "ORMAutomationBucket | None": 

95 """Evaluates an Automation, either triggered by a specific event or proactively 

96 on a time interval. Evaluating a Automation updates the associated counters for 

97 each automation, and will fire the associated action if it has met the threshold.""" 

98 automation = trigger.automation 

99 

100 logging_context = { 

101 "automation": automation.id, 

102 "trigger": trigger.id, 

103 "bucketing_key": bucket.bucketing_key, 

104 "bucket_start": bucket.start, 

105 "bucket_end": bucket.end, 

106 "bucket_initial_count": bucket.count, 

107 "bucket_last_operation": bucket.last_operation, 

108 "now": now, 

109 "triggering_event": ( 

110 { 

111 "id": triggering_event.id, 

112 "event": triggering_event.event, 

113 } 

114 if triggering_event 

115 else None 

116 ), 

117 } 

118 

119 # Implementation notes: 

120 # 

121 # This triggering algorithm maintains an invariant that there is exactly one 

122 # time-based "bucket" open and collecting events for each automation at a time. When 

123 # an event comes in that matches the automation, one of four things can happen: 

124 # 

125 # 1. The event would have matched an older bucket that has either expired or has 

126 # already filled up, and thus is no longer relevant; 

127 # 2. The event matches the current bucket, but the bucket does not meet its 

128 # threshold yet; 

129 # 3. The event matches the current bucket, causes it to meet the the threshold, so 

130 # we fire immediately and advance the bucket to the next time period. 

131 # 4. The event matches the current bucket, but the event is for a future time after 

132 # the current bucket has expired, so we will start the new bucket and re-evaluate 

133 # 

134 # Automations are also evaluated proactively without an event to see if they have 

135 # met their proactive threshold (where not enough events have happened in the time 

136 # period) 

137 

138 # If there was a triggering event, then we need to "spend" this count somewhere, 

139 # either in the currently open bucket or in the next time period's bucket 

140 count = 1 if triggering_event else 0 

141 

142 if now < bucket.start: 

143 # This is an older out-of-order message or a redundant event for a reactive 

144 # trigger that has has already fired, so it should not should not affect the 

145 # current bucket. We can safely ignore this event/timestamp entirely. Case #1 

146 # from the implementation notes above. 

147 logger.debug( 

148 "Automation %s (%r) trigger %s got a late event for keys (%r)", 

149 automation.id, 

150 automation.name, 

151 trigger.id, 

152 bucket.bucketing_key, 

153 extra=logging_context, 

154 ) 

155 return bucket 

156 

157 if count and (trigger.immediate or bucket.start <= now < bucket.end): 

158 # we are still within the automation time period, so spend the count in the 

159 # current bucket 

160 bucket = await increment_bucket(session, bucket, count, triggering_event) 

161 count = 0 

162 

163 # Reactive automations will fire "eagerly", meaning they will fire as _soon_ as the 

164 # threshold is reached, then wait for the rest of their time period before 

165 # firing again. This is done by creating a new bucket in the future after the 

166 # trigger has fired. 

167 

168 # Proactive automations must wait until the whole bucket has expired before they can 

169 # fire, because we can't know if we'll get one late event just in time to cause the 

170 # automation _not_ to fire. 

171 

172 ready_to_fire = trigger.posture == Posture.Reactive or bucket.end <= now 

173 meets_threshold = trigger.meets_threshold(bucket.count) 

174 

175 if ready_to_fire and meets_threshold: 

176 logger.debug( 

177 ( 

178 "Automation %s (%r) trigger %s triggered for keys (%r) %s, " 

179 "having occurred %s times between %s and %s" 

180 ), 

181 automation.id, 

182 automation.name, 

183 trigger.id, 

184 bucket.bucketing_key, 

185 "reactively" if triggering_event else "proactively", 

186 bucket.count, 

187 bucket.start, 

188 bucket.end, 

189 extra=logging_context, 

190 ) 

191 

192 firing = Firing( 

193 trigger=trigger, 

194 trigger_states={TriggerState.Triggered}, 

195 triggered=prefect.types._datetime.now("UTC"), 

196 triggering_labels={ 

197 label: value 

198 for label, value in zip(sorted(trigger.for_each), bucket.bucketing_key) 

199 }, 

200 triggering_event=triggering_event or bucket.last_event, 

201 ) 

202 

203 await fire(session, firing) 

204 

205 # when acting, remove the current bucket from the database immediately to avoid 

206 # potentially double-acting in the case of a crash between now and the next 

207 # time we backup buckets to the database 

208 await remove_bucket(session, bucket) 

209 

210 elif now < bucket.end: 210 ↛ 255line 210 didn't jump to line 255 because the condition on line 210 was always true

211 # We didn't fire this time, and also the bucket still has more time, so leave 

212 # before setting up a new future bucket. Case #2 from the implementation notes 

213 # above. 

214 logger.debug( 

215 "Automation %s (%r) trigger %s has more time for keys (%r)", 

216 automation.id, 

217 automation.name, 

218 trigger.id, 

219 bucket.bucketing_key, 

220 extra={ 

221 **logging_context, 

222 "ready_to_fire": ready_to_fire, 

223 "meets_threshold": meets_threshold, 

224 }, 

225 ) 

226 

227 # Special case of a proactive trigger for which the same event satisfies 

228 # both `after` and `expect`. For example, using flow run heartbeats for crash 

229 # detection, after the first heartbeat we expect a subsequent heartbeat or 

230 # terminal state within a given time. 

231 # 

232 # If we've already reached the proactive threshold, we need to remove the 

233 # current bucket and start a new one for the latest event. 

234 if ( 234 ↛ 241line 234 didn't jump to line 241 because the condition on line 234 was never true

235 triggering_event 

236 and trigger.posture == Posture.Proactive 

237 and not meets_threshold 

238 and trigger.starts_after(triggering_event.event) 

239 and trigger.expects(triggering_event.event) 

240 ): 

241 await remove_bucket(session, bucket) 

242 return await start_new_bucket( 

243 session, 

244 trigger, 

245 bucketing_key=bucket.bucketing_key, 

246 start=triggering_event.occurred, 

247 end=triggering_event.occurred + trigger.within, 

248 count=0, 

249 last_event=triggering_event, 

250 ) 

251 

252 return bucket 

253 else: 

254 # Case #2 from the implementation notes above. 

255 logger.debug( 

256 "Automation %s (%r) trigger %s not ready to fire for keys (%r)", 

257 automation.id, 

258 automation.name, 

259 trigger.id, 

260 bucket.bucketing_key, 

261 extra={ 

262 **logging_context, 

263 "ready_to_fire": ready_to_fire, 

264 "meets_threshold": meets_threshold, 

265 "bucket_current_count": bucket.count, 

266 }, 

267 ) 

268 

269 # We are now outside of the automation's time period or we triggered this 

270 # time. That means it's time to start a new bucket for the next possible time 

271 # window (if this automation does not require an event to start `after`): 

272 if trigger.after: 272 ↛ 275line 272 didn't jump to line 275 because the condition on line 272 was never true

273 # remove the bucket because it should only get re-created if we see another 

274 # appropriate starting event 

275 return await remove_bucket(session, bucket) 

276 else: 

277 if trigger.within == timedelta(seconds=0): 277 ↛ 280line 277 didn't jump to line 280 because the condition on line 277 was always true

278 return None 

279 

280 start = prefect.types._datetime.create_datetime_instance(max(bucket.end, now)) 

281 end = start + trigger.within 

282 

283 # If we're processing a reactive trigger and leaving the function with a count 

284 # that we've just spent in the bucket for the next time window, it means that we 

285 # just processed an event that was in the future. It's possible that this event 

286 # was sufficient enough to cause the trigger to fire, so we need to evaluate one 

287 # more time to see if that's the case. This is case #4 from the implementation 

288 # notes above. 

289 if trigger.posture == Posture.Reactive and count > 0: 

290 bucket = await start_new_bucket( 

291 session, 

292 trigger, 

293 bucketing_key=tuple(bucket.bucketing_key), 

294 start=start, 

295 end=end, 

296 count=0, 

297 ) 

298 return await evaluate(session, trigger, bucket, now, triggering_event) 

299 else: 

300 return await start_new_bucket( 

301 session, 

302 trigger, 

303 bucketing_key=tuple(bucket.bucketing_key), 

304 start=start, 

305 end=end, 

306 count=count, 

307 ) 

308 

309 

310async def fire(session: AsyncSession, firing: Firing) -> None: 

311 if isinstance(firing.trigger.parent, Automation): 

312 await act(firing) 

313 elif isinstance(firing.trigger.parent, CompositeTrigger): 313 ↛ 316line 313 didn't jump to line 316 because the condition on line 313 was always true

314 await evaluate_composite_trigger(session, firing) 

315 else: 

316 raise NotImplementedError( 

317 f"Cannot fire {firing} with parent trigger type {type(firing.trigger.parent)}" 

318 ) 

319 

320 

321async def evaluate_composite_trigger(session: AsyncSession, firing: Firing) -> None: 

322 automation = firing.trigger.automation 

323 

324 assert isinstance(firing.trigger.parent, CompositeTrigger) 

325 trigger: CompositeTrigger = firing.trigger.parent 

326 

327 # If we only need to see 1 child firing, 

328 # then the parent trigger can fire immediately. 

329 if trigger.num_expected_firings == 1: 329 ↛ 358line 329 didn't jump to line 358 because the condition on line 329 was always true

330 logger.info( 

331 "Automation %s (%r) %s trigger %s fired (shortcut)", 

332 automation.id, 

333 automation.name, 

334 trigger.type, 

335 trigger.id, 

336 extra={ 

337 "automation": automation.id, 

338 "trigger": trigger.id, 

339 "trigger_type": trigger.type, 

340 "firings": firing.id, 

341 }, 

342 ) 

343 await fire( 

344 session, 

345 Firing( 

346 trigger=trigger, 

347 trigger_states={TriggerState.Triggered}, 

348 triggered=prefect.types._datetime.now("UTC"), 

349 triggering_firings=[firing], 

350 triggering_event=firing.triggering_event, 

351 ), 

352 ) 

353 return 

354 

355 # Acquire an advisory lock to serialize concurrent evaluations for this 

356 # compound trigger. This prevents a race condition where multiple child 

357 # triggers fire concurrently and neither transaction sees both firings. 

358 await acquire_composite_trigger_lock(session, trigger) 

359 

360 # If we're only looking within a certain time horizon, remove any older firings that 

361 # should no longer be considered as satisfying this trigger 

362 if trigger.within is not None: 

363 await clear_old_child_firings( 

364 session, trigger, firing.triggered - trigger.within 

365 ) 

366 

367 # Otherwise we need N child firings. We'll upsert this firing and then check 

368 # what the current state of the world is. If we have enough firings, we'll 

369 # fire the parent trigger. 

370 await upsert_child_firing(session, firing) 

371 firings: list[Firing] = [ 

372 cf.child_firing for cf in await get_child_firings(session, trigger) 

373 ] 

374 firing_ids: set[UUID] = {f.id for f in firings} 

375 

376 # If our current firing no longer exists when we read firings 

377 # another firing has superseded it, and we should defer to that one 

378 if firing.id not in firing_ids: 378 ↛ anywhereline 378 didn't jump anywhere: it always raised an exception.

379 return 

380 

381 if trigger.ready_to_fire(firings): 381 ↛ exitline 381 didn't return from function 'evaluate_composite_trigger' because the condition on line 381 was always true

382 logger.info( 

383 "Automation %s (%r) %s trigger %s fired", 

384 automation.id, 

385 automation.name, 

386 trigger.type, 

387 trigger.id, 

388 extra={ 

389 "automation": automation.id, 

390 "trigger": trigger.id, 

391 "trigger_type": trigger.type, 

392 "firings": ",".join(str(f.id) for f in firings), 

393 }, 

394 ) 

395 

396 # Clear by firing id, and only proceed if we won the race to claim them. 

397 # This prevents double-firing when multiple workers evaluate concurrently. 

398 deleted_ids = await clear_child_firings( 

399 session, trigger, firing_ids=list(firing_ids) 

400 ) 

401 

402 if deleted_ids != firing_ids: 

403 logger.debug( 

404 "Composite trigger %s skipped fire; expected to delete %s firings, " 

405 "actually deleted %s (another worker likely claimed them)", 

406 trigger.id, 

407 len(firing_ids), 

408 len(deleted_ids), 

409 extra={ 

410 "automation": automation.id, 

411 "trigger": trigger.id, 

412 "expected_firing_ids": sorted(str(f) for f in firing_ids), 

413 "deleted_firing_ids": sorted(str(f) for f in deleted_ids), 

414 }, 

415 ) 

416 return 

417 

418 await fire( 

419 session, 

420 Firing( 

421 trigger=trigger, 

422 trigger_states={TriggerState.Triggered}, 

423 triggered=prefect.types._datetime.now("UTC"), 

424 triggering_firings=firings, 

425 triggering_event=firing.triggering_event, 

426 ), 

427 ) 

428 

429 

430async def act(firing: Firing) -> None: 

431 """Given a Automation that has been triggered, the triggering labels and event 

432 (if there was one), publish an action for the `actions` service to process.""" 

433 automation = firing.trigger.automation 

434 

435 state_change_events: Dict[TriggerState, ReceivedEvent] = { 

436 trigger_state: firing.trigger.create_automation_state_change_event( 

437 firing=firing, 

438 trigger_state=trigger_state, 

439 ) 

440 for trigger_state in sorted(firing.trigger_states, key=list(TriggerState).index) 

441 } 

442 await messaging.publish(state_change_events.values()) 

443 

444 # Determine the primary state change event ID for linking action events back to 

445 # the automation.triggered or automation.resolved event. Prefer Triggered over 

446 # Resolved when both are present. 

447 primary_state_change_event = state_change_events.get( 

448 TriggerState.Triggered 

449 ) or state_change_events.get(TriggerState.Resolved) 

450 primary_state_change_event_id = ( 

451 primary_state_change_event.id if primary_state_change_event else None 

452 ) 

453 

454 # By default, all `automation.actions` are fired. Each tuple contains: 

455 # (triggering_event, action, automation_triggered_event_id) 

456 source_actions: List[ 

457 Tuple[Optional[ReceivedEvent], ServerActionTypes, UUID | None] 

458 ] = [ 

459 (firing.triggering_event, action, primary_state_change_event_id) 

460 for action in automation.actions 

461 ] 

462 

463 # Conditionally add in actions that fire on specific trigger states 

464 if TriggerState.Triggered in firing.trigger_states: 464 ↛ 471line 464 didn't jump to line 471 because the condition on line 464 was always true

465 triggered_event = state_change_events[TriggerState.Triggered] 

466 source_actions += [ 

467 (triggered_event, action, triggered_event.id) 

468 for action in automation.actions_on_trigger 

469 ] 

470 

471 if TriggerState.Resolved in firing.trigger_states: 471 ↛ 472line 471 didn't jump to line 472 because the condition on line 471 was never true

472 resolved_event = state_change_events[TriggerState.Resolved] 

473 source_actions += [ 

474 (resolved_event, action, resolved_event.id) 

475 for action in automation.actions_on_resolve 

476 ] 

477 

478 actions = [ 

479 TriggeredAction( 

480 automation=automation, 

481 firing=firing, 

482 triggered=firing.triggered, 

483 triggering_labels=firing.triggering_labels, 

484 triggering_event=action_triggering_event, 

485 action=action, 

486 action_index=index, 

487 automation_triggered_event_id=automation_triggered_event_id, 

488 ) 

489 for index, ( 

490 action_triggering_event, 

491 action, 

492 automation_triggered_event_id, 

493 ) in enumerate(source_actions) 

494 ] 

495 

496 async with messaging.create_actions_publisher() as publisher: 

497 for action in actions: 

498 await publisher.publish_data(action.model_dump_json().encode(), {}) 

499 

500 

501__events_clock_lock: Optional[asyncio.Lock] = None 

502_events_clock: Optional[float] = None 

503_events_clock_updated: Optional[float] = None 

504 

505 

506def _events_clock_lock() -> asyncio.Lock: 

507 global __events_clock_lock 

508 if __events_clock_lock is None: 

509 __events_clock_lock = asyncio.Lock() 

510 return __events_clock_lock 

511 

512 

513async def update_events_clock(event: ReceivedEvent) -> None: 

514 global _events_clock, _events_clock_updated 

515 async with _events_clock_lock(): 

516 # we want the offset to be negative to represent that we are always 

517 # processing events behind realtime... 

518 now = prefect.types._datetime.now("UTC").timestamp() 

519 event_timestamp = event.occurred.timestamp() 

520 offset = event_timestamp - now 

521 

522 # ...and we should clamp this value to zero so we don't inadvertently look like 

523 # we are processing the future 

524 if offset > 0.0: 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true

525 event_timestamp = now 

526 

527 if not _events_clock or event_timestamp >= _events_clock: 

528 _events_clock = event_timestamp 

529 

530 _events_clock_updated = now 

531 

532 

533async def get_events_clock() -> Optional[float]: 

534 global _events_clock 

535 return _events_clock 

536 

537 

538async def get_events_clock_offset() -> float: 

539 """Calculate the current clock offset. This takes into account both the `occurred` 

540 of the last event, as well as the time we _saw_ the last event. This helps to 

541 ensure that in low volume environments, we don't end up getting huge offsets.""" 

542 global _events_clock, _events_clock_updated 

543 

544 async with _events_clock_lock(): 

545 if _events_clock is None or _events_clock_updated is None: 

546 return 0.0 

547 

548 now: float = prefect.types._datetime.now("UTC").timestamp() 

549 offset = (_events_clock - now) + (now - _events_clock_updated) 

550 

551 return offset 

552 

553 

554async def reset_events_clock() -> None: 

555 global _events_clock, _events_clock_updated 

556 async with _events_clock_lock(): 

557 _events_clock = None 

558 _events_clock_updated = None 

559 

560 

561async def reactive_evaluation(event: ReceivedEvent, depth: int = 0) -> None: 

562 """ 

563 Evaluate all automations that may apply to this event. 

564 

565 Args: 

566 event (ReceivedEvent): The event to evaluate. This object contains all the necessary information 

567 about the event, including its type, associated resources, and metadata. 

568 depth (int, optional): The current recursion depth. This is used to prevent infinite recursion 

569 due to cyclic event dependencies. Defaults to 0 and is incremented with 

570 each recursive call. 

571 

572 """ 

573 async with AsyncExitStack() as stack: 

574 await update_events_clock(event) 

575 await stack.enter_async_context( 

576 get_triggers_causal_ordering().preceding_event_confirmed( 

577 reactive_evaluation, event, depth 

578 ) 

579 ) 

580 

581 interested_triggers = find_interested_triggers(event) 

582 if not interested_triggers: 

583 return 

584 

585 for trigger in interested_triggers: 

586 logger.info( 

587 "Automation %s, trigger %s covers event %r (%s) for %r at %r", 

588 trigger.automation.id, 

589 trigger.id, 

590 event.event, 

591 event.id, 

592 event.resource.id, 

593 event.occurred.isoformat(), 

594 ) 

595 

596 bucketing_key = trigger.bucketing_key(event) 

597 

598 async with automations_session(begin_transaction=True) as session: 

599 try: 

600 bucket: Optional["ORMAutomationBucket"] = None 

601 

602 if trigger.after and trigger.starts_after(event.event): 602 ↛ 618line 602 didn't jump to line 618 because the condition on line 602 was never true

603 # When an event matches both the after and expect, each event 

604 # can both start a new bucket and increment the bucket that was 

605 # started by the previous event. Here we offset the bucket to 

606 # start at -1 so that the first event will leave the bucket at 0 

607 # after evaluation. See the tests: 

608 # 

609 # test_same_event_in_expect_and_after_never_reacts_immediately 

610 # test_same_event_in_expect_and_after_reacts_after_threshold_is_met 

611 # test_same_event_in_expect_and_after_proactively_does_not_fire 

612 # test_same_event_in_expect_and_after_proactively_fires 

613 # 

614 # in test_triggers_regressions.py for examples of how we expect 

615 # this to behave. 

616 # 

617 # https://github.com/PrefectHQ/nebula/issues/4201 

618 initial_count = -1 if trigger.expects(event.event) else 0 

619 bucket = await ensure_bucket( 

620 session, 

621 trigger, 

622 bucketing_key, 

623 start=event.occurred, 

624 end=event.occurred + trigger.within, 

625 last_event=event, 

626 initial_count=initial_count, 

627 ) 

628 

629 if ( 

630 not bucket 

631 and not trigger.after 

632 and trigger.expects(event.event) 

633 ): 

634 # When ensuring a bucket and _creating it for the first time_, 

635 # use an old time so that we can catch any other events flowing 

636 # through the system at the same time even if they are out of 

637 # order. After the trigger fires and creates its next bucket, 

638 # time will start from that point forward. We'll use our 

639 # preceding event lookback variable as the horizon that we'll 

640 # accept these older events. 

641 # 

642 # https://github.com/PrefectHQ/nebula/issues/7230 

643 start = event.occurred - PRECEDING_EVENT_LOOKBACK 

644 

645 bucket = await ensure_bucket( 

646 session, 

647 trigger, 

648 bucketing_key=bucketing_key, 

649 start=start, 

650 end=event.occurred + trigger.within, 

651 last_event=event, 

652 ) 

653 

654 if not trigger.expects(event.event): 654 ↛ 655line 654 didn't jump to line 655 because the condition on line 654 was never true

655 continue 

656 

657 if not bucket: 

658 bucket = await read_bucket(session, trigger, bucketing_key) 

659 if not bucket: 659 ↛ 662line 659 didn't jump to line 662 because the condition on line 659 was always true

660 continue 

661 

662 await evaluate( 

663 session, 

664 trigger, 

665 bucket, 

666 event.occurred, 

667 triggering_event=event, 

668 ) 

669 finally: 

670 await session.commit() 

671 

672 

673# retry on operational errors to account for db flakiness with sqlite 

674@retry_async_fn(max_attempts=3, retry_on_exceptions=(sa.exc.OperationalError,)) 

675async def get_lost_followers() -> List[ReceivedEvent]: 

676 """Get followers that have been sitting around longer than our lookback""" 

677 return await get_triggers_causal_ordering().get_lost_followers() 

678 

679 

680async def periodic_evaluation(now: prefect.types._datetime.DateTime) -> None: 

681 """Periodic tasks that should be run regularly, but not as often as every event""" 

682 offset = await get_events_clock_offset() 

683 as_of = now + timedelta(seconds=offset) 

684 

685 logger.debug("Running periodic evaluation as of %s (offset %ss)", as_of, offset) 

686 

687 await reconcile_automations() 

688 

689 # Any followers that have been sitting around longer than our lookback are never 

690 # going to see their leader event (maybe it was lost or took too long to arrive). 

691 # These events can just be evaluated now in the order they occurred. 

692 for event in await get_lost_followers(): 692 ↛ 693line 692 didn't jump to line 693 because the loop on line 692 never started

693 await reactive_evaluation(event) 

694 

695 async with automations_session() as session: 

696 await sweep_closed_buckets( 

697 session, 

698 as_of - PREFECT_EVENTS_EXPIRED_BUCKET_BUFFER.value(), 

699 ) 

700 await session.commit() 

701 

702 

703async def evaluate_periodically(periodic_granularity: timedelta) -> None: 

704 """Runs periodic evaluation on the given interval""" 

705 logger.debug( 

706 "Starting periodic evaluation task every %s seconds", 

707 periodic_granularity.total_seconds(), 

708 ) 

709 while True: 

710 try: 

711 await periodic_evaluation(prefect.types._datetime.now("UTC")) 

712 except Exception: 

713 logger.exception("Error running periodic evaluation") 

714 finally: 

715 await asyncio.sleep(periodic_granularity.total_seconds()) 

716 

717 

718# The currently loaded automations for this shard, organized both by ID and by 

719# account and workspace 

720automations_by_id: Dict[UUID, Automation] = {} 

721triggers: Dict[TriggerID, EventTrigger] = {} 

722# Index from expect/after pattern → trigger IDs, so find_interested_triggers 

723# checks a handful of patterns instead of scanning all triggers. 

724_triggers_by_expect: Dict[str, Set[TriggerID]] = {} 

725next_proactive_runs: Dict[TriggerID, prefect.types._datetime.DateTime] = {} 

726automation_state_snapshot: Optional[AutomationStateSnapshot] = None 

727 

728# This lock governs any changes to the set of loaded automations; any routine that will 

729# add/remove automations must be holding this lock when it does so. It's best to use 

730# the methods below to access the loaded set of automations. 

731__automations_lock: Optional[asyncio.Lock] = None 

732 

733 

734def _automations_lock() -> asyncio.Lock: 

735 global __automations_lock 

736 if __automations_lock is None: 

737 __automations_lock = asyncio.Lock() 

738 return __automations_lock 

739 

740 

741def find_interested_triggers(event: ReceivedEvent) -> Collection[EventTrigger]: 

742 candidate_ids: set[TriggerID] = set() 

743 for expect_pattern, trigger_ids in _triggers_by_expect.items(): 

744 if fnmatch.fnmatchcase(event.event, expect_pattern): 

745 candidate_ids.update(trigger_ids) 

746 candidates = [triggers[tid] for tid in candidate_ids if tid in triggers] 

747 return [trigger for trigger in candidates if trigger.covers(event)] 

748 

749 

750def clear_loaded_automations() -> None: 

751 automations_by_id.clear() 

752 triggers.clear() 

753 _triggers_by_expect.clear() 

754 next_proactive_runs.clear() 

755 

756 

757def _index_keys_for(trigger: EventTrigger) -> set[str]: 

758 # An empty `expect` means "match any event" (event_pattern is `.+`), so the 

759 # trigger must be reachable for every event regardless of its `after` set. 

760 if not trigger.expect: 

761 return {"*"} 

762 return trigger.expect | trigger.after 

763 

764 

765def _index_trigger(trigger: EventTrigger) -> None: 

766 for key in _index_keys_for(trigger): 

767 _triggers_by_expect.setdefault(key, set()).add(trigger.id) 

768 

769 

770def _unindex_trigger(trigger: EventTrigger) -> None: 

771 for key in _index_keys_for(trigger): 

772 if key in _triggers_by_expect: 772 ↛ 771line 772 didn't jump to line 771 because the condition on line 772 was always true

773 _triggers_by_expect[key].discard(trigger.id) 

774 if not _triggers_by_expect[key]: 

775 del _triggers_by_expect[key] 

776 

777 

778def load_automation(automation: Optional[Automation]) -> None: 

779 """Loads the given automation into memory so that it is available for evaluations""" 

780 if not automation: 780 ↛ 781line 780 didn't jump to line 781 because the condition on line 780 was never true

781 return 

782 

783 event_triggers = automation.triggers_of_type(EventTrigger) 

784 

785 if not automation.enabled or not event_triggers: 

786 forget_automation(automation.id) 

787 return 

788 

789 automations_by_id[automation.id] = automation 

790 

791 for trigger in event_triggers: 

792 triggers[trigger.id] = trigger 

793 next_proactive_runs.pop(trigger.id, None) 

794 _index_trigger(trigger) 

795 

796 

797def forget_automation(automation_id: UUID) -> None: 

798 """Unloads the given automation from memory""" 

799 if automation := automations_by_id.pop(automation_id, None): 

800 for trigger in automation.triggers(): 

801 triggers.pop(trigger.id, None) 

802 next_proactive_runs.pop(trigger.id, None) 

803 for trigger in automation.triggers_of_type(EventTrigger): 

804 _unindex_trigger(trigger) 

805 

806 

807async def automation_changed( 

808 automation_id: UUID, 

809 event: Literal["automation__created", "automation__updated", "automation__deleted"], 

810) -> None: 

811 global automation_state_snapshot 

812 

813 async with _automations_lock(): 

814 if event in ("automation__deleted", "automation__updated"): 

815 forget_automation(automation_id) 

816 

817 async with automations_session() as session: 

818 if event in ("automation__created", "automation__updated"): 

819 automation = await read_automation(session, automation_id) 

820 load_automation(automation) 

821 automation_state_snapshot = await read_automation_state_snapshot(session) 

822 

823 

824@db_injector 

825async def load_automations(db: PrefectDBInterface, session: AsyncSession): 

826 """Loads all automations for the given set of accounts""" 

827 query = sa.select(db.Automation) 

828 

829 logger.debug("Loading automations") 

830 

831 result = await session.execute(query) 

832 for automation in result.scalars().all(): 

833 load_automation(Automation.model_validate(automation, from_attributes=True)) 

834 

835 logger.debug( 

836 "Loaded %s automations with %s triggers", len(automations_by_id), len(triggers) 

837 ) 

838 

839 

840@db_injector 

841async def read_automation_state_snapshot( 

842 db: PrefectDBInterface, session: AsyncSession 

843) -> AutomationStateSnapshot: 

844 query = sa.select( 

845 sa.func.max(db.Automation.updated), 

846 sa.func.count(db.Automation.id), 

847 ).select_from(db.Automation) 

848 

849 latest_updated, count = (await session.execute(query)).one() 

850 

851 return ( 

852 prefect.types._datetime.create_datetime_instance(latest_updated) 

853 if latest_updated 

854 else None, 

855 count or 0, 

856 ) 

857 

858 

859async def reconcile_automations(force: bool = False) -> bool: 

860 global automation_state_snapshot 

861 

862 async with _automations_lock(): 

863 async with automations_session() as session: 

864 current_snapshot = await read_automation_state_snapshot(session) 

865 if not force and current_snapshot == automation_state_snapshot: 

866 return False 

867 

868 previous_automations = automations_by_id.copy() 

869 previous_triggers = triggers.copy() 

870 previous_triggers_by_expect = { 

871 k: v.copy() for k, v in _triggers_by_expect.items() 

872 } 

873 previous_next_proactive_runs = next_proactive_runs.copy() 

874 

875 clear_loaded_automations() 

876 

877 try: 

878 await load_automations(session) 

879 except Exception: 

880 clear_loaded_automations() 

881 automations_by_id.update(previous_automations) 

882 triggers.update(previous_triggers) 

883 _triggers_by_expect.update(previous_triggers_by_expect) 

884 next_proactive_runs.update(previous_next_proactive_runs) 

885 raise 

886 

887 automation_state_snapshot = current_snapshot 

888 return True 

889