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

prefect/server/api/flows.py

@@ -0,0 +1,210 @@
+"""
+Routes for interacting with flow objects.
+"""
+
+from typing import List, Optional
+from uuid import UUID
+
+from fastapi import Depends, HTTPException, Path, Response, status
+from fastapi.param_functions import Body
+
+import prefect.server.api.dependencies as dependencies
+import prefect.server.models as models
+import prefect.server.schemas as schemas
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.schemas.responses import FlowPaginationResponse
+from prefect.server.utilities.server import PrefectRouter
+from prefect.types._datetime import now
+
+router: PrefectRouter = PrefectRouter(prefix="/flows", tags=["Flows"])
+
+
+@router.post("/")
+async def create_flow(
+    flow: schemas.actions.FlowCreate,
+    response: Response,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.Flow:
+    """Gracefully creates a new flow from the provided schema. If a flow with the
+    same name already exists, the existing flow is returned.
+    """
+    # hydrate the input model into a full flow model
+    flow = schemas.core.Flow(**flow.model_dump())
+
+    right_now = now("UTC")
+
+    async with db.session_context(begin_transaction=True) as session:
+        model = await models.flows.create_flow(session=session, flow=flow)
+
+    if model.created >= right_now:
+        response.status_code = status.HTTP_201_CREATED
+    return model
+
+
+@router.patch("/{id}", status_code=status.HTTP_204_NO_CONTENT)
+async def update_flow(
+    flow: schemas.actions.FlowUpdate,
+    flow_id: UUID = Path(..., description="The flow id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    """
+    Updates a flow.
+    """
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.flows.update_flow(
+            session=session, flow=flow, flow_id=flow_id
+        )
+    if not result:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Flow not found"
+        )
+
+
+@router.post("/count")
+async def count_flows(
+    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),
+) -> int:
+    """
+    Count flows.
+    """
+    async with db.session_context() as session:
+        return await models.flows.count_flows(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            work_pool_filter=work_pools,
+        )
+
+
+@router.get("/name/{name}")
+async def read_flow_by_name(
+    name: str = Path(..., description="The name of the flow"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.Flow:
+    """
+    Get a flow by name.
+    """
+    async with db.session_context() as session:
+        flow = await models.flows.read_flow_by_name(session=session, name=name)
+    if not flow:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Flow not found"
+        )
+    return flow
+
+
+@router.get("/{id}")
+async def read_flow(
+    flow_id: UUID = Path(..., description="The flow id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.Flow:
+    """
+    Get a flow by id.
+    """
+    async with db.session_context() as session:
+        flow = await models.flows.read_flow(session=session, flow_id=flow_id)
+    if not flow:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Flow not found"
+        )
+    return flow
+
+
+@router.post("/filter")
+async def read_flows(
+    limit: int = dependencies.LimitBody(),
+    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,
+    sort: schemas.sorting.FlowSort = Body(schemas.sorting.FlowSort.NAME_ASC),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.core.Flow]:
+    """
+    Query for flows.
+    """
+    async with db.session_context() as session:
+        return await models.flows.read_flows(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            work_pool_filter=work_pools,
+            sort=sort,
+            offset=offset,
+            limit=limit,
+        )
+
+
+@router.delete("/{id}", status_code=status.HTTP_204_NO_CONTENT)
+async def delete_flow(
+    flow_id: UUID = Path(..., description="The flow id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    """
+    Delete a flow by id.
+    """
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.flows.delete_flow(session=session, flow_id=flow_id)
+    if not result:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Flow not found"
+        )
+
+
+@router.post("/paginate")
+async def paginate_flows(
+    limit: int = dependencies.LimitBody(),
+    page: int = Body(1, ge=1),
+    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,
+    work_pools: Optional[schemas.filters.WorkPoolFilter] = None,
+    sort: schemas.sorting.FlowSort = Body(schemas.sorting.FlowSort.NAME_ASC),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> FlowPaginationResponse:
+    """
+    Pagination query for flows.
+    """
+    offset = (page - 1) * limit
+
+    async with db.session_context() as session:
+        results = await models.flows.read_flows(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            work_pool_filter=work_pools,
+            sort=sort,
+            offset=offset,
+            limit=limit,
+        )
+
+        count = await models.flows.count_flows(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            work_pool_filter=work_pools,
+        )
+
+    return FlowPaginationResponse(
+        results=results,
+        count=count,
+        limit=limit,
+        pages=(count + limit - 1) // limit,
+        page=page,
+    )

prefect/server/api/logs.py

@@ -0,0 +1,43 @@
+"""
+Routes for interacting with log objects.
+"""
+
+from typing import List
+
+from fastapi import Body, Depends, status
+
+import prefect.server.api.dependencies as dependencies
+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="/logs", tags=["Logs"])
+
+
+@router.post("/", status_code=status.HTTP_201_CREATED)
+async def create_logs(
+    logs: List[schemas.actions.LogCreate],
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    """Create new logs from the provided schema."""
+    for batch in models.logs.split_logs_into_batches(logs):
+        async with db.session_context(begin_transaction=True) as session:
+            await models.logs.create_logs(session=session, logs=batch)
+
+
+@router.post("/filter")
+async def read_logs(
+    limit: int = dependencies.LimitBody(),
+    offset: int = Body(0, ge=0),
+    logs: schemas.filters.LogFilter = None,
+    sort: schemas.sorting.LogSort = Body(schemas.sorting.LogSort.TIMESTAMP_ASC),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.core.Log]:
+    """
+    Query for logs.
+    """
+    async with db.session_context() as session:
+        return await models.logs.read_logs(
+            session=session, log_filter=logs, offset=offset, limit=limit, sort=sort
+        )

prefect/server/api/middleware.py

@@ -0,0 +1,73 @@
+from typing import Awaitable, Callable
+
+from fastapi import status
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+
+from prefect import settings
+from prefect.server import models
+from prefect.server.database import provide_database_interface
+
+NextMiddlewareFunction = Callable[[Request], Awaitable[Response]]
+
+
+class CsrfMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware for CSRF protection. This middleware will check for a CSRF token
+    in the headers of any POST, PUT, PATCH, or DELETE request. If the token is
+    not present or does not match the token stored in the database for the
+    client, the request will be rejected with a 403 status code.
+    """
+
+    async def dispatch(
+        self, request: Request, call_next: NextMiddlewareFunction
+    ) -> Response:
+        """
+        Dispatch method for the middleware. This method will check for the
+        presence of a CSRF token in the headers of the request and compare it
+        to the token stored in the database for the client. If the token is not
+        present or does not match, the request will be rejected with a 403
+        status code.
+        """
+
+        request_needs_csrf_protection = request.method in {
+            "POST",
+            "PUT",
+            "PATCH",
+            "DELETE",
+        }
+
+        if (
+            settings.PREFECT_SERVER_CSRF_PROTECTION_ENABLED.value()
+            and request_needs_csrf_protection
+        ):
+            incoming_token = request.headers.get("Prefect-Csrf-Token")
+            incoming_client = request.headers.get("Prefect-Csrf-Client")
+
+            if incoming_token is None:
+                return JSONResponse(
+                    {"detail": "Missing CSRF token."},
+                    status_code=status.HTTP_403_FORBIDDEN,
+                )
+
+            if incoming_client is None:
+                return JSONResponse(
+                    {"detail": "Missing client identifier."},
+                    status_code=status.HTTP_403_FORBIDDEN,
+                )
+
+            db = provide_database_interface()
+            async with db.session_context() as session:
+                token = await models.csrf_token.read_token_for_client(
+                    session=session, client=incoming_client
+                )
+
+                if token is None or token.token != incoming_token:
+                    return JSONResponse(
+                        {"detail": "Invalid CSRF token or client identifier."},
+                        status_code=status.HTTP_403_FORBIDDEN,
+                        headers={"Access-Control-Allow-Origin": "*"},
+                    )
+
+        return await call_next(request)

prefect/server/api/root.py

@@ -0,0 +1,35 @@
+"""
+Contains the `hello` route for testing and healthcheck purposes.
+"""
+
+from fastapi import Depends, status
+from fastapi.responses import JSONResponse
+
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.utilities.server import PrefectRouter
+
+router: PrefectRouter = PrefectRouter(prefix="", tags=["Root"])
+
+
+@router.get("/hello")
+async def hello() -> str:
+    """Say hello!"""
+    return "👋"
+
+
+@router.get("/ready")
+async def perform_readiness_check(
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> JSONResponse:
+    is_db_connectable = await db.is_db_connectable()
+
+    if is_db_connectable:
+        return JSONResponse(
+            status_code=status.HTTP_200_OK,
+            content={"message": "OK"},
+        )
+
+    return JSONResponse(
+        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+        content={"message": "Database is not available"},
+    )

prefect/server/api/run_history.py

@@ -0,0 +1,170 @@
+"""
+Utilities for querying flow and task run history.
+"""
+
+import datetime
+import json
+from typing import TYPE_CHECKING, List, Optional
+
+import pydantic
+import sqlalchemy as sa
+from typing_extensions import Literal
+
+import prefect.server.models as models
+import prefect.server.schemas as schemas
+from prefect.logging import get_logger
+from prefect.server.database import PrefectDBInterface, db_injector
+from prefect.types import DateTime
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("server.api")
+
+
+@db_injector
+async def run_history(
+    db: PrefectDBInterface,
+    session: sa.orm.Session,
+    run_type: Literal["flow_run", "task_run"],
+    history_start: DateTime,
+    history_end: DateTime,
+    history_interval: datetime.timedelta,
+    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,
+    work_pools: Optional[schemas.filters.WorkPoolFilter] = None,
+    work_queues: Optional[schemas.filters.WorkQueueFilter] = None,
+) -> List[schemas.responses.HistoryResponse]:
+    """
+    Produce a history of runs aggregated by interval and state
+    """
+
+    # SQLite has issues with very small intervals
+    # (by 0.001 seconds it stops incrementing the interval)
+    if history_interval < datetime.timedelta(seconds=1):
+        raise ValueError("History interval must not be less than 1 second.")
+
+    # prepare run-specific models
+    if run_type == "flow_run":
+        run_model = db.FlowRun
+        run_filter_function = models.flow_runs._apply_flow_run_filters
+    elif run_type == "task_run":
+        run_model = db.TaskRun
+        run_filter_function = models.task_runs._apply_task_run_filters
+    else:
+        raise ValueError(
+            f"Unknown run type {run_type!r}. Expected 'flow_run' or 'task_run'."
+        )
+
+    # create a CTE for timestamp intervals
+    intervals = db.queries.make_timestamp_intervals(
+        history_start,
+        history_end,
+        history_interval,
+    ).cte("intervals")
+
+    # apply filters to the flow runs (and related states)
+    runs = (
+        await run_filter_function(
+            db,
+            sa.select(
+                run_model.id,
+                run_model.expected_start_time,
+                run_model.estimated_run_time,
+                run_model.estimated_start_time_delta,
+                run_model.state_type,
+                run_model.state_name,
+            ).select_from(run_model),
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+            work_pool_filter=work_pools,
+            work_queue_filter=work_queues,
+        )
+    ).alias("runs")
+    # outer join intervals to the filtered runs to create a dataset composed of
+    # every interval and the aggregate of all its runs. The runs aggregate is represented
+    # by a descriptive JSON object
+    counts = (
+        sa.select(
+            intervals.c.interval_start,
+            intervals.c.interval_end,
+            # build a JSON object, ignoring the case where the count of runs is 0
+            sa.case(
+                (sa.func.count(runs.c.id) == 0, None),
+                else_=db.queries.build_json_object(
+                    "state_type",
+                    runs.c.state_type,
+                    "state_name",
+                    runs.c.state_name,
+                    "count_runs",
+                    sa.func.count(runs.c.id),
+                    # estimated run times only includes positive run times (to avoid any unexpected corner cases)
+                    "sum_estimated_run_time",
+                    sa.func.sum(
+                        sa.func.greatest(
+                            0, sa.extract("epoch", runs.c.estimated_run_time)
+                        )
+                    ),
+                    # estimated lateness is the sum of any positive start time deltas
+                    "sum_estimated_lateness",
+                    sa.func.sum(
+                        sa.func.greatest(
+                            0, sa.extract("epoch", runs.c.estimated_start_time_delta)
+                        )
+                    ),
+                ),
+            ).label("state_agg"),
+        )
+        .select_from(intervals)
+        .join(
+            runs,
+            sa.and_(
+                runs.c.expected_start_time >= intervals.c.interval_start,
+                runs.c.expected_start_time < intervals.c.interval_end,
+            ),
+            isouter=True,
+        )
+        .group_by(
+            intervals.c.interval_start,
+            intervals.c.interval_end,
+            runs.c.state_type,
+            runs.c.state_name,
+        )
+    ).alias("counts")
+
+    # aggregate all state-aggregate objects into a single array for each interval,
+    # ensuring that intervals with no runs have an empty array
+    query = (
+        sa.select(
+            counts.c.interval_start,
+            counts.c.interval_end,
+            sa.func.coalesce(
+                db.queries.json_arr_agg(
+                    db.queries.cast_to_json(counts.c.state_agg)
+                ).filter(counts.c.state_agg.is_not(None)),
+                sa.literal("[]", literal_execute=True),
+            ).label("states"),
+        )
+        .group_by(counts.c.interval_start, counts.c.interval_end)
+        .order_by(counts.c.interval_start)
+        # return no more than 500 bars
+        .limit(500)
+    )
+
+    # issue the query
+    result = await session.execute(query)
+    records = result.mappings()
+
+    # load and parse the record if the database returns JSON as strings
+    if db.queries.uses_json_strings:
+        records = [dict(r) for r in records]
+        for r in records:
+            r["states"] = json.loads(r["states"])
+
+    return pydantic.TypeAdapter(
+        List[schemas.responses.HistoryResponse]
+    ).validate_python(records)

prefect/server/api/saved_searches.py

@@ -0,0 +1,99 @@
+"""
+Routes for interacting with saved search objects.
+"""
+
+from typing import List
+from uuid import UUID
+
+from fastapi import Body, Depends, HTTPException, Path, Response, status
+
+import prefect.server.api.dependencies as dependencies
+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
+from prefect.types._datetime import now
+
+router: PrefectRouter = PrefectRouter(prefix="/saved_searches", tags=["SavedSearches"])
+
+
+@router.put("/")
+async def create_saved_search(
+    saved_search: schemas.actions.SavedSearchCreate,
+    response: Response,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.SavedSearch:
+    """Gracefully creates a new saved search from the provided schema.
+
+    If a saved search with the same name already exists, the saved search's fields are
+    replaced.
+    """
+
+    # hydrate the input model into a full model
+    saved_search = schemas.core.SavedSearch(**saved_search.model_dump())
+
+    right_now = now("UTC")
+
+    async with db.session_context(begin_transaction=True) as session:
+        model = await models.saved_searches.create_saved_search(
+            session=session, saved_search=saved_search
+        )
+
+    if model.created >= right_now:
+        response.status_code = status.HTTP_201_CREATED
+
+    return model
+
+
+@router.get("/{id}")
+async def read_saved_search(
+    saved_search_id: UUID = Path(..., description="The saved search id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.SavedSearch:
+    """
+    Get a saved search by id.
+    """
+    async with db.session_context() as session:
+        saved_search = await models.saved_searches.read_saved_search(
+            session=session, saved_search_id=saved_search_id
+        )
+    if not saved_search:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Saved search not found"
+        )
+    return saved_search
+
+
+@router.post("/filter")
+async def read_saved_searches(
+    limit: int = dependencies.LimitBody(),
+    offset: int = Body(0, ge=0),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.core.SavedSearch]:
+    """
+    Query for saved searches.
+    """
+    async with db.session_context() as session:
+        return await models.saved_searches.read_saved_searches(
+            session=session,
+            offset=offset,
+            limit=limit,
+        )
+
+
+@router.delete("/{id}", status_code=status.HTTP_204_NO_CONTENT)
+async def delete_saved_search(
+    saved_search_id: UUID = Path(..., description="The saved search id", alias="id"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    """
+    Delete a saved search by id.
+    """
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.saved_searches.delete_saved_search(
+            session=session, saved_search_id=saved_search_id
+        )
+    if not result:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Saved search not found"
+        )