Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/events/models/automations.py: 67%
181 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1import logging
2from contextlib import asynccontextmanager
3from typing import AsyncGenerator, Optional, Sequence, Union
4from uuid import UUID
6import orjson
7import sqlalchemy as sa
8from sqlalchemy.ext.asyncio import AsyncSession
9from typing_extensions import Literal, TypeAlias
11from prefect.logging import get_logger
12from prefect.server.database import PrefectDBInterface, db_injector
13from prefect.server.events import clients, filters
14from prefect.server.events.schemas import lifecycle
15from prefect.server.events.schemas.automations import (
16 Automation,
17 AutomationPartialUpdate,
18 AutomationSort,
19 AutomationUpdate,
20)
21from prefect.server.utilities.database import get_dialect
22from prefect.types._datetime import now
23from prefect.utilities.asyncutils import run_coro_as_sync
25logger: logging.Logger = get_logger(__name__)
28async def emit_automation_created_event(automation: Automation) -> None:
29 """Emit an event when an automation is created."""
30 async with clients.PrefectServerEventsClient() as events_client:
31 await events_client.emit(
32 lifecycle.automation_created_event(automation, now("UTC"))
33 )
36async def emit_automation_updated_event(automation: Automation) -> None:
37 """Emit an event when an automation is updated."""
38 async with clients.PrefectServerEventsClient() as events_client:
39 await events_client.emit(
40 lifecycle.automation_updated_event(automation, now("UTC"))
41 )
44async def emit_automation_deleted_event(automation: Automation) -> None:
45 """Emit an event when an automation is deleted."""
46 async with clients.PrefectServerEventsClient() as events_client:
47 await events_client.emit(
48 lifecycle.automation_deleted_event(automation, now("UTC"))
49 )
52AutomationChangeEvent: TypeAlias = Literal[
53 "automation__created", "automation__updated", "automation__deleted"
54]
55AUTOMATION_CHANGES_CHANNEL = "prefect_automation_changes"
58@asynccontextmanager
59@db_injector
60async def automations_session(
61 db: PrefectDBInterface, begin_transaction: bool = False
62) -> AsyncGenerator[AsyncSession, None]:
63 async with db.session_context(begin_transaction=begin_transaction) as session:
64 yield session
67@db_injector
68async def read_automations_for_workspace(
69 db: PrefectDBInterface,
70 session: AsyncSession,
71 sort: AutomationSort = AutomationSort.NAME_ASC,
72 limit: Optional[int] = None,
73 offset: Optional[int] = None,
74 automation_filter: Optional[filters.AutomationFilter] = None,
75) -> Sequence[Automation]:
76 query = sa.select(db.Automation)
78 query = query.order_by(db.Automation.sort_expression(sort))
80 if automation_filter:
81 query = query.where(automation_filter.as_sql_filter())
82 if limit is not None: 82 ↛ 84line 82 didn't jump to line 84 because the condition on line 82 was always true
83 query = query.limit(limit)
84 if offset is not None: 84 ↛ 87line 84 didn't jump to line 87 because the condition on line 84 was always true
85 query = query.offset(offset)
87 result = await session.execute(query)
89 return [
90 Automation.model_validate(a, from_attributes=True)
91 for a in result.scalars().all()
92 ]
95@db_injector
96async def count_automations_for_workspace(
97 db: PrefectDBInterface,
98 session: AsyncSession,
99) -> int:
100 query = sa.select(sa.func.count(None)).select_from(db.Automation)
102 result = await session.execute(query)
104 return result.scalar() or 0
107@db_injector
108async def read_automation(
109 db: PrefectDBInterface,
110 session: AsyncSession,
111 automation_id: UUID,
112) -> Optional[Automation]:
113 automation = await session.scalar(
114 sa.select(db.Automation).where(db.Automation.id == automation_id)
115 )
116 if not automation: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true
117 return None
118 return Automation.model_validate(automation, from_attributes=True)
121@db_injector
122async def read_automation_by_id(
123 db: PrefectDBInterface, session: AsyncSession, automation_id: UUID
124) -> Optional[Automation]:
125 automation = await session.scalar(
126 sa.select(db.Automation).where(
127 db.Automation.id == automation_id,
128 )
129 )
130 if not automation:
131 return None
132 return Automation.model_validate(automation, from_attributes=True)
135async def _notify(session: AsyncSession, automation: Automation, event: str):
136 from prefect.server.events.triggers import automation_changed
138 event_key: AutomationChangeEvent
139 if event == "created":
140 event_key = "automation__created"
141 elif event == "updated":
142 event_key = "automation__updated"
143 elif event == "deleted": 143 ↛ 146line 143 didn't jump to line 146 because the condition on line 143 was always true
144 event_key = "automation__deleted"
145 else:
146 logger.error(
147 f"Unknown event type '{event}' in _notify for automation {automation.id}"
148 )
149 return
151 # Handle cache updates based on database type
152 sync_session = session.sync_session
153 dialect_name = get_dialect(sync_session).name
155 if dialect_name == "postgresql": 155 ↛ 182line 155 didn't jump to line 182 because the condition on line 155 was always true
156 # For PostgreSQL, only send NOTIFY - the listener will update the cache
157 try:
158 payload_json = (
159 orjson.dumps(
160 {
161 "automation_id": str(automation.id),
162 "event_type": event,
163 }
164 )
165 .decode()
166 .replace("'", "''")
167 )
168 await session.execute(
169 sa.text(f"NOTIFY {AUTOMATION_CHANGES_CHANNEL}, '{payload_json}'")
170 )
172 logger.debug(
173 f"Sent Postgres NOTIFY on channel '{AUTOMATION_CHANGES_CHANNEL}' for automation {automation.id}, event: {event}"
174 )
175 except Exception as e:
176 logger.error(
177 f"Failed to send Postgres NOTIFY for automation {automation.id} on channel {AUTOMATION_CHANGES_CHANNEL}: {e}",
178 exc_info=True,
179 )
180 else:
181 # For SQLite, we need to update the cache after commit
182 @sa.event.listens_for(sync_session, "after_commit", once=True)
183 def update_cache_after_commit(session):
184 try:
185 run_coro_as_sync(automation_changed(automation.id, event_key))
186 except Exception as e:
187 logger.error(
188 f"Failed to update in-memory cache for automation {automation.id}, event: {event}: {e}",
189 exc_info=True,
190 )
193@db_injector
194async def create_automation(
195 db: PrefectDBInterface, session: AsyncSession, automation: Automation
196) -> Automation:
197 new_automation = db.Automation(**automation.model_dump())
198 session.add(new_automation)
199 await session.flush()
200 automation = Automation.model_validate(new_automation, from_attributes=True)
202 await _sync_automation_related_resources(session, new_automation.id, automation)
204 await _notify(session, automation, "created")
205 await emit_automation_created_event(automation)
206 return automation
209@db_injector
210async def update_automation(
211 db: PrefectDBInterface,
212 session: AsyncSession,
213 automation_update: Union[AutomationUpdate, AutomationPartialUpdate],
214 automation_id: UUID,
215) -> bool:
216 if not isinstance(automation_update, (AutomationUpdate, AutomationPartialUpdate)): 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 raise TypeError(
218 "automation_update must be an AutomationUpdate or AutomationPartialUpdate, "
219 f"not {type(automation_update)}"
220 )
222 automation = await read_automation(session, automation_id)
223 if not automation:
224 return False
226 if isinstance(automation_update, AutomationPartialUpdate):
227 # Partial updates won't go through the full Automation/AutomationCore
228 # validation, which could change due to one of these updates. Here we attempt
229 # to apply and parse the final effect of the partial update to the existing
230 # automation to see if anything fails validation.
231 Automation.model_validate(
232 {
233 **automation.model_dump(mode="json"),
234 **automation_update.model_dump(mode="json"),
235 }
236 )
238 result = await session.execute(
239 sa.update(db.Automation)
240 .where(db.Automation.id == automation_id)
241 .values(**automation_update.model_dump_for_orm(exclude_unset=True))
242 )
244 if isinstance(automation_update, AutomationUpdate):
245 await _sync_automation_related_resources(
246 session, automation_id, automation_update
247 )
249 await _notify(session, automation, "updated")
251 updated_automation = await read_automation(session, automation_id)
252 if updated_automation is not None:
253 await emit_automation_updated_event(updated_automation)
255 return result.rowcount > 0 # type: ignore
258@db_injector
259async def delete_automation(
260 db: PrefectDBInterface,
261 session: AsyncSession,
262 automation_id: UUID,
263) -> bool:
264 automation = await read_automation(session, automation_id)
265 if not automation:
266 return False
268 # Delete child tables in a consistent order to prevent deadlocks
269 # when multiple automations are deleted concurrently
270 await session.execute(
271 sa.delete(db.AutomationBucket).where(
272 db.AutomationBucket.automation_id == automation_id,
273 )
274 )
275 await session.execute(
276 sa.delete(db.AutomationRelatedResource).where(
277 db.AutomationRelatedResource.automation_id == automation_id,
278 )
279 )
280 await session.execute(
281 sa.delete(db.CompositeTriggerChildFiring).where(
282 db.CompositeTriggerChildFiring.automation_id == automation_id,
283 )
284 )
286 # Now delete the parent automation
287 await session.execute(
288 sa.delete(db.Automation).where(
289 db.Automation.id == automation_id,
290 )
291 )
292 await _sync_automation_related_resources(session, automation_id, None)
294 await _notify(session, automation, "deleted")
295 await emit_automation_deleted_event(automation)
296 return True
299@db_injector
300async def delete_automations_for_workspace(
301 db: PrefectDBInterface,
302 session: AsyncSession,
303) -> bool:
304 automations = await read_automations_for_workspace(
305 session,
306 )
308 # Delete child tables in a consistent order to prevent deadlocks
309 # when multiple workspace deletions occur concurrently
310 await session.execute(sa.delete(db.AutomationBucket))
311 await session.execute(sa.delete(db.AutomationRelatedResource))
312 await session.execute(sa.delete(db.CompositeTriggerChildFiring))
314 # Now delete all automations
315 result = await session.execute(sa.delete(db.Automation))
316 for automation in automations:
317 await _notify(session, automation, "deleted")
318 return result.rowcount > 0
321@db_injector
322async def disable_automations_for_workspace(
323 db: PrefectDBInterface,
324 session: AsyncSession,
325) -> bool:
326 automations = await read_automations_for_workspace(session)
327 result = await session.execute(sa.update(db.Automation).values(enabled=False))
328 for automation in automations:
329 await _notify(session, automation, "updated")
330 return result.rowcount > 0
333@db_injector
334async def disable_automation(
335 db: PrefectDBInterface, session: AsyncSession, automation_id: UUID
336) -> bool:
337 automation = await read_automation_by_id(
338 session=session,
339 automation_id=automation_id,
340 )
341 if not automation:
342 raise ValueError(f"Automation with ID {automation_id} not found")
344 result = await session.execute(
345 sa.update(db.Automation)
346 .where(db.Automation.id == automation_id)
347 .values(enabled=False)
348 )
349 await _notify(session, automation, "updated")
350 return result.rowcount > 0
353@db_injector
354async def _sync_automation_related_resources(
355 db: PrefectDBInterface,
356 session: AsyncSession,
357 automation_id: UUID,
358 automation: Optional[Union[Automation, AutomationUpdate]],
359):
360 """Actively maintains the set of related resources for an automation"""
361 from prefect.server.events import actions
363 await session.execute(
364 sa.delete(db.AutomationRelatedResource).where(
365 db.AutomationRelatedResource.automation_id == automation_id,
366 db.AutomationRelatedResource.resource_id.like("prefect.deployment.%"),
367 db.AutomationRelatedResource.automation_owned_by_resource.is_(False),
368 ),
369 execution_options={"synchronize_session": False},
370 )
372 if not automation:
373 return
375 deployment_ids = set(
376 action.deployment_id
377 for action in automation.actions
378 if isinstance(action, actions.RunDeployment) and action.source == "selected"
379 )
380 for deployment_id in deployment_ids:
381 await relate_automation_to_resource(
382 session, automation_id, f"prefect.deployment.{deployment_id}", False
383 )
386@db_injector
387async def relate_automation_to_resource(
388 db: PrefectDBInterface,
389 session: AsyncSession,
390 automation_id: UUID,
391 resource_id: str,
392 owned_by_resource: bool,
393) -> None:
394 await session.execute(
395 db.queries.insert(db.AutomationRelatedResource)
396 .values(
397 automation_id=automation_id,
398 resource_id=resource_id,
399 automation_owned_by_resource=owned_by_resource,
400 )
401 .on_conflict_do_update(
402 index_elements=[
403 db.AutomationRelatedResource.automation_id,
404 db.AutomationRelatedResource.resource_id,
405 ],
406 set_=dict(
407 automation_owned_by_resource=sa.or_(
408 db.AutomationRelatedResource.automation_owned_by_resource,
409 sa.true() if owned_by_resource else sa.false(),
410 ),
411 updated=now("UTC"),
412 ),
413 )
414 )
417@db_injector
418async def read_automations_related_to_resource(
419 db: PrefectDBInterface,
420 session: AsyncSession,
421 resource_id: str,
422 owned_by_resource: Optional[bool] = None,
423 automation_filter: Optional[filters.AutomationFilter] = None,
424) -> Sequence[Automation]:
425 query = (
426 sa.select(db.Automation)
427 .join(db.Automation.related_resources)
428 .where(
429 db.AutomationRelatedResource.resource_id == resource_id,
430 )
431 )
432 if owned_by_resource is not None:
433 query = query.where(
434 db.AutomationRelatedResource.automation_owned_by_resource
435 == owned_by_resource
436 )
438 if automation_filter:
439 query = query.where(automation_filter.as_sql_filter())
441 result = await session.execute(query)
442 return [
443 Automation.model_validate(a, from_attributes=True)
444 for a in result.scalars().all()
445 ]
448@db_injector
449async def delete_automations_owned_by_resource(
450 db: PrefectDBInterface,
451 session: AsyncSession,
452 resource_id: str,
453 automation_filter: Optional[filters.AutomationFilter] = None,
454) -> Sequence[UUID]:
455 automations = await read_automations_related_to_resource(
456 session=session,
457 resource_id=resource_id,
458 owned_by_resource=True,
459 automation_filter=automation_filter,
460 )
462 automation_ids = [automation.id for automation in automations]
464 await session.execute(
465 sa.delete(db.Automation).where(db.Automation.id.in_(automation_ids))
466 )
468 for automation in automations:
469 await _notify(session, automation, "deleted")
471 return automation_ids