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

prefect/server/api/ui/flows.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Dict, List, Optional
+from uuid import UUID
+
+import sqlalchemy as sa
+from fastapi import Body, Depends
+from pydantic import Field, field_validator
+
+from prefect.logging import get_logger
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.schemas.states import StateType
+from prefect.server.utilities.database import UUID as UUIDTypeDecorator
+from prefect.server.utilities.schemas import PrefectBaseModel
+from prefect.server.utilities.server import PrefectRouter
+from prefect.types import DateTime
+from prefect.types._datetime import create_datetime_instance
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger()
+
+router: PrefectRouter = PrefectRouter(prefix="/ui/flows", tags=["Flows", "UI"])
+
+
+class SimpleNextFlowRun(PrefectBaseModel):
+    id: UUID = Field(default=..., description="The flow run id.")
+    flow_id: UUID = Field(default=..., description="The flow id.")
+    name: str = Field(default=..., description="The flow run name")
+    state_name: str = Field(default=..., description="The state name.")
+    state_type: StateType = Field(default=..., description="The state type.")
+    next_scheduled_start_time: DateTime = Field(
+        default=..., description="The next scheduled start time"
+    )
+
+    @field_validator("next_scheduled_start_time", mode="before")
+    @classmethod
+    def validate_next_scheduled_start_time(cls, v: DateTime | datetime) -> DateTime:
+        if isinstance(v, datetime):
+            return create_datetime_instance(v)
+        return v
+
+
+@router.post("/count-deployments")
+async def count_deployments_by_flow(
+    flow_ids: List[UUID] = Body(default=..., embed=True, max_items=200),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> Dict[UUID, int]:
+    """
+    Get deployment counts by flow id.
+    """
+    async with db.session_context() as session:
+        query = (
+            sa.select(
+                db.Deployment.flow_id,
+                sa.func.count(db.Deployment.id).label("deployment_count"),
+            )
+            .where(db.Deployment.flow_id.in_(flow_ids))
+            .group_by(db.Deployment.flow_id)
+        )
+
+        results = await session.execute(query)
+
+        deployment_counts_by_flow = {
+            flow_id: deployment_count for flow_id, deployment_count in results.all()
+        }
+
+        return {
+            flow_id: deployment_counts_by_flow.get(flow_id, 0) for flow_id in flow_ids
+        }
+
+
+def _get_postgres_next_runs_query(flow_ids: List[UUID]):
+    # Here we use the raw query because CROSS LATERAL JOINS are very
+    # difficult to express correctly in sqlalchemy.
+    raw_query = sa.text(
+        """
+        SELECT fr.id, fr.name, fr.flow_id, fr.state_name, fr.state_type, fr.state_name, fr.next_scheduled_start_time
+        FROM (
+            SELECT DISTINCT flow_id FROM flow_run
+            WHERE flow_id IN :flow_ids
+            AND state_type = 'SCHEDULED'
+            ) AS unique_flows
+        CROSS JOIN LATERAL (
+            SELECT *
+            FROM flow_run fr
+            WHERE fr.flow_id = unique_flows.flow_id
+            AND fr.state_type = 'SCHEDULED'
+            ORDER BY fr.next_scheduled_start_time ASC
+            LIMIT 1
+        ) fr;
+        """
+    )
+
+    bindparams = [
+        sa.bindparam(
+            "flow_ids",
+            flow_ids,
+            expanding=True,
+            type_=UUIDTypeDecorator,
+        ),
+    ]
+
+    query = raw_query.bindparams(*bindparams)
+    return query
+
+
+def _get_sqlite_next_runs_query(flow_ids: List[UUID]):
+    raw_query = sa.text(
+        """
+        WITH min_times AS (
+            SELECT flow_id, MIN(next_scheduled_start_time) AS min_next_scheduled_start_time
+            FROM flow_run
+            WHERE flow_id IN :flow_ids
+            AND state_type = 'SCHEDULED'
+            GROUP BY flow_id
+        )
+        SELECT fr.id, fr.name, fr.flow_id, fr.state_name, fr.state_type, fr.next_scheduled_start_time
+        FROM flow_run fr
+        JOIN min_times mt ON fr.flow_id = mt.flow_id AND fr.next_scheduled_start_time = mt.min_next_scheduled_start_time
+        WHERE fr.state_type = 'SCHEDULED';
+
+        """
+    )
+
+    bindparams = [
+        sa.bindparam(
+            "flow_ids",
+            flow_ids,
+            expanding=True,
+            type_=UUIDTypeDecorator,
+        ),
+    ]
+
+    query = raw_query.bindparams(*bindparams)
+    return query
+
+
+@router.post("/next-runs")
+async def next_runs_by_flow(
+    flow_ids: List[UUID] = Body(default=..., embed=True, max_items=200),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> Dict[UUID, Optional[SimpleNextFlowRun]]:
+    """
+    Get the next flow run by flow id.
+    """
+
+    async with db.session_context() as session:
+        if db.dialect.name == "postgresql":
+            query = _get_postgres_next_runs_query(flow_ids=flow_ids)
+        else:
+            query = _get_sqlite_next_runs_query(flow_ids=flow_ids)
+
+        results = await session.execute(query)
+
+        results_by_flow_id = {
+            UUID(str(result.flow_id)): SimpleNextFlowRun(
+                id=result.id,
+                flow_id=result.flow_id,
+                name=result.name,
+                state_name=result.state_name,
+                state_type=result.state_type,
+                next_scheduled_start_time=result.next_scheduled_start_time,
+            )
+            for result in results.all()
+        }
+
+        response = {
+            flow_id: results_by_flow_id.get(flow_id, None) for flow_id in flow_ids
+        }
+        return response

prefect/server/api/ui/schemas.py

@@ -0,0 +1,63 @@
+from typing import TYPE_CHECKING, Any
+
+from fastapi import Body, Depends, HTTPException, status
+
+from prefect.logging import get_logger
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.schemas.responses import SchemaValuesValidationResponse
+from prefect.server.utilities.server import APIRouter
+from prefect.utilities.schema_tools.hydration import HydrationContext, hydrate
+from prefect.utilities.schema_tools.validation import (
+    CircularSchemaRefError,
+    build_error_obj,
+    is_valid_schema,
+    preprocess_schema,
+    validate,
+)
+
+if TYPE_CHECKING:
+    import logging
+
+router: APIRouter = APIRouter(prefix="/ui/schemas", tags=["UI", "Schemas"])
+
+logger: "logging.Logger" = get_logger("server.api.ui.schemas")
+
+
+@router.post("/validate")
+async def validate_obj(
+    json_schema: dict[str, Any] = Body(
+        ...,
+        embed=True,
+        alias="schema",
+        json_schema_extra={"additionalProperties": True},
+    ),
+    values: dict[str, Any] = Body(
+        ..., embed=True, json_schema_extra={"additionalProperties": True}
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> SchemaValuesValidationResponse:
+    schema = preprocess_schema(json_schema)
+
+    try:
+        is_valid_schema(schema, preprocess=False)
+    except ValueError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)
+        )
+
+    async with db.session_context() as session:
+        ctx = await HydrationContext.build(
+            session=session, render_jinja=False, render_workspace_variables=True
+        )
+
+    hydrated_values = hydrate(values, ctx)
+    try:
+        errors = validate(hydrated_values, schema, preprocess=False)
+    except CircularSchemaRefError:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="Invalid schema: Unable to validate schema with circular references.",
+        )
+    error_obj = build_error_obj(errors)
+
+    return error_obj

prefect/server/api/ui/task_runs.py

@@ -0,0 +1,175 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, List, Optional
+
+import sqlalchemy as sa
+from fastapi import Depends, HTTPException, status
+from pydantic import Field, model_serializer
+
+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._datetime import end_of_period, now
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("server.api.ui.task_runs")
+
+router: PrefectRouter = PrefectRouter(prefix="/ui/task_runs", tags=["Task Runs", "UI"])
+
+FAILED_STATES = [schemas.states.StateType.CRASHED, schemas.states.StateType.FAILED]
+
+
+class TaskRunCount(PrefectBaseModel):
+    completed: int = Field(
+        default=..., description="The number of completed task runs."
+    )
+    failed: int = Field(default=..., description="The number of failed task runs.")
+
+    @model_serializer
+    def ser_model(self) -> dict[str, int]:
+        return {
+            "completed": int(self.completed),
+            "failed": int(self.failed),
+        }
+
+
+@router.post("/dashboard/counts")
+async def read_dashboard_task_run_counts(
+    task_runs: schemas.filters.TaskRunFilter,
+    flows: Optional[schemas.filters.FlowFilter] = None,
+    flow_runs: Optional[schemas.filters.FlowRunFilter] = None,
+    deployments: Optional[schemas.filters.DeploymentFilter] = None,
+    work_pools: Optional[schemas.filters.WorkPoolFilter] = None,
+    work_queues: Optional[schemas.filters.WorkQueueFilter] = None,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[TaskRunCount]:
+    if task_runs.start_time is None or task_runs.start_time.after_ is None:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="task_runs.start_time.after_ is required",
+        )
+
+    # We only care about task runs that are in a terminal state, all others
+    # should be ignored.
+    task_runs.state = schemas.filters.TaskRunFilterState(
+        type=schemas.filters.TaskRunFilterStateType(
+            any_=list(schemas.states.TERMINAL_STATES)
+        )
+    )
+
+    bucket_count = 20
+    start_time = task_runs.start_time.after_.start_of("minute")
+    end_time = (
+        end_of_period(task_runs.start_time.before_, "minute")
+        if task_runs.start_time.before_
+        else end_of_period(now("UTC"), "minute")
+    )
+    window = end_time - start_time
+    delta = window.as_timedelta() / bucket_count
+
+    async with db.session_context(begin_transaction=False) as session:
+        # Gather the raw counts. The counts are divided into buckets of time
+        # and each bucket contains the number of successful and failed task
+        # runs.
+        # SQLAlchemy doesn't play nicely with our DateTime type so we convert it
+        # to a datetime object.
+        start_datetime = datetime(
+            start_time.year,
+            start_time.month,
+            start_time.day,
+            start_time.hour,
+            start_time.minute,
+            start_time.second,
+            start_time.microsecond,
+            start_time.timezone,
+        )
+        bucket_expression = sa.func.floor(
+            sa.func.date_diff_seconds(db.TaskRun.start_time, start_datetime)
+            / delta.total_seconds()
+        ).label("bucket")
+
+        raw_counts = (
+            (
+                await models.task_runs._apply_task_run_filters(
+                    db,
+                    sa.select(
+                        bucket_expression,
+                        sa.func.min(db.TaskRun.end_time).label("oldest"),
+                        sa.func.sum(
+                            sa.case(
+                                (
+                                    db.TaskRun.state_type.in_(FAILED_STATES),
+                                    1,
+                                ),
+                                else_=0,
+                            )
+                        ).label("failed_count"),
+                        sa.func.sum(
+                            sa.case(
+                                (
+                                    db.TaskRun.state_type.notin_(FAILED_STATES),
+                                    1,
+                                ),
+                                else_=0,
+                            )
+                        ).label("successful_count"),
+                    ),
+                    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,
+                )
+            )
+            .group_by("bucket", db.TaskRun.start_time)
+            .subquery()
+        )
+
+        # Aggregate the raw counts by bucket
+        query = (
+            sa.select(
+                raw_counts.c.bucket.label("bucket"),
+                sa.func.min(raw_counts.c.oldest).label("oldest"),
+                sa.func.sum(raw_counts.c.failed_count).label("failed_count"),
+                sa.func.sum(raw_counts.c.successful_count).label("successful_count"),
+            )
+            .select_from(raw_counts)
+            .group_by(raw_counts.c.bucket)
+            .order_by(sa.asc("oldest"))
+        )
+
+        result = await session.execute(query)
+
+    # Ensure that all buckets of time are present in the result even if no
+    # matching task runs occurred during the given time period.
+    buckets = [TaskRunCount(completed=0, failed=0) for _ in range(bucket_count)]
+
+    for row in result:
+        index = int(row.bucket)
+        buckets[index].completed = row.successful_count
+        buckets[index].failed = row.failed_count
+
+    return buckets
+
+
+@router.post("/count")
+async def read_task_run_counts_by_state(
+    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),
+) -> schemas.states.CountByState:
+    async with db.session_context(begin_transaction=False) as session:
+        return await models.task_runs.count_task_runs_by_state(
+            session=session,
+            flow_filter=flows,
+            flow_run_filter=flow_runs,
+            task_run_filter=task_runs,
+            deployment_filter=deployments,
+        )