Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/models/artifacts.py: 72%
159 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1from typing import Any, Optional, Sequence, TypeVar, Union
2from uuid import UUID
4import sqlalchemy as sa
5from sqlalchemy import select
6from sqlalchemy.ext.asyncio import AsyncSession
7from sqlalchemy.sql import Select
9from prefect.server.database import PrefectDBInterface, db_injector, orm_models
10from prefect.server.events import clients
11from prefect.server.events.schemas import lifecycle
12from prefect.server.schemas import actions, filters, sorting
13from prefect.server.schemas.core import Artifact
14from prefect.types._datetime import DateTime, now
16T = TypeVar("T", bound=tuple[Any, ...])
19async def emit_artifact_collection_created_event(
20 artifact_collection: orm_models.ArtifactCollection,
21) -> None:
22 """Emit an event when an artifact collection is created."""
23 async with clients.PrefectServerEventsClient() as events_client:
24 await events_client.emit(
25 lifecycle.artifact_collection_created_event(artifact_collection, now("UTC"))
26 )
29async def emit_artifact_collection_updated_event(
30 artifact_collection: orm_models.ArtifactCollection,
31) -> None:
32 """Emit an event when an artifact collection's latest artifact changes."""
33 async with clients.PrefectServerEventsClient() as events_client:
34 await events_client.emit(
35 lifecycle.artifact_collection_updated_event(artifact_collection, now("UTC"))
36 )
39async def emit_artifact_collection_deleted_event(
40 artifact_collection: orm_models.ArtifactCollection,
41) -> None:
42 """Emit an event when an artifact collection is deleted."""
43 async with clients.PrefectServerEventsClient() as events_client:
44 await events_client.emit(
45 lifecycle.artifact_collection_deleted_event(artifact_collection, now("UTC"))
46 )
49@db_injector
50async def _insert_into_artifact_collection(
51 db: PrefectDBInterface,
52 session: AsyncSession,
53 artifact: Artifact,
54 now: Optional[DateTime] = None,
55) -> orm_models.ArtifactCollection:
56 """
57 Inserts a new artifact into the artifact_collection table or updates it.
58 """
59 insert_values = artifact.model_dump_for_orm(
60 exclude_unset=True, exclude={"id", "updated", "created"}
61 )
62 upsert_new_latest_id = (
63 db.queries.insert(db.ArtifactCollection)
64 .values(latest_id=artifact.id, updated=now, created=now, **insert_values)
65 .on_conflict_do_update(
66 index_elements=db.orm.artifact_collection_unique_upsert_columns,
67 set_=dict(
68 latest_id=artifact.id,
69 updated=now,
70 **insert_values,
71 ),
72 )
73 )
75 await session.execute(upsert_new_latest_id)
77 query = (
78 sa.select(db.ArtifactCollection)
79 .where(sa.and_(db.ArtifactCollection.key == artifact.key))
80 .execution_options(populate_existing=True)
81 )
83 result = await session.execute(query)
85 model = result.scalar()
87 if model is not None:
88 if model.latest_id != artifact.id:
89 raise ValueError(
90 f"Artifact {artifact.id} was not inserted into the artifact collection"
91 " table."
92 )
93 if model is None:
94 raise ValueError(
95 f"Artifact {artifact.id} was not inserted into the artifact collection"
96 " table."
97 )
99 if model.created == model.updated:
100 await emit_artifact_collection_created_event(model)
101 else:
102 await emit_artifact_collection_updated_event(model)
104 return model
107@db_injector
108async def _insert_into_artifact(
109 db: PrefectDBInterface,
110 session: AsyncSession,
111 artifact: Artifact,
112 now: Optional[DateTime] = None,
113) -> orm_models.Artifact:
114 """
115 Inserts a new artifact into the artifact table.
116 """
117 artifact_id = artifact.id
118 insert_stmt = db.queries.insert(db.Artifact).values(
119 created=now,
120 updated=now,
121 **artifact.model_dump_for_orm(exclude={"created", "updated"}),
122 )
123 await session.execute(insert_stmt)
125 query = (
126 sa.select(db.Artifact)
127 .where(db.Artifact.id == artifact_id)
128 .limit(1)
129 .execution_options(populate_existing=True)
130 )
132 result = await session.execute(query)
133 return result.scalar_one()
136async def create_artifact(
137 session: AsyncSession,
138 artifact: Artifact,
139) -> orm_models.Artifact:
140 right_now = now("UTC")
142 if artifact.key is not None:
143 await _insert_into_artifact_collection(
144 session=session, now=right_now, artifact=artifact
145 )
147 result = await _insert_into_artifact(
148 session=session,
149 now=right_now,
150 artifact=artifact,
151 )
153 return result
156@db_injector
157async def read_latest_artifact(
158 db: PrefectDBInterface,
159 session: AsyncSession,
160 key: str,
161) -> Union[orm_models.ArtifactCollection, None]:
162 """
163 Reads the latest artifact by key.
164 Args:
165 session: A database session
166 key: The artifact key
167 Returns:
168 Artifact: The latest artifact
169 """
170 latest_artifact_query = sa.select(db.ArtifactCollection).where(
171 db.ArtifactCollection.key == key
172 )
173 result = await session.execute(latest_artifact_query)
174 return result.scalar()
177@db_injector
178async def read_artifact(
179 db: PrefectDBInterface,
180 session: AsyncSession,
181 artifact_id: UUID,
182) -> Union[orm_models.Artifact, None]:
183 """
184 Reads an artifact by id.
185 """
187 query = sa.select(db.Artifact).where(db.Artifact.id == artifact_id)
189 result = await session.execute(query)
190 return result.scalar()
193async def _apply_artifact_filters(
194 db: PrefectDBInterface,
195 query: Select[T],
196 flow_run_filter: Optional[filters.FlowRunFilter] = None,
197 task_run_filter: Optional[filters.TaskRunFilter] = None,
198 artifact_filter: Optional[filters.ArtifactFilter] = None,
199 deployment_filter: Optional[filters.DeploymentFilter] = None,
200 flow_filter: Optional[filters.FlowFilter] = None,
201) -> Select[T]:
202 """Applies filters to an artifact query as a combination of EXISTS subqueries."""
203 if artifact_filter:
204 query = query.where(artifact_filter.as_sql_filter())
206 if flow_filter or flow_run_filter or deployment_filter:
207 flow_run_exists_clause = select(db.FlowRun).where(
208 db.Artifact.flow_run_id == db.FlowRun.id
209 )
210 if flow_run_filter:
211 flow_run_exists_clause = flow_run_exists_clause.where(
212 flow_run_filter.as_sql_filter()
213 )
215 if flow_filter:
216 flow_run_exists_clause = flow_run_exists_clause.join(
217 db.Flow, db.Flow.id == db.FlowRun.flow_id
218 ).where(flow_filter.as_sql_filter())
220 if deployment_filter:
221 flow_run_exists_clause = flow_run_exists_clause.join(
222 db.Deployment, db.Deployment.id == db.FlowRun.deployment_id
223 ).where(deployment_filter.as_sql_filter())
225 query = query.where(flow_run_exists_clause.exists())
227 if task_run_filter:
228 task_run_exists_clause = select(db.TaskRun).where(
229 db.Artifact.task_run_id == db.TaskRun.id
230 )
231 task_run_exists_clause = task_run_exists_clause.where(
232 task_run_filter.as_sql_filter()
233 )
235 query = query.where(task_run_exists_clause.exists())
237 return query
240async def _apply_artifact_collection_filters(
241 db: PrefectDBInterface,
242 query: Select[T],
243 flow_run_filter: Optional[filters.FlowRunFilter] = None,
244 task_run_filter: Optional[filters.TaskRunFilter] = None,
245 artifact_filter: Optional[filters.ArtifactCollectionFilter] = None,
246 deployment_filter: Optional[filters.DeploymentFilter] = None,
247 flow_filter: Optional[filters.FlowFilter] = None,
248) -> Select[T]:
249 """Applies filters to an artifact collection query as a combination of EXISTS subqueries."""
250 if artifact_filter:
251 query = query.where(artifact_filter.as_sql_filter())
253 if flow_filter or flow_run_filter or deployment_filter:
254 flow_run_exists_clause = select(db.FlowRun).where(
255 db.ArtifactCollection.flow_run_id == db.FlowRun.id
256 )
257 if flow_run_filter:
258 flow_run_exists_clause = flow_run_exists_clause.where(
259 flow_run_filter.as_sql_filter()
260 )
262 if flow_filter:
263 flow_run_exists_clause = flow_run_exists_clause.join(
264 db.Flow, db.Flow.id == db.FlowRun.flow_id
265 ).where(flow_filter.as_sql_filter())
267 if deployment_filter:
268 flow_run_exists_clause = flow_run_exists_clause.join(
269 db.Deployment, db.Deployment.id == db.FlowRun.deployment_id
270 ).where(deployment_filter.as_sql_filter())
272 query = query.where(flow_run_exists_clause.exists())
274 if task_run_filter:
275 task_run_exists_clause = select(db.TaskRun).where(
276 db.ArtifactCollection.task_run_id == db.TaskRun.id
277 )
278 task_run_exists_clause = task_run_exists_clause.where(
279 task_run_filter.as_sql_filter()
280 )
282 query = query.where(task_run_exists_clause.exists())
284 return query
287@db_injector
288async def read_artifacts(
289 db: PrefectDBInterface,
290 session: AsyncSession,
291 offset: Optional[int] = None,
292 limit: Optional[int] = None,
293 artifact_filter: Optional[filters.ArtifactFilter] = None,
294 flow_run_filter: Optional[filters.FlowRunFilter] = None,
295 task_run_filter: Optional[filters.TaskRunFilter] = None,
296 deployment_filter: Optional[filters.DeploymentFilter] = None,
297 flow_filter: Optional[filters.FlowFilter] = None,
298 sort: sorting.ArtifactSort = sorting.ArtifactSort.ID_DESC,
299) -> Sequence[orm_models.Artifact]:
300 """
301 Reads artifacts.
303 Args:
304 session: A database session
305 offset: Query offset
306 limit: Query limit
307 artifact_filter: Only select artifacts matching this filter
308 flow_run_filter: Only select artifacts whose flow runs matching this filter
309 task_run_filter: Only select artifacts whose task runs matching this filter
310 deployment_filter: Only select artifacts whose flow runs belong to deployments matching this filter
311 flow_filter: Only select artifacts whose flow runs belong to flows matching this filter
312 """
313 query = sa.select(db.Artifact).order_by(*sort.as_sql_sort())
315 query = await _apply_artifact_filters(
316 db,
317 query,
318 artifact_filter=artifact_filter,
319 flow_run_filter=flow_run_filter,
320 task_run_filter=task_run_filter,
321 deployment_filter=deployment_filter,
322 flow_filter=flow_filter,
323 )
325 if offset is not None: 325 ↛ 327line 325 didn't jump to line 327 because the condition on line 325 was always true
326 query = query.offset(offset)
327 if limit is not None: 327 ↛ 330line 327 didn't jump to line 330 because the condition on line 327 was always true
328 query = query.limit(limit)
330 result = await session.execute(query)
331 return result.scalars().unique().all()
334@db_injector
335async def read_latest_artifacts(
336 db: PrefectDBInterface,
337 session: AsyncSession,
338 offset: Optional[int] = None,
339 limit: Optional[int] = None,
340 artifact_filter: Optional[filters.ArtifactCollectionFilter] = None,
341 flow_run_filter: Optional[filters.FlowRunFilter] = None,
342 task_run_filter: Optional[filters.TaskRunFilter] = None,
343 deployment_filter: Optional[filters.DeploymentFilter] = None,
344 flow_filter: Optional[filters.FlowFilter] = None,
345 sort: sorting.ArtifactCollectionSort = sorting.ArtifactCollectionSort.ID_DESC,
346) -> Sequence[orm_models.ArtifactCollection]:
347 """
348 Reads artifacts.
350 Args:
351 session: A database session
352 offset: Query offset
353 limit: Query limit
354 artifact_filter: Only select artifacts matching this filter
355 flow_run_filter: Only select artifacts whose flow runs matching this filter
356 task_run_filter: Only select artifacts whose task runs matching this filter
357 deployment_filter: Only select artifacts whose flow runs belong to deployments matching this filter
358 flow_filter: Only select artifacts whose flow runs belong to flows matching this filter
359 """
360 query = sa.select(db.ArtifactCollection).order_by(*sort.as_sql_sort())
361 query = await _apply_artifact_collection_filters(
362 db,
363 query,
364 artifact_filter=artifact_filter,
365 flow_run_filter=flow_run_filter,
366 task_run_filter=task_run_filter,
367 deployment_filter=deployment_filter,
368 flow_filter=flow_filter,
369 )
371 if offset is not None: 371 ↛ 373line 371 didn't jump to line 373 because the condition on line 371 was always true
372 query = query.offset(offset)
373 if limit is not None: 373 ↛ 376line 373 didn't jump to line 376 because the condition on line 373 was always true
374 query = query.limit(limit)
376 result = await session.execute(query)
377 return result.scalars().unique().all()
380@db_injector
381async def count_artifacts(
382 db: PrefectDBInterface,
383 session: AsyncSession,
384 artifact_filter: Optional[filters.ArtifactFilter] = None,
385 flow_run_filter: Optional[filters.FlowRunFilter] = None,
386 task_run_filter: Optional[filters.TaskRunFilter] = None,
387 deployment_filter: Optional[filters.DeploymentFilter] = None,
388 flow_filter: Optional[filters.FlowFilter] = None,
389) -> int:
390 """
391 Counts artifacts.
392 Args:
393 session: A database session
394 artifact_filter: Only select artifacts matching this filter
395 flow_run_filter: Only select artifacts whose flow runs matching this filter
396 task_run_filter: Only select artifacts whose task runs matching this filter
397 """
398 query = sa.select(sa.func.count(db.Artifact.id))
400 query = await _apply_artifact_filters(
401 db,
402 query,
403 artifact_filter=artifact_filter,
404 flow_run_filter=flow_run_filter,
405 task_run_filter=task_run_filter,
406 deployment_filter=deployment_filter,
407 flow_filter=flow_filter,
408 )
410 result = await session.execute(query)
411 return result.scalar_one()
414@db_injector
415async def count_latest_artifacts(
416 db: PrefectDBInterface,
417 session: AsyncSession,
418 artifact_filter: Optional[filters.ArtifactCollectionFilter] = None,
419 flow_run_filter: Optional[filters.FlowRunFilter] = None,
420 task_run_filter: Optional[filters.TaskRunFilter] = None,
421 deployment_filter: Optional[filters.DeploymentFilter] = None,
422 flow_filter: Optional[filters.FlowFilter] = None,
423) -> int:
424 """
425 Counts artifacts.
426 Args:
427 session: A database session
428 artifact_filter: Only select artifacts matching this filter
429 flow_run_filter: Only select artifacts whose flow runs matching this filter
430 task_run_filter: Only select artifacts whose task runs matching this filter
431 """
432 query = sa.select(sa.func.count(db.ArtifactCollection.id))
434 query = await _apply_artifact_collection_filters(
435 db,
436 query,
437 artifact_filter=artifact_filter,
438 flow_run_filter=flow_run_filter,
439 task_run_filter=task_run_filter,
440 deployment_filter=deployment_filter,
441 flow_filter=flow_filter,
442 )
444 result = await session.execute(query)
445 return result.scalar_one()
448@db_injector
449async def update_artifact(
450 db: PrefectDBInterface,
451 session: AsyncSession,
452 artifact_id: UUID,
453 artifact: actions.ArtifactUpdate,
454) -> bool:
455 """
456 Updates an artifact by id.
458 Args:
459 session: A database session
460 artifact_id (UUID): The artifact id to update
461 artifact: An artifact model
463 Returns:
464 bool: True if the update was successful, False otherwise
465 """
466 update_artifact_data = artifact.model_dump_for_orm(exclude_unset=True)
468 update_artifact_stmt = (
469 sa.update(db.Artifact)
470 .where(db.Artifact.id == artifact_id)
471 .values(**update_artifact_data)
472 )
474 artifact_result = await session.execute(update_artifact_stmt)
476 update_artifact_collection_data = artifact.model_dump_for_orm(exclude_unset=True)
477 update_artifact_collection_stmt = (
478 sa.update(db.ArtifactCollection)
479 .where(db.ArtifactCollection.latest_id == artifact_id)
480 .values(**update_artifact_collection_data)
481 )
482 collection_result = await session.execute(update_artifact_collection_stmt)
484 return artifact_result.rowcount + collection_result.rowcount > 0
487@db_injector
488async def delete_artifact(
489 db: PrefectDBInterface, session: AsyncSession, artifact_id: UUID
490) -> bool:
491 """
492 Deletes an artifact by id.
494 The ArtifactCollection table is used to track the latest version of an artifact
495 by key. If we are deleting the latest version of an artifact from the Artifact
496 table, we need to first update the latest version referenced in ArtifactCollection
497 so that it points to the next latest version of the artifact.
499 Example:
500 If we have the following artifacts in Artifact:
501 - key: "foo", id: 1, created: 2020-01-01
502 - key: "foo", id: 2, created: 2020-01-02
503 - key: "foo", id: 3, created: 2020-01-03
505 the ArtifactCollection table has the following entry:
506 - key: "foo", latest_id: 3
508 If we delete the artifact with id 3, we need to update the latest version of the
509 artifact with key "foo" to be the artifact with id 2.
511 Args:
512 session: A database session
513 artifact_id (UUID): The artifact id to delete
515 Returns:
516 bool: True if the delete was successful, False otherwise
517 """
518 artifact = await session.get(db.Artifact, artifact_id)
519 if artifact is None:
520 return False
522 is_latest_version = (
523 await session.execute(
524 sa.select(db.ArtifactCollection)
525 .where(db.ArtifactCollection.key == artifact.key)
526 .where(db.ArtifactCollection.latest_id == artifact_id)
527 )
528 ).scalar_one_or_none() is not None
530 if is_latest_version:
531 next_latest_version = (
532 await session.execute(
533 sa.select(db.Artifact)
534 .where(db.Artifact.key == artifact.key)
535 .where(db.Artifact.id != artifact_id)
536 .order_by(db.Artifact.created.desc())
537 .limit(1)
538 )
539 ).scalar_one_or_none()
541 collection = (
542 await session.execute(
543 sa.select(db.ArtifactCollection).where(
544 db.ArtifactCollection.key == artifact.key
545 )
546 )
547 ).scalar_one_or_none()
549 if next_latest_version is not None:
550 set_next_latest_version = (
551 sa.update(db.ArtifactCollection)
552 .where(db.ArtifactCollection.key == artifact.key)
553 .values(
554 latest_id=next_latest_version.id,
555 data=next_latest_version.data,
556 description=next_latest_version.description,
557 type=next_latest_version.type,
558 created=next_latest_version.created,
559 updated=next_latest_version.updated,
560 flow_run_id=next_latest_version.flow_run_id,
561 task_run_id=next_latest_version.task_run_id,
562 metadata_=next_latest_version.metadata_,
563 )
564 )
565 await session.execute(set_next_latest_version)
567 if collection is not None:
568 await session.refresh(collection)
569 await emit_artifact_collection_updated_event(collection)
571 else:
572 if collection is not None:
573 await emit_artifact_collection_deleted_event(collection)
574 await session.execute(
575 sa.delete(db.ArtifactCollection)
576 .where(db.ArtifactCollection.key == artifact.key)
577 .where(db.ArtifactCollection.latest_id == artifact_id)
578 )
580 delete_stmt = sa.delete(db.Artifact).where(db.Artifact.id == artifact_id)
582 result = await session.execute(delete_stmt)
583 return result.rowcount > 0