prefect-client 3.2.2__py3-none-any.whl → 3.2.4__py3-none-any.whl

prefect/server/api/task_run_states.py

@@ -0,0 +1,52 @@
+"""
+Routes for interacting with task run state objects.
+"""
+
+from typing import List
+from uuid import UUID
+
+from fastapi import Depends, HTTPException, Path, status
+
+import prefect.server.models as models
+import prefect.server.schemas as schemas
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.utilities.server import PrefectRouter
+
+router: PrefectRouter = PrefectRouter(
+    prefix="/task_run_states", tags=["Task Run States"]
+)
+
+
+@router.get("/{id}")
+async def read_task_run_state(
+    task_run_state_id: UUID = Path(
+        ..., description="The task run state id", alias="id"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.states.State:
+    """
+    Get a task run state by id.
+    """
+    async with db.session_context() as session:
+        task_run_state = await models.task_run_states.read_task_run_state(
+            session=session, task_run_state_id=task_run_state_id
+        )
+    if not task_run_state:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Flow run state not found"
+        )
+    return task_run_state
+
+
+@router.get("/")
+async def read_task_run_states(
+    task_run_id: UUID,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.states.State]:
+    """
+    Get states associated with a task run.
+    """
+    async with db.session_context() as session:
+        return await models.task_run_states.read_task_run_states(
+            session=session, task_run_id=task_run_id
+        )

prefect/server/api/task_runs.py

@@ -0,0 +1,342 @@
+"""
+Routes for interacting with task run objects.
+"""
+
+import asyncio
+import datetime
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from uuid import UUID
+
+from fastapi import (
+    Body,
+    Depends,
+    HTTPException,
+    Path,
+    Response,
+    WebSocket,
+    status,
+)
+from starlette.websockets import WebSocketDisconnect
+
+import prefect.server.api.dependencies as dependencies
+import prefect.server.models as models
+import prefect.server.schemas as schemas
+from prefect.logging import get_logger
+from prefect.server.api.run_history import run_history
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.orchestration import dependencies as orchestration_dependencies
+from prefect.server.orchestration.core_policy import CoreTaskPolicy
+from prefect.server.orchestration.policies import TaskRunOrchestrationPolicy
+from prefect.server.schemas.responses import OrchestrationResult
+from prefect.server.task_queue import MultiQueue, TaskQueue
+from prefect.server.utilities import subscriptions
+from prefect.server.utilities.server import PrefectRouter
+from prefect.types import DateTime
+from prefect.types._datetime import now
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("server.api")
+
+router: PrefectRouter = PrefectRouter(prefix="/task_runs", tags=["Task Runs"])
+
+
+@router.post("/")
+async def create_task_run(
+    task_run: schemas.actions.TaskRunCreate,
+    response: Response,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+    orchestration_parameters: Dict[str, Any] = Depends(
+        orchestration_dependencies.provide_task_orchestration_parameters
+    ),
+) -> schemas.core.TaskRun:
+    """
+    Create a task run. If a task run with the same flow_run_id,
+    task_key, and dynamic_key already exists, the existing task
+    run will be returned.
+
+    If no state is provided, the task run will be created in a PENDING state.
+    """
+    # hydrate the input model into a full task run / state model
+    task_run_dict = task_run.model_dump()
+    if not task_run_dict.get("id"):
+        task_run_dict.pop("id", None)
+    task_run = schemas.core.TaskRun(**task_run_dict)
+
+    if not task_run.state:
+        task_run.state = schemas.states.Pending()
+
+    right_now = now("UTC")
+
+    async with db.session_context(begin_transaction=True) as session:
+        model = await models.task_runs.create_task_run(
+            session=session,
+            task_run=task_run,
+            orchestration_parameters=orchestration_parameters,
+        )
+
+    if model.created >= right_now:
+        response.status_code = status.HTTP_201_CREATED
+
+    new_task_run: schemas.core.TaskRun = schemas.core.TaskRun.model_validate(model)
+
+    return new_task_run
+
+
+@router.patch("/{id}", status_code=status.HTTP_204_NO_CONTENT)
+async def update_task_run(
+    task_run: schemas.actions.TaskRunUpdate,
+    task_run_id: UUID = Path(..., description="The task run id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    """
+    Updates a task run.
+    """
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.task_runs.update_task_run(
+            session=session, task_run=task_run, task_run_id=task_run_id
+        )
+    if not result:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Task run not found")
+
+
+@router.post("/count")
+async def count_task_runs(
+    db: PrefectDBInterface = Depends(provide_database_interface),
+    flows: schemas.filters.FlowFilter = None,
+    flow_runs: schemas.filters.FlowRunFilter = None,
+    task_runs: schemas.filters.TaskRunFilter = None,
+    deployments: schemas.filters.DeploymentFilter = None,
+) -> int:
+    """
+    Count task runs.
+    """
+    async with db.session_context() as session:
+        return await models.task_runs.count_task_runs(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+        )
+
+
+@router.post("/history")
+async def task_run_history(
+    history_start: DateTime = Body(..., description="The history's start time."),
+    history_end: DateTime = Body(..., description="The history's end time."),
+    # Workaround for the fact that FastAPI does not let us configure ser_json_timedelta
+    # to represent timedeltas as floats in JSON.
+    history_interval: float = Body(
+        ...,
+        description=(
+            "The size of each history interval, in seconds. Must be at least 1 second."
+        ),
+        json_schema_extra={"format": "time-delta"},
+        alias="history_interval_seconds",
+    ),
+    flows: schemas.filters.FlowFilter = None,
+    flow_runs: schemas.filters.FlowRunFilter = None,
+    task_runs: schemas.filters.TaskRunFilter = None,
+    deployments: schemas.filters.DeploymentFilter = None,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.responses.HistoryResponse]:
+    """
+    Query for task run history data across a given range and interval.
+    """
+    if isinstance(history_interval, float):
+        history_interval = datetime.timedelta(seconds=history_interval)
+
+    if history_interval < datetime.timedelta(seconds=1):
+        raise HTTPException(
+            status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="History interval must not be less than 1 second.",
+        )
+
+    async with db.session_context() as session:
+        return await run_history(
+            session=session,
+            run_type="task_run",
+            history_start=history_start,
+            history_end=history_end,
+            history_interval=history_interval,
+            flows=flows,
+            flow_runs=flow_runs,
+            task_runs=task_runs,
+            deployments=deployments,
+        )
+
+
+@router.get("/{id}")
+async def read_task_run(
+    task_run_id: UUID = Path(..., description="The task run id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.TaskRun:
+    """
+    Get a task run by id.
+    """
+    async with db.session_context() as session:
+        task_run = await models.task_runs.read_task_run(
+            session=session, task_run_id=task_run_id
+        )
+    if not task_run:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Task not found")
+    return task_run
+
+
+@router.post("/filter")
+async def read_task_runs(
+    sort: schemas.sorting.TaskRunSort = Body(schemas.sorting.TaskRunSort.ID_DESC),
+    limit: int = dependencies.LimitBody(),
+    offset: int = Body(0, ge=0),
+    flows: Optional[schemas.filters.FlowFilter] = None,
+    flow_runs: Optional[schemas.filters.FlowRunFilter] = None,
+    task_runs: Optional[schemas.filters.TaskRunFilter] = None,
+    deployments: Optional[schemas.filters.DeploymentFilter] = None,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.core.TaskRun]:
+    """
+    Query for task runs.
+    """
+    async with db.session_context() as session:
+        return await models.task_runs.read_task_runs(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            offset=offset,
+            limit=limit,
+            sort=sort,
+        )
+
+
+@router.delete("/{id}", status_code=status.HTTP_204_NO_CONTENT)
+async def delete_task_run(
+    task_run_id: UUID = Path(..., description="The task run id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    """
+    Delete a task run by id.
+    """
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.task_runs.delete_task_run(
+            session=session, task_run_id=task_run_id
+        )
+    if not result:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Task not found")
+
+
+@router.post("/{id}/set_state")
+async def set_task_run_state(
+    task_run_id: UUID = Path(..., description="The task run id", alias="id"),
+    state: schemas.actions.StateCreate = Body(..., description="The intended state."),
+    force: bool = Body(
+        False,
+        description=(
+            "If false, orchestration rules will be applied that may alter or prevent"
+            " the state transition. If True, orchestration rules are not applied."
+        ),
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+    response: Response = None,
+    task_policy: TaskRunOrchestrationPolicy = Depends(
+        orchestration_dependencies.provide_task_policy
+    ),
+    orchestration_parameters: Dict[str, Any] = Depends(
+        orchestration_dependencies.provide_task_orchestration_parameters
+    ),
+) -> OrchestrationResult:
+    """Set a task run state, invoking any orchestration rules."""
+
+    right_now = now("UTC")
+
+    # create the state
+    async with db.session_context(
+        begin_transaction=True, with_for_update=True
+    ) as session:
+        orchestration_result = await models.task_runs.set_task_run_state(
+            session=session,
+            task_run_id=task_run_id,
+            state=schemas.states.State.model_validate(
+                state
+            ),  # convert to a full State object
+            force=force,
+            task_policy=CoreTaskPolicy,
+            orchestration_parameters=orchestration_parameters,
+        )
+
+    # set the 201 if a new state was created
+    if orchestration_result.state and orchestration_result.state.timestamp >= right_now:
+        response.status_code = status.HTTP_201_CREATED
+    else:
+        response.status_code = status.HTTP_200_OK
+
+    return orchestration_result
+
+
+@router.websocket("/subscriptions/scheduled")
+async def scheduled_task_subscription(websocket: WebSocket) -> None:
+    websocket = await subscriptions.accept_prefect_socket(websocket)
+    if not websocket:
+        return
+
+    try:
+        subscription = await websocket.receive_json()
+    except subscriptions.NORMAL_DISCONNECT_EXCEPTIONS:
+        return
+
+    if subscription.get("type") != "subscribe":
+        return await websocket.close(
+            code=4001, reason="Protocol violation: expected 'subscribe' message"
+        )
+
+    task_keys = subscription.get("keys", [])
+    if not task_keys:
+        return await websocket.close(
+            code=4001, reason="Protocol violation: expected 'keys' in subscribe message"
+        )
+
+    if not (client_id := subscription.get("client_id")):
+        return await websocket.close(
+            code=4001,
+            reason="Protocol violation: expected 'client_id' in subscribe message",
+        )
+
+    subscribed_queue = MultiQueue(task_keys)
+
+    logger.info(f"Task worker {client_id!r} subscribed to task keys {task_keys!r}")
+
+    while True:
+        try:
+            # observe here so that all workers with active websockets are tracked
+            await models.task_workers.observe_worker(task_keys, client_id)
+            task_run = await asyncio.wait_for(subscribed_queue.get(), timeout=1)
+        except asyncio.TimeoutError:
+            if not await subscriptions.still_connected(websocket):
+                await models.task_workers.forget_worker(client_id)
+                return
+            continue
+
+        try:
+            await websocket.send_json(task_run.model_dump(mode="json"))
+
+            acknowledgement = await websocket.receive_json()
+            ack_type = acknowledgement.get("type")
+            if ack_type != "ack":
+                if ack_type == "quit":
+                    return await websocket.close()
+
+                raise WebSocketDisconnect(
+                    code=4001, reason="Protocol violation: expected 'ack' message"
+                )
+
+            await models.task_workers.observe_worker([task_run.task_key], client_id)
+
+        except subscriptions.NORMAL_DISCONNECT_EXCEPTIONS:
+            # If sending fails or pong fails, put the task back into the retry queue
+            await asyncio.shield(TaskQueue.for_key(task_run.task_key).retry(task_run))
+            return
+        finally:
+            await models.task_workers.forget_worker(client_id)

prefect/server/api/task_workers.py

@@ -0,0 +1,31 @@
+from typing import List, Optional
+
+from fastapi import Body
+from pydantic import BaseModel
+
+from prefect.server import models
+from prefect.server.models.task_workers import TaskWorkerResponse
+from prefect.server.utilities.server import PrefectRouter
+
+router: PrefectRouter = PrefectRouter(prefix="/task_workers", tags=["Task Workers"])
+
+
+class TaskWorkerFilter(BaseModel):
+    task_keys: List[str]
+
+
+@router.post("/filter")
+async def read_task_workers(
+    task_worker_filter: Optional[TaskWorkerFilter] = Body(
+        default=None, description="The task worker filter", embed=True
+    ),
+) -> List[TaskWorkerResponse]:
+    """Read active task workers. Optionally filter by task keys."""
+
+    if task_worker_filter and task_worker_filter.task_keys:
+        return await models.task_workers.get_workers_for_task_keys(
+            task_keys=task_worker_filter.task_keys,
+        )
+
+    else:
+        return await models.task_workers.get_all_workers()

prefect/server/api/templates.py

@@ -0,0 +1,35 @@
+import orjson
+from fastapi import Body, Response
+from jinja2.exceptions import TemplateSyntaxError
+
+from prefect.server.utilities.server import PrefectRouter
+from prefect.server.utilities.user_templates import (
+    TemplateSecurityError,
+    validate_user_template,
+)
+
+router: PrefectRouter = PrefectRouter(prefix="/templates", tags=["Automations"])
+
+
+@router.post(
+    "/validate",
+    response_class=Response,
+)
+def validate_template(template: str = Body(default="")) -> Response:
+    try:
+        validate_user_template(template)
+        return Response(content="", status_code=204)
+    except (TemplateSyntaxError, TemplateSecurityError) as e:
+        return Response(
+            status_code=422,
+            media_type="application/json",
+            content=orjson.dumps(
+                {
+                    "error": {
+                        "line": e.lineno,
+                        "message": e.message,
+                        "source": template,
+                    },
+                }
+            ),
+        )

prefect/server/api/ui/__init__.py

@@ -0,0 +1,3 @@
+"""Routes primarily for use by the UI"""
+
+from . import flows, flow_runs, schemas, task_runs

prefect/server/api/ui/flow_runs.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import datetime
+from typing import TYPE_CHECKING, List
+from uuid import UUID
+
+import sqlalchemy as sa
+from fastapi import Body, Depends
+from pydantic import Field
+
+import prefect.server.schemas as schemas
+from prefect.logging import get_logger
+from prefect.server import models
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.utilities.schemas.bases import PrefectBaseModel
+from prefect.server.utilities.server import PrefectRouter
+from prefect.types import DateTime
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("server.api.ui.flow_runs")
+
+router: PrefectRouter = PrefectRouter(prefix="/ui/flow_runs", tags=["Flow Runs", "UI"])
+
+
+class SimpleFlowRun(PrefectBaseModel):
+    id: UUID = Field(default=..., description="The flow run id.")
+    state_type: schemas.states.StateType = Field(
+        default=..., description="The state type."
+    )
+    timestamp: DateTime = Field(
+        default=...,
+        description=(
+            "The start time of the run, or the expected start time "
+            "if it hasn't run yet."
+        ),
+    )
+    duration: datetime.timedelta = Field(
+        default=..., description="The total run time of the run."
+    )
+    lateness: datetime.timedelta = Field(
+        default=..., description="The delay between the expected and actual start time."
+    )
+
+
+@router.post("/history")
+async def read_flow_run_history(
+    sort: schemas.sorting.FlowRunSort = Body(
+        schemas.sorting.FlowRunSort.EXPECTED_START_TIME_DESC
+    ),
+    limit: int = Body(1000, le=1000),
+    offset: int = Body(0, ge=0),
+    flows: schemas.filters.FlowFilter = None,
+    flow_runs: schemas.filters.FlowRunFilter = None,
+    task_runs: schemas.filters.TaskRunFilter = None,
+    deployments: schemas.filters.DeploymentFilter = None,
+    work_pools: schemas.filters.WorkPoolFilter = None,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[SimpleFlowRun]:
+    columns = [
+        db.FlowRun.id,
+        db.FlowRun.state_type,
+        db.FlowRun.start_time,
+        db.FlowRun.expected_start_time,
+        db.FlowRun.total_run_time,
+        # Although it isn't returned, we need to select
+        # this field in order to compute `estimated_run_time`
+        db.FlowRun.state_timestamp,
+    ]
+    async with db.session_context() as session:
+        result = await models.flow_runs.read_flow_runs(
+            columns=columns,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            work_pool_filter=work_pools,
+            sort=sort,
+            limit=limit,
+            offset=offset,
+            session=session,
+        )
+    return [
+        SimpleFlowRun(
+            id=r.id,
+            state_type=r.state_type,
+            timestamp=r.start_time or r.expected_start_time,
+            duration=r.estimated_run_time,
+            lateness=r.estimated_start_time_delta,
+        )
+        for r in result
+    ]
+
+
+@router.post("/count-task-runs")
+async def count_task_runs_by_flow_run(
+    flow_run_ids: list[UUID] = Body(default=..., embed=True, max_items=200),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> dict[UUID, int]:
+    """
+    Get task run counts by flow run id.
+    """
+    async with db.session_context() as session:
+        query = (
+            sa.select(
+                db.TaskRun.flow_run_id,
+                sa.func.count(db.TaskRun.id).label("task_run_count"),
+            )
+            .where(
+                sa.and_(
+                    db.TaskRun.flow_run_id.in_(flow_run_ids),
+                    sa.not_(db.TaskRun.subflow_run.has()),
+                )
+            )
+            .group_by(db.TaskRun.flow_run_id)
+        )
+
+        results = await session.execute(query)
+
+        task_run_counts_by_flow_run = {
+            flow_run_id: task_run_count for flow_run_id, task_run_count in results.t
+        }
+
+        return {
+            flow_run_id: task_run_counts_by_flow_run.get(flow_run_id, 0)
+            for flow_run_id in flow_run_ids
+        }