Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/models/work_queues.py: 54%
253 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1"""
2Functions for interacting with work queue ORM objects.
3Intended for internal use by the Prefect REST API.
4"""
6import datetime
7from typing import (
8 Any,
9 Awaitable,
10 Callable,
11 Dict,
12 Iterable,
13 Optional,
14 Sequence,
15 Tuple,
16 Union,
17 cast,
18)
19from uuid import UUID
21import sqlalchemy as sa
22from docket import Depends, Retry
23from pydantic import TypeAdapter
24from sqlalchemy import delete, select
25from sqlalchemy.ext.asyncio import AsyncSession
27import prefect.server.models as models
28import prefect.server.schemas as schemas
29from prefect.server.database import (
30 PrefectDBInterface,
31 db_injector,
32 orm_models,
33 provide_database_interface,
34)
35from prefect.server.events import clients
36from prefect.server.events.clients import PrefectServerEventsClient
37from prefect.server.exceptions import ObjectNotFoundError
38from prefect.server.models.events import (
39 work_queue_created_event,
40 work_queue_deleted_event,
41 work_queue_status_event,
42 work_queue_updated_event,
43)
44from prefect.server.models.workers import (
45 DEFAULT_AGENT_WORK_POOL_NAME,
46 bulk_update_work_queue_priorities,
47)
48from prefect.server.schemas.states import StateType
49from prefect.server.schemas.statuses import WorkQueueStatus
50from prefect.server.utilities.database import UUID as PrefectUUID
51from prefect.types._datetime import DateTime, now
53WORK_QUEUE_LAST_POLLED_TIMEOUT = datetime.timedelta(seconds=60)
56@db_injector
57async def create_work_queue(
58 db: PrefectDBInterface,
59 session: AsyncSession,
60 work_queue: Union[schemas.core.WorkQueue, schemas.actions.WorkQueueCreate],
61) -> orm_models.WorkQueue:
62 """
63 Inserts a WorkQueue.
65 If a WorkQueue with the same name exists, an error will be thrown.
67 Args:
68 session (AsyncSession): a database session
69 work_queue (schemas.core.WorkQueue): a WorkQueue model
71 Returns:
72 orm_models.WorkQueue: the newly-created or updated WorkQueue
74 """
75 data = work_queue.model_dump()
77 if data.get("work_pool_id") is None: 77 ↛ 104line 77 didn't jump to line 104 because the condition on line 77 was always true
78 # If no work pool is provided, get or create the default agent work pool
79 default_agent_work_pool = await models.workers.read_work_pool_by_name(
80 session=session, work_pool_name=DEFAULT_AGENT_WORK_POOL_NAME
81 )
82 if default_agent_work_pool:
83 data["work_pool_id"] = default_agent_work_pool.id
84 else:
85 default_agent_work_pool = await models.workers.create_work_pool(
86 session=session,
87 work_pool=schemas.actions.WorkPoolCreate(
88 name=DEFAULT_AGENT_WORK_POOL_NAME, type="prefect-agent"
89 ),
90 )
91 if work_queue.name == "default":
92 # If the desired work queue name is default, it was created when the
93 # work pool was created. We can just return it.
94 default_work_queue = await models.workers.read_work_queue(
95 session=session,
96 work_queue_id=default_agent_work_pool.default_queue_id,
97 )
98 assert default_work_queue
99 return default_work_queue
100 data["work_pool_id"] = default_agent_work_pool.id
102 # Set the priority to be the max priority + 1
103 # This will make the new queue the lowest priority
104 if data["priority"] is None:
105 # Set the priority to be the first priority value that isn't already taken
106 priorities_query = sa.select(db.WorkQueue.priority).where(
107 db.WorkQueue.work_pool_id == data["work_pool_id"]
108 )
109 priorities = (await session.execute(priorities_query)).scalars().all()
111 priority = None
112 for i, p in enumerate(sorted(priorities)):
113 # if a rank was skipped (e.g. the set priority is different than the
114 # enumerated priority) then we can "take" that spot for this work
115 # queue
116 if i + 1 != p:
117 priority = i + 1
118 break
120 # otherwise take the maximum priority plus one
121 if priority is None:
122 priority = max(priorities, default=0) + 1
124 data["priority"] = priority
126 model = db.WorkQueue(**data)
128 session.add(model)
129 await session.flush()
130 await session.refresh(model)
132 if work_queue.priority:
133 await bulk_update_work_queue_priorities(
134 session=session,
135 work_pool_id=data["work_pool_id"],
136 new_priorities={model.id: work_queue.priority},
137 )
139 await emit_work_queue_created_event(session=session, work_queue=model)
141 return model
144@db_injector
145async def read_work_queue(
146 db: PrefectDBInterface,
147 session: AsyncSession,
148 work_queue_id: Union[UUID, PrefectUUID],
149) -> Optional[orm_models.WorkQueue]:
150 """
151 Reads a WorkQueue by id.
153 Args:
154 session (AsyncSession): A database session
155 work_queue_id (str): a WorkQueue id
157 Returns:
158 orm_models.WorkQueue: the WorkQueue
159 """
161 return await session.get(db.WorkQueue, work_queue_id)
164# States counted against work queue concurrency by the ORM query in
165# query_components.get_scheduled_flow_runs_from_work_queues.
166SLOT_OCCUPYING_STATES = {
167 StateType.PENDING,
168 StateType.RUNNING,
169 StateType.CANCELLING,
170}
172# States counted against legacy tag-based work queue concurrency limits
173# in _legacy_get_runs_in_work_queue.
174LEGACY_SLOT_OCCUPYING_STATES = {
175 StateType.PENDING,
176 StateType.RUNNING,
177}
180async def _count_legacy_queue_active_slots(
181 session: AsyncSession,
182 queue_filter: schemas.core.QueueFilter,
183) -> int:
184 """Count active slots for a legacy tag-based queue using its filter criteria."""
185 return await models.flow_runs.count_flow_runs(
186 session=session,
187 flow_run_filter=schemas.filters.FlowRunFilter(
188 tags=schemas.filters.FlowRunFilterTags(all_=queue_filter.tags),
189 deployment_id=schemas.filters.FlowRunFilterDeploymentId(
190 any_=queue_filter.deployment_ids, is_null_=False
191 ),
192 state=schemas.filters.FlowRunFilterState(
193 type=schemas.filters.FlowRunFilterStateType(
194 any_=list(LEGACY_SLOT_OCCUPYING_STATES)
195 )
196 ),
197 ),
198 )
201@db_injector
202async def count_work_queue_active_slots(
203 db: PrefectDBInterface,
204 session: AsyncSession,
205 work_queue_id: UUID,
206) -> int:
207 """
208 Count flow runs occupying concurrency slots for a given work queue.
210 For standard queues (including pool-backed and default-agent queues),
211 counts Pending/Running/Cancelling flow runs by work_queue_id FK.
213 For legacy tag-based queues, counts Pending/Running flow runs matching
214 the queue's tag/deployment filter (matching _legacy_get_runs_in_work_queue).
215 """
216 work_queue = await session.get(db.WorkQueue, work_queue_id)
217 if work_queue is None: 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true
218 return 0
220 if work_queue.filter is not None:
221 queue_filter = TypeAdapter(schemas.core.QueueFilter).validate_python(
222 work_queue.filter
223 )
224 return await _count_legacy_queue_active_slots(session, queue_filter)
226 query = (
227 select(sa.func.count())
228 .select_from(db.FlowRun)
229 .where(
230 db.FlowRun.work_queue_id == work_queue_id,
231 db.FlowRun.state_type.in_(SLOT_OCCUPYING_STATES),
232 )
233 )
234 result = await session.execute(query)
235 return result.scalar_one()
238@db_injector
239async def count_work_queue_active_slots_bulk(
240 db: PrefectDBInterface,
241 session: AsyncSession,
242 work_queue_ids: Sequence[UUID],
243) -> dict[UUID, int]:
244 """
245 Count active slots for multiple work queues. Standard queues are counted
246 in a single bulk GROUP BY query; legacy tag-based queues fall back to
247 per-queue counting since each has its own filter criteria.
248 """
249 if not work_queue_ids: 249 ↛ 250line 249 didn't jump to line 250 because the condition on line 249 was never true
250 return {}
252 # Load only id and filter to classify queue types without full ORM hydration
253 query = select(db.WorkQueue.id, db.WorkQueue.filter).where(
254 db.WorkQueue.id.in_(work_queue_ids)
255 )
256 result = await session.execute(query)
257 rows = result.all()
259 standard_ids: list[UUID] = []
260 legacy_filters: list[tuple[UUID, Any]] = []
261 for wq_id, wq_filter in rows:
262 if wq_filter is not None:
263 legacy_filters.append((wq_id, wq_filter))
264 else:
265 standard_ids.append(wq_id)
267 counts: dict[UUID, int] = {}
269 if standard_ids:
270 bulk_query = (
271 select(
272 db.FlowRun.work_queue_id,
273 sa.func.count(db.FlowRun.id),
274 )
275 .select_from(db.FlowRun)
276 .where(
277 db.FlowRun.work_queue_id.in_(standard_ids),
278 db.FlowRun.state_type.in_(SLOT_OCCUPYING_STATES),
279 )
280 .group_by(db.FlowRun.work_queue_id)
281 )
282 bulk_result = await session.execute(bulk_query)
283 counts.update(dict(bulk_result.all()))
285 for wq_id, wq_filter in legacy_filters:
286 queue_filter = TypeAdapter(schemas.core.QueueFilter).validate_python(wq_filter)
287 counts[wq_id] = await _count_legacy_queue_active_slots(session, queue_filter)
289 return counts
292@db_injector
293async def read_work_queue_by_name(
294 db: PrefectDBInterface, session: AsyncSession, name: str
295) -> Optional[orm_models.WorkQueue]:
296 """
297 Reads a WorkQueue by name.
299 Args:
300 session (AsyncSession): A database session
301 name (str): a WorkQueue name
303 Returns:
304 orm_models.WorkQueue: the WorkQueue
305 """
306 default_work_pool = await models.workers.read_work_pool_by_name(
307 session=session, work_pool_name=DEFAULT_AGENT_WORK_POOL_NAME
308 )
309 # Logic to make sure this functionality doesn't break during migration
310 if default_work_pool is not None:
311 query = select(db.WorkQueue).filter_by(
312 name=name, work_pool_id=default_work_pool.id
313 )
314 else:
315 query = select(db.WorkQueue).filter_by(name=name)
316 result = await session.execute(query)
317 return result.scalar()
320@db_injector
321async def read_work_queues(
322 db: PrefectDBInterface,
323 session: AsyncSession,
324 offset: Optional[int] = None,
325 limit: Optional[int] = None,
326 work_queue_filter: Optional[schemas.filters.WorkQueueFilter] = None,
327) -> Sequence[orm_models.WorkQueue]:
328 """
329 Read WorkQueues.
331 Args:
332 session: A database session
333 offset: Query offset
334 limit: Query limit
335 work_queue_filter: only select work queues matching these filters
336 Returns:
337 Sequence[orm_models.WorkQueue]: WorkQueues
338 """
340 query = select(db.WorkQueue).order_by(db.WorkQueue.name)
342 if offset is not None: 342 ↛ 344line 342 didn't jump to line 344 because the condition on line 342 was always true
343 query = query.offset(offset)
344 if limit is not None: 344 ↛ 346line 344 didn't jump to line 346 because the condition on line 344 was always true
345 query = query.limit(limit)
346 if work_queue_filter:
347 query = query.where(work_queue_filter.as_sql_filter())
349 result = await session.execute(query)
350 return result.scalars().unique().all()
353def is_last_polled_recent(last_polled: Optional[DateTime]) -> bool:
354 if last_polled is None:
355 return False
356 return (now("UTC") - last_polled) <= WORK_QUEUE_LAST_POLLED_TIMEOUT
359@db_injector
360async def update_work_queue(
361 db: PrefectDBInterface,
362 session: AsyncSession,
363 work_queue_id: UUID,
364 work_queue: schemas.actions.WorkQueueUpdate,
365 emit_status_change: Optional[
366 Callable[[orm_models.WorkQueue], Awaitable[None]]
367 ] = None,
368) -> bool:
369 """
370 Update a WorkQueue by id.
372 Args:
373 session (AsyncSession): A database session
374 work_queue: the work queue data
375 work_queue_id (str): a WorkQueue id
377 Returns:
378 bool: whether or not the WorkQueue was updated
379 """
380 # exclude_unset=True allows us to only update values provided by
381 # the user, ignoring any defaults on the model
382 update_data = work_queue.model_dump_for_orm(exclude_unset=True)
383 current_work_queue = await read_work_queue(
384 session=session, work_queue_id=work_queue_id
385 )
386 if current_work_queue is None: 386 ↛ 387line 386 didn't jump to line 387 because the condition on line 386 was never true
387 return False
389 session.expunge(current_work_queue)
391 if "is_paused" in update_data:
392 # Only update the status to paused if it's not already paused. This ensures a work queue that is already
393 # paused will not get a status update if it's paused again
394 if (
395 update_data.get("is_paused")
396 and current_work_queue.status != WorkQueueStatus.PAUSED
397 ):
398 update_data["status"] = WorkQueueStatus.PAUSED
400 # If unpausing, only update status if it's currently paused. This ensures a work queue that is already
401 # unpaused will not get a status update if it's unpaused again
402 if ( 402 ↛ 407line 402 didn't jump to line 407 because the condition on line 402 was never true
403 update_data.get("is_paused") is False
404 and current_work_queue.status == WorkQueueStatus.PAUSED
405 ):
406 # Default status if unpaused
407 update_data["status"] = WorkQueueStatus.NOT_READY
409 # Determine source of last_polled: update_data or database
410 last_polled: Optional[DateTime]
411 if "last_polled" in update_data:
412 last_polled = cast(DateTime, update_data["last_polled"])
413 else:
414 last_polled = current_work_queue.last_polled
416 # Check if last polled is recent and set status to READY if so
417 if is_last_polled_recent(last_polled):
418 update_data["status"] = schemas.statuses.WorkQueueStatus.READY
420 update_stmt = (
421 sa.update(db.WorkQueue)
422 .where(db.WorkQueue.id == work_queue_id)
423 .values(**update_data)
424 )
425 result = await session.execute(update_stmt)
426 updated = result.rowcount > 0
428 if updated:
429 wq = await read_work_queue(session=session, work_queue_id=work_queue_id)
430 assert wq is not None
431 assert current_work_queue is not wq
432 WORK_QUEUE_EVENT_FIELDS = {
433 "name",
434 "description",
435 "concurrency_limit",
436 "priority",
437 # Exclude "is_paused" - handled with status
438 # Exclude "last_polled" - usually auto-updated
439 # Exclude "filter" - deprecated
440 }
441 # Detect which fields actually changed
442 changed_fields = {}
443 for field in update_data.keys():
444 if field not in WORK_QUEUE_EVENT_FIELDS or field == "status":
445 continue
447 old_value = getattr(current_work_queue, field, None)
448 new_value = getattr(wq, field, None)
450 if old_value != new_value:
451 changed_fields[field] = {
452 "from": old_value,
453 "to": new_value,
454 }
456 # Emit event for non-status field changes
457 if changed_fields:
458 await emit_work_queue_updated_event(
459 session=session,
460 work_queue=wq,
461 changed_fields=changed_fields,
462 )
464 if "status" in update_data and emit_status_change:
465 await emit_status_change(wq)
467 return updated
470@db_injector
471async def delete_work_queue(
472 db: PrefectDBInterface, session: AsyncSession, work_queue_id: UUID
473) -> bool:
474 """
475 Delete a WorkQueue by id.
477 Args:
478 session (AsyncSession): A database session
479 work_queue_id (str): a WorkQueue id
481 Returns:
482 bool: whether or not the WorkQueue was deleted
483 """
484 work_queue = await session.get(db.WorkQueue, work_queue_id)
485 if work_queue is None: 485 ↛ 486line 485 didn't jump to line 486 because the condition on line 485 was never true
486 return False
488 await emit_work_queue_deleted_event(session=session, work_queue=work_queue)
490 await session.execute(delete(db.WorkQueue).where(db.WorkQueue.id == work_queue_id))
491 return True
494@db_injector
495async def get_runs_in_work_queue(
496 db: PrefectDBInterface,
497 session: AsyncSession,
498 work_queue_id: UUID,
499 limit: Optional[int] = None,
500 scheduled_before: Optional[datetime.datetime] = None,
501) -> Tuple[orm_models.WorkQueue, Sequence[orm_models.FlowRun]]:
502 """
503 Get runs from a work queue.
505 Args:
506 session: A database session. work_queue_id: The work queue id.
507 scheduled_before: Only return runs scheduled to start before this time.
508 limit: An optional limit for the number of runs to return from the
509 queue. This limit applies to the request only. It does not affect
510 the work queue's concurrency limit. If `limit` exceeds the work
511 queue's concurrency limit, it will be ignored.
513 """
514 work_queue = await read_work_queue(session=session, work_queue_id=work_queue_id)
515 if not work_queue:
516 raise ObjectNotFoundError(f"Work queue with id {work_queue_id} not found.")
518 if work_queue.filter is None:
519 query = db.queries.get_scheduled_flow_runs_from_work_queues(
520 limit_per_queue=limit,
521 work_queue_ids=[work_queue_id],
522 scheduled_before=scheduled_before,
523 )
524 result = await session.execute(query)
525 return work_queue, result.scalars().unique().all()
527 # if the work queue has a filter, it's a deprecated tag-based work queue
528 # and uses an old approach
529 else:
530 return work_queue, await _legacy_get_runs_in_work_queue(
531 session=session,
532 work_queue_id=work_queue_id,
533 scheduled_before=scheduled_before,
534 limit=limit,
535 )
538async def _legacy_get_runs_in_work_queue(
539 session: AsyncSession,
540 work_queue_id: UUID,
541 scheduled_before: Optional[datetime.datetime] = None,
542 limit: Optional[int] = None,
543) -> Sequence[orm_models.FlowRun]:
544 """
545 DEPRECATED method for getting runs from a tag-based work queue
547 Args:
548 session: A database session.
549 work_queue_id: The work queue id.
550 scheduled_before: Only return runs scheduled to start before this time.
551 limit: An optional limit for the number of runs to return from the queue.
552 This limit applies to the request only. It does not affect the
553 work queue's concurrency limit. If `limit` exceeds the work queue's
554 concurrency limit, it will be ignored.
556 """
558 work_queue = await read_work_queue(session=session, work_queue_id=work_queue_id)
559 if not work_queue: 559 ↛ 560line 559 didn't jump to line 560 because the condition on line 559 was never true
560 raise ObjectNotFoundError(f"Work queue with id {work_queue_id} not found.")
562 if work_queue.is_paused: 562 ↛ 568line 562 didn't jump to line 568 because the condition on line 562 was always true
563 return []
565 # ensure the filter object is fully hydrated
566 # SQLAlchemy caching logic can result in a dict type instead
567 # of the full pydantic model
568 work_queue_filter = TypeAdapter(schemas.core.QueueFilter).validate_python(
569 work_queue.filter
570 )
571 flow_run_filter = dict(
572 tags=dict(all_=work_queue_filter.tags),
573 deployment_id=dict(any_=work_queue_filter.deployment_ids, is_null_=False),
574 )
576 # if the work queue has a concurrency limit, check how many runs are currently
577 # executing and compare that count to the concurrency limit
578 if work_queue.concurrency_limit is not None:
579 # Note this does not guarantee race conditions won't be hit
580 running_frs = await models.flow_runs.count_flow_runs(
581 session=session,
582 flow_run_filter=schemas.filters.FlowRunFilter(
583 **flow_run_filter,
584 state=dict(type=dict(any_=[StateType.PENDING, StateType.RUNNING])),
585 ),
586 )
588 # compute the available concurrency slots
589 open_concurrency_slots = max(0, work_queue.concurrency_limit - running_frs)
591 # if a limit override was given, ensure we return no more
592 # than that limit
593 if limit is not None:
594 limit = min(open_concurrency_slots, limit)
595 else:
596 limit = open_concurrency_slots
598 return await models.flow_runs.read_flow_runs(
599 session=session,
600 flow_run_filter=schemas.filters.FlowRunFilter(
601 **flow_run_filter,
602 state=dict(type=dict(any_=[StateType.SCHEDULED])),
603 next_scheduled_start_time=dict(before_=scheduled_before),
604 ),
605 limit=limit,
606 sort=schemas.sorting.FlowRunSort.NEXT_SCHEDULED_START_TIME_ASC,
607 )
610async def ensure_work_queue_exists(
611 session: AsyncSession, name: str
612) -> orm_models.WorkQueue:
613 """
614 Checks if a work queue exists and creates it if it does not.
616 Useful when working with deployments, agents, and flow runs that automatically create work queues.
618 Will also create a work pool queue in the default agent pool to facilitate migration to work pools.
619 """
620 # read work queue
621 work_queue = await models.work_queues.read_work_queue_by_name(
622 session=session, name=name
623 )
624 if not work_queue:
625 default_pool = await models.workers.read_work_pool_by_name(
626 session=session, work_pool_name=DEFAULT_AGENT_WORK_POOL_NAME
627 )
629 if default_pool is None:
630 work_queue = await models.work_queues.create_work_queue(
631 session=session,
632 work_queue=schemas.actions.WorkQueueCreate(name=name, priority=1),
633 )
634 else:
635 if name != "default":
636 work_queue = await models.workers.create_work_queue(
637 session=session,
638 work_pool_id=default_pool.id,
639 work_queue=schemas.actions.WorkQueueCreate(name=name, priority=1),
640 )
641 else:
642 work_queue = await models.work_queues.read_work_queue(
643 session=session, work_queue_id=default_pool.default_queue_id
644 )
645 assert work_queue, "Default work queue not found"
647 return work_queue
650async def read_work_queue_status(
651 session: AsyncSession, work_queue_id: UUID
652) -> schemas.core.WorkQueueStatusDetail:
653 """
654 Get work queue status by id.
656 Args:
657 session (AsyncSession): A database session
658 work_queue_id (str): a WorkQueue id
660 Returns:
661 Information about the status of the work queue.
662 """
664 work_queue = await read_work_queue(session=session, work_queue_id=work_queue_id)
665 if not work_queue: 665 ↛ 668line 665 didn't jump to line 668 because the condition on line 665 was always true
666 raise ObjectNotFoundError(f"Work queue with id {work_queue_id} not found")
668 work_queue_late_runs_count = await models.flow_runs.count_flow_runs(
669 session=session,
670 flow_run_filter=schemas.filters.FlowRunFilter(
671 state=schemas.filters.FlowRunFilterState(name={"any_": ["Late"]}),
672 ),
673 work_queue_filter=schemas.filters.WorkQueueFilter(
674 id=schemas.filters.WorkQueueFilterId(any_=[work_queue_id])
675 ),
676 )
678 # All work queues use the default policy for now
679 health_check_policy = schemas.core.WorkQueueHealthPolicy(
680 maximum_late_runs=0, maximum_seconds_since_last_polled=60
681 )
683 healthy = health_check_policy.evaluate_health_status(
684 late_runs_count=work_queue_late_runs_count,
685 last_polled=work_queue.last_polled, # type: ignore
686 )
688 return schemas.core.WorkQueueStatusDetail(
689 healthy=healthy,
690 late_runs_count=work_queue_late_runs_count,
691 last_polled=work_queue.last_polled,
692 health_check_policy=health_check_policy,
693 )
696@db_injector
697async def record_work_queue_polls(
698 db: PrefectDBInterface,
699 session: AsyncSession,
700 polled_work_queue_ids: Sequence[UUID],
701 ready_work_queue_ids: Sequence[UUID],
702) -> None:
703 """Record that the given work queues were polled, and also update the given
704 ready_work_queue_ids to READY."""
705 polled = now("UTC")
707 if polled_work_queue_ids:
708 await session.execute(
709 sa.update(db.WorkQueue)
710 .where(db.WorkQueue.id.in_(polled_work_queue_ids))
711 .values(last_polled=polled)
712 )
714 if ready_work_queue_ids: 714 ↛ exitline 714 didn't return from function 'record_work_queue_polls' because the condition on line 714 was always true
715 await session.execute(
716 sa.update(db.WorkQueue)
717 .where(db.WorkQueue.id.in_(ready_work_queue_ids))
718 .values(last_polled=polled, status=WorkQueueStatus.READY)
719 )
722async def mark_work_queues_ready(
723 *,
724 db: PrefectDBInterface = Depends(provide_database_interface),
725 polled_work_queue_ids: Sequence[UUID],
726 ready_work_queue_ids: Sequence[UUID],
727 retry: Retry = Retry(attempts=5, delay=datetime.timedelta(seconds=0.5)),
728) -> None:
729 async with db.session_context(begin_transaction=True) as session:
730 await record_work_queue_polls(
731 session=session,
732 polled_work_queue_ids=polled_work_queue_ids,
733 ready_work_queue_ids=ready_work_queue_ids,
734 )
736 # Emit events for any work queues that have transitioned to ready during this poll
737 # Uses a separate transaction to avoid keeping locks open longer from the updates
738 # in the previous transaction
739 if not ready_work_queue_ids:
740 return
742 async with db.session_context(begin_transaction=True) as session:
743 newly_ready_work_queues = await session.execute(
744 sa.select(db.WorkQueue).where(db.WorkQueue.id.in_(ready_work_queue_ids))
745 )
747 events = [
748 await work_queue_status_event(
749 session=session,
750 work_queue=work_queue,
751 occurred=now("UTC"),
752 )
753 for work_queue in newly_ready_work_queues.scalars().all()
754 ]
756 async with PrefectServerEventsClient() as events_client:
757 for event in events:
758 await events_client.emit(event)
761@db_injector
762async def mark_work_queues_not_ready(
763 db: PrefectDBInterface,
764 work_queue_ids: Iterable[UUID],
765) -> None:
766 if not work_queue_ids:
767 return
769 async with db.session_context(begin_transaction=True) as session:
770 await session.execute(
771 sa.update(db.WorkQueue)
772 .where(db.WorkQueue.id.in_(work_queue_ids))
773 .values(status=WorkQueueStatus.NOT_READY)
774 )
776 # Emit events for any work queues that have transitioned to ready during this poll
777 # Uses a separate transaction to avoid keeping locks open longer from the updates
778 # in the previous transaction
780 async with db.session_context(begin_transaction=True) as session:
781 newly_unready_work_queues = await session.execute(
782 sa.select(db.WorkQueue).where(db.WorkQueue.id.in_(work_queue_ids))
783 )
785 events = [
786 await work_queue_status_event(
787 session=session,
788 work_queue=work_queue,
789 occurred=now("UTC"),
790 )
791 for work_queue in newly_unready_work_queues.scalars().all()
792 ]
794 async with PrefectServerEventsClient() as events_client:
795 for event in events:
796 await events_client.emit(event)
799@db_injector
800async def emit_work_queue_status_event(
801 db: PrefectDBInterface,
802 work_queue: orm_models.WorkQueue,
803) -> None:
804 """Emit an event when work queue fields are updated."""
805 async with db.session_context() as session:
806 event = await work_queue_status_event(
807 session=session,
808 work_queue=work_queue,
809 occurred=now("UTC"),
810 )
811 async with PrefectServerEventsClient() as events_client:
812 await events_client.emit(event)
815async def emit_work_queue_updated_event(
816 session: AsyncSession,
817 work_queue: orm_models.WorkQueue,
818 changed_fields: Dict[str, Dict[str, Any]],
819) -> None:
820 if not changed_fields: 820 ↛ 821line 820 didn't jump to line 821 because the condition on line 820 was never true
821 return
823 async with PrefectServerEventsClient() as events_client:
824 await events_client.emit(
825 await work_queue_updated_event(
826 session=session,
827 work_queue=work_queue,
828 changed_fields=changed_fields,
829 occurred=now("UTC"),
830 )
831 )
834async def emit_work_queue_created_event(
835 session: AsyncSession,
836 work_queue: orm_models.WorkQueue,
837) -> None:
838 """Emit an event when a work queue is created."""
839 async with clients.PrefectServerEventsClient() as events_client:
840 await events_client.emit(
841 await work_queue_created_event(
842 session=session,
843 work_queue=work_queue,
844 occurred=now("UTC"),
845 )
846 )
849async def emit_work_queue_deleted_event(
850 session: AsyncSession,
851 work_queue: orm_models.WorkQueue,
852) -> None:
853 """Emit an event when a work queue is deleted."""
854 async with clients.PrefectServerEventsClient() as events_client:
855 await events_client.emit(
856 await work_queue_deleted_event(
857 session=session,
858 work_queue=work_queue,
859 occurred=now("UTC"),
860 )
861 )