prefect.server.models.work_queues

Functions for interacting with work queue ORM objects. Intended for internal use by the Prefect REST API.

create_work_queue(session, work_queue) async

Inserts a WorkQueue.

If a WorkQueue with the same name exists, an error will be thrown.

Parameters:

Name	Type	Description	Default
session	AsyncSession	a database session	required
work_queue	WorkQueue	a WorkQueue model	required

Returns:

Type	Description
WorkQueue	orm_models.WorkQueue: the newly-created or updated WorkQueue

Source code in src/prefect/server/models/work_queues.py

async def create_work_queue(
    session: AsyncSession,
    work_queue: Union[schemas.core.WorkQueue, schemas.actions.WorkQueueCreate],
) -> orm_models.WorkQueue:
    """
    Inserts a WorkQueue.

    If a WorkQueue with the same name exists, an error will be thrown.

    Args:
        session (AsyncSession): a database session
        work_queue (schemas.core.WorkQueue): a WorkQueue model

    Returns:
        orm_models.WorkQueue: the newly-created or updated WorkQueue

    """
    data = work_queue.model_dump()

    if data.get("work_pool_id") is None:
        # If no work pool is provided, get or create the default agent work pool
        default_agent_work_pool = await models.workers.read_work_pool_by_name(
            session=session, work_pool_name=DEFAULT_AGENT_WORK_POOL_NAME
        )
        if default_agent_work_pool:
            data["work_pool_id"] = default_agent_work_pool.id
        else:
            default_agent_work_pool = await models.workers.create_work_pool(
                session=session,
                work_pool=schemas.actions.WorkPoolCreate(
                    name=DEFAULT_AGENT_WORK_POOL_NAME, type="prefect-agent"
                ),
            )
            if work_queue.name == "default":
                # If the desired work queue name is default, it was created when the
                # work pool was created. We can just return it.
                default_work_queue = await models.workers.read_work_queue(
                    session=session,
                    work_queue_id=default_agent_work_pool.default_queue_id,
                )
                assert default_work_queue
                return default_work_queue
            data["work_pool_id"] = default_agent_work_pool.id

    # Set the priority to be the max priority + 1
    # This will make the new queue the lowest priority
    if data["priority"] is None:
        # Set the priority to be the first priority value that isn't already taken
        priorities_query = sa.select(orm_models.WorkQueue.priority).where(
            orm_models.WorkQueue.work_pool_id == data["work_pool_id"]
        )
        priorities = (await session.execute(priorities_query)).scalars().all()

        priority = None
        for i, p in enumerate(sorted(priorities)):
            # if a rank was skipped (e.g. the set priority is different than the
            # enumerated priority) then we can "take" that spot for this work
            # queue
            if i + 1 != p:
                priority = i + 1
                break

        # otherwise take the maximum priority plus one
        if priority is None:
            priority = max(priorities, default=0) + 1

        data["priority"] = priority

    model = orm_models.WorkQueue(**data)

    session.add(model)
    await session.flush()
    await session.refresh(model)

    if work_queue.priority:
        await bulk_update_work_queue_priorities(
            session=session,
            work_pool_id=data["work_pool_id"],
            new_priorities={model.id: work_queue.priority},
        )

    return model

delete_work_queue(session, work_queue_id) async

Delete a WorkQueue by id.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session	required
work_queue_id	str	a WorkQueue id	required

Returns:

Name	Type	Description
bool	bool	whether or not the WorkQueue was deleted

Source code in src/prefect/server/models/work_queues.py

async def delete_work_queue(session: AsyncSession, work_queue_id: UUID) -> bool:
    """
    Delete a WorkQueue by id.

    Args:
        session (AsyncSession): A database session
        work_queue_id (str): a WorkQueue id

    Returns:
        bool: whether or not the WorkQueue was deleted
    """
    result = await session.execute(
        delete(orm_models.WorkQueue).where(orm_models.WorkQueue.id == work_queue_id)
    )

    return result.rowcount > 0

ensure_work_queue_exists(session, name) async

Checks if a work queue exists and creates it if it does not.

Useful when working with deployments, agents, and flow runs that automatically create work queues.

Will also create a work pool queue in the default agent pool to facilitate migration to work pools.

Source code in src/prefect/server/models/work_queues.py

async def ensure_work_queue_exists(
    session: AsyncSession, name: str
) -> orm_models.WorkQueue:
    """
    Checks if a work queue exists and creates it if it does not.

    Useful when working with deployments, agents, and flow runs that automatically create work queues.

    Will also create a work pool queue in the default agent pool to facilitate migration to work pools.
    """
    # read work queue
    work_queue = await models.work_queues.read_work_queue_by_name(
        session=session, name=name
    )
    if not work_queue:
        default_pool = await models.workers.read_work_pool_by_name(
            session=session, work_pool_name=DEFAULT_AGENT_WORK_POOL_NAME
        )

        if default_pool is None:
            work_queue = await models.work_queues.create_work_queue(
                session=session,
                work_queue=schemas.actions.WorkQueueCreate(name=name, priority=1),
            )
        else:
            if name != "default":
                work_queue = await models.workers.create_work_queue(
                    session=session,
                    work_pool_id=default_pool.id,
                    work_queue=schemas.actions.WorkQueueCreate(name=name, priority=1),
                )
            else:
                work_queue = await models.work_queues.read_work_queue(
                    session=session, work_queue_id=default_pool.default_queue_id
                )
                assert work_queue, "Default work queue not found"

    return work_queue

get_runs_in_work_queue(db, session, work_queue_id, limit=None, scheduled_before=None) async

Get runs from a work queue.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session. work_queue_id: The work queue id.	required
scheduled_before	Optional[datetime]	Only return runs scheduled to start before this time.	None
limit	Optional[int]	An optional limit for the number of runs to return from the queue. This limit applies to the request only. It does not affect the work queue's concurrency limit. If limit exceeds the work queue's concurrency limit, it will be ignored.	None

Source code in src/prefect/server/models/work_queues.py

@db_injector
async def get_runs_in_work_queue(
    db: PrefectDBInterface,
    session: AsyncSession,
    work_queue_id: UUID,
    limit: Optional[int] = None,
    scheduled_before: Optional[datetime.datetime] = None,
) -> Tuple[orm_models.WorkQueue, Sequence[orm_models.FlowRun]]:
    """
    Get runs from a work queue.

    Args:
        session: A database session. work_queue_id: The work queue id.
        scheduled_before: Only return runs scheduled to start before this time.
        limit: An optional limit for the number of runs to return from the
            queue. This limit applies to the request only. It does not affect
            the work queue's concurrency limit. If `limit` exceeds the work
            queue's concurrency limit, it will be ignored.

    """
    work_queue = await read_work_queue(session=session, work_queue_id=work_queue_id)
    if not work_queue:
        raise ObjectNotFoundError(f"Work queue with id {work_queue_id} not found.")

    if work_queue.filter is None:
        query = db.queries.get_scheduled_flow_runs_from_work_queues(
            limit_per_queue=limit,
            work_queue_ids=[work_queue_id],
            scheduled_before=scheduled_before,
        )
        result = await session.execute(query)
        return work_queue, result.scalars().unique().all()

    # if the work queue has a filter, it's a deprecated tag-based work queue
    # and uses an old approach
    else:
        return work_queue, await _legacy_get_runs_in_work_queue(
            session=session,
            work_queue_id=work_queue_id,
            scheduled_before=scheduled_before,
            limit=limit,
        )

read_work_queue(session, work_queue_id) async

Reads a WorkQueue by id.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session	required
work_queue_id	str	a WorkQueue id	required

Returns:

Type	Description
Optional[WorkQueue]	orm_models.WorkQueue: the WorkQueue

Source code in src/prefect/server/models/work_queues.py

async def read_work_queue(
    session: AsyncSession, work_queue_id: Union[UUID, PrefectUUID]
) -> Optional[orm_models.WorkQueue]:
    """
    Reads a WorkQueue by id.

    Args:
        session (AsyncSession): A database session
        work_queue_id (str): a WorkQueue id

    Returns:
        orm_models.WorkQueue: the WorkQueue
    """

    return await session.get(orm_models.WorkQueue, work_queue_id)

read_work_queue_by_name(session, name) async

Reads a WorkQueue by id.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session	required
work_queue_id	str	a WorkQueue id	required

Returns:

Type	Description
Optional[WorkQueue]	orm_models.WorkQueue: the WorkQueue

Source code in src/prefect/server/models/work_queues.py

async def read_work_queue_by_name(
    session: AsyncSession, name: str
) -> Optional[orm_models.WorkQueue]:
    """
    Reads a WorkQueue by id.

    Args:
        session (AsyncSession): A database session
        work_queue_id (str): a WorkQueue id

    Returns:
        orm_models.WorkQueue: the WorkQueue
    """
    default_work_pool = await models.workers.read_work_pool_by_name(
        session=session, work_pool_name=DEFAULT_AGENT_WORK_POOL_NAME
    )
    # Logic to make sure this functionality doesn't break during migration
    if default_work_pool is not None:
        query = select(orm_models.WorkQueue).filter_by(
            name=name, work_pool_id=default_work_pool.id
        )
    else:
        query = select(orm_models.WorkQueue).filter_by(name=name)
    result = await session.execute(query)
    return result.scalar()

read_work_queue_status(session, work_queue_id) async

Get work queue status by id.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session	required
work_queue_id	str	a WorkQueue id	required

Returns:

Type	Description
WorkQueueStatusDetail	Information about the status of the work queue.

Source code in src/prefect/server/models/work_queues.py

async def read_work_queue_status(
    session: AsyncSession, work_queue_id: UUID
) -> schemas.core.WorkQueueStatusDetail:
    """
    Get work queue status by id.

    Args:
        session (AsyncSession): A database session
        work_queue_id (str): a WorkQueue id

    Returns:
        Information about the status of the work queue.
    """

    work_queue = await read_work_queue(session=session, work_queue_id=work_queue_id)
    if not work_queue:
        raise ObjectNotFoundError(f"Work queue with id {work_queue_id} not found")

    work_queue_late_runs_count = await models.flow_runs.count_flow_runs(
        session=session,
        flow_run_filter=schemas.filters.FlowRunFilter(
            state=schemas.filters.FlowRunFilterState(name={"any_": ["Late"]}),
        ),
        work_queue_filter=schemas.filters.WorkQueueFilter(
            id=schemas.filters.WorkQueueFilterId(any_=[work_queue_id])
        ),
    )

    # All work queues use the default policy for now
    health_check_policy = schemas.core.WorkQueueHealthPolicy(
        maximum_late_runs=0, maximum_seconds_since_last_polled=60
    )

    healthy = health_check_policy.evaluate_health_status(
        late_runs_count=work_queue_late_runs_count,
        last_polled=work_queue.last_polled,  # type: ignore
    )

    return schemas.core.WorkQueueStatusDetail(
        healthy=healthy,
        late_runs_count=work_queue_late_runs_count,
        last_polled=work_queue.last_polled,
        health_check_policy=health_check_policy,
    )

read_work_queues(session, offset=None, limit=None, work_queue_filter=None) async

Read WorkQueues.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session	required
offset	Optional[int]	Query offset	None
limit	Optional[int]	Query limit	None
work_queue_filter	Optional[WorkQueueFilter]	only select work queues matching these filters	None

Source code in src/prefect/server/models/work_queues.py

async def read_work_queues(
    session: AsyncSession,
    offset: Optional[int] = None,
    limit: Optional[int] = None,
    work_queue_filter: Optional[schemas.filters.WorkQueueFilter] = None,
) -> Sequence[orm_models.WorkQueue]:
    """
    Read WorkQueues.

    Args:
        session: A database session
        offset: Query offset
        limit: Query limit
        work_queue_filter: only select work queues matching these filters
    Returns:
        Sequence[orm_models.WorkQueue]: WorkQueues
    """

    query = select(orm_models.WorkQueue).order_by(orm_models.WorkQueue.name)

    if offset is not None:
        query = query.offset(offset)
    if limit is not None:
        query = query.limit(limit)
    if work_queue_filter:
        query = query.where(work_queue_filter.as_sql_filter())

    result = await session.execute(query)
    return result.scalars().unique().all()

record_work_queue_polls(session, polled_work_queue_ids, ready_work_queue_ids) async

Record that the given work queues were polled, and also update the given ready_work_queue_ids to READY.

Source code in src/prefect/server/models/work_queues.py

async def record_work_queue_polls(
    session: AsyncSession,
    polled_work_queue_ids: Sequence[UUID],
    ready_work_queue_ids: Sequence[UUID],
) -> None:
    """Record that the given work queues were polled, and also update the given
    ready_work_queue_ids to READY."""
    polled = pendulum.now("UTC")

    if polled_work_queue_ids:
        await session.execute(
            sa.update(orm_models.WorkQueue)
            .where(orm_models.WorkQueue.id.in_(polled_work_queue_ids))
            .values(last_polled=polled)
        )

    if ready_work_queue_ids:
        await session.execute(
            sa.update(orm_models.WorkQueue)
            .where(orm_models.WorkQueue.id.in_(ready_work_queue_ids))
            .values(last_polled=polled, status=WorkQueueStatus.READY)
        )

update_work_queue(session, work_queue_id, work_queue, emit_status_change=None) async

Update a WorkQueue by id.

Parameters:

Name	Type	Description	Default
session	AsyncSession	A database session	required
work_queue	WorkQueueUpdate	the work queue data	required
work_queue_id	str	a WorkQueue id	required

Returns:

Name	Type	Description
bool	bool	whether or not the WorkQueue was updated

Source code in src/prefect/server/models/work_queues.py

async def update_work_queue(
    session: AsyncSession,
    work_queue_id: UUID,
    work_queue: schemas.actions.WorkQueueUpdate,
    emit_status_change: Optional[
        Callable[[orm_models.WorkQueue], Awaitable[None]]
    ] = None,
) -> bool:
    """
    Update a WorkQueue by id.

    Args:
        session (AsyncSession): A database session
        work_queue: the work queue data
        work_queue_id (str): a WorkQueue id

    Returns:
        bool: whether or not the WorkQueue was updated
    """
    # exclude_unset=True allows us to only update values provided by
    # the user, ignoring any defaults on the model
    update_data = work_queue.model_dump_for_orm(exclude_unset=True)

    if "is_paused" in update_data:
        wq = await read_work_queue(session=session, work_queue_id=work_queue_id)
        if wq is None:
            return False

        # Only update the status to paused if it's not already paused. This ensures a work queue that is already
        # paused will not get a status update if it's paused again
        if update_data.get("is_paused") and wq.status != WorkQueueStatus.PAUSED:
            update_data["status"] = WorkQueueStatus.PAUSED

        # If unpausing, only update status if it's currently paused. This ensures a work queue that is already
        # unpaused will not get a status update if it's unpaused again
        if (
            update_data.get("is_paused") is False
            and wq.status == WorkQueueStatus.PAUSED
        ):
            # Default status if unpaused
            update_data["status"] = WorkQueueStatus.NOT_READY

            # Determine source of last_polled: update_data or database
            last_polled: Optional[pendulum.DateTime]
            if "last_polled" in update_data:
                last_polled = cast(pendulum.DateTime, update_data["last_polled"])
            else:
                last_polled = wq.last_polled

            # Check if last polled is recent and set status to READY if so
            if is_last_polled_recent(last_polled):
                update_data["status"] = schemas.statuses.WorkQueueStatus.READY

    update_stmt = (
        sa.update(orm_models.WorkQueue)
        .where(orm_models.WorkQueue.id == work_queue_id)
        .values(**update_data)
    )
    result = await session.execute(update_stmt)
    updated = result.rowcount > 0

    if updated:
        if "status" in update_data and emit_status_change:
            wq = await read_work_queue(session=session, work_queue_id=work_queue_id)
            assert wq
            await emit_status_change(wq)

    return updated