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

prefect/server/api/clients.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from urllib.parse import quote
+from uuid import UUID
+
+import httpx
+import pydantic
+from httpx import Response
+from starlette import status
+from typing_extensions import Self
+
+from prefect.client.base import PrefectHttpxAsyncClient
+from prefect.exceptions import ObjectNotFound
+from prefect.logging import get_logger
+from prefect.server.schemas.actions import DeploymentFlowRunCreate, StateCreate
+from prefect.server.schemas.core import WorkPool
+from prefect.server.schemas.filters import VariableFilter, VariableFilterName
+from prefect.server.schemas.responses import DeploymentResponse, OrchestrationResult
+from prefect.settings import PREFECT_SERVER_API_AUTH_STRING
+from prefect.types import StrictVariableValue
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger(__name__)
+
+
+class BaseClient:
+    _http_client: PrefectHttpxAsyncClient
+
+    def __init__(self, additional_headers: dict[str, str] | None = None):
+        from prefect.server.api.server import create_app
+
+        additional_headers = additional_headers or {}
+
+        # create_app caches application instances, and invoking it with no arguments
+        # will point it to the the currently running server instance
+        api_app = create_app()
+
+        # we pull the auth string from _server_ settings because this client is run on the server
+        auth_string = PREFECT_SERVER_API_AUTH_STRING.value()
+
+        if auth_string:
+            token = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")
+            additional_headers.setdefault("Authorization", f"Basic {token}")
+
+        self._http_client = PrefectHttpxAsyncClient(
+            transport=httpx.ASGITransport(app=api_app, raise_app_exceptions=False),
+            headers={**additional_headers},
+            base_url="http://prefect-in-memory/api",
+            enable_csrf_support=False,
+            raise_on_all_errors=False,
+        )
+
+    async def __aenter__(self) -> Self:
+        await self._http_client.__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self._http_client.__aexit__(*args)
+
+
+class OrchestrationClient(BaseClient):
+    async def read_deployment_raw(self, deployment_id: UUID) -> Response:
+        return await self._http_client.get(f"/deployments/{deployment_id}")
+
+    async def read_deployment(
+        self, deployment_id: UUID
+    ) -> Optional[DeploymentResponse]:
+        try:
+            response = await self.read_deployment_raw(deployment_id)
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
+                return None
+            raise
+        return DeploymentResponse.model_validate(response.json())
+
+    async def read_flow_raw(self, flow_id: UUID) -> Response:
+        return await self._http_client.get(f"/flows/{flow_id}")
+
+    async def create_flow_run(
+        self, deployment_id: UUID, flow_run_create: DeploymentFlowRunCreate
+    ) -> Response:
+        return await self._http_client.post(
+            f"/deployments/{deployment_id}/create_flow_run",
+            json=flow_run_create.model_dump(mode="json"),
+        )
+
+    async def read_flow_run_raw(self, flow_run_id: UUID) -> Response:
+        return await self._http_client.get(f"/flow_runs/{flow_run_id}")
+
+    async def read_task_run_raw(self, task_run_id: UUID) -> Response:
+        return await self._http_client.get(f"/task_runs/{task_run_id}")
+
+    async def resume_flow_run(self, flow_run_id: UUID) -> OrchestrationResult:
+        response = await self._http_client.post(
+            f"/flow_runs/{flow_run_id}/resume",
+        )
+        response.raise_for_status()
+        return OrchestrationResult.model_validate(response.json())
+
+    async def pause_deployment(self, deployment_id: UUID) -> Response:
+        return await self._http_client.post(
+            f"/deployments/{deployment_id}/pause_deployment",
+        )
+
+    async def resume_deployment(self, deployment_id: UUID) -> Response:
+        return await self._http_client.post(
+            f"/deployments/{deployment_id}/resume_deployment",
+        )
+
+    async def set_flow_run_state(
+        self, flow_run_id: UUID, state: StateCreate
+    ) -> Response:
+        return await self._http_client.post(
+            f"/flow_runs/{flow_run_id}/set_state",
+            json={
+                "state": state.model_dump(mode="json"),
+                "force": False,
+            },
+        )
+
+    async def pause_work_pool(self, work_pool_name: str) -> Response:
+        return await self._http_client.patch(
+            f"/work_pools/{quote(work_pool_name)}", json={"is_paused": True}
+        )
+
+    async def resume_work_pool(self, work_pool_name: str) -> Response:
+        return await self._http_client.patch(
+            f"/work_pools/{quote(work_pool_name)}", json={"is_paused": False}
+        )
+
+    async def read_work_pool_raw(self, work_pool_id: UUID) -> Response:
+        return await self._http_client.post(
+            "/work_pools/filter",
+            json={"work_pools": {"id": {"any_": [str(work_pool_id)]}}},
+        )
+
+    async def read_work_pool(self, work_pool_id: UUID) -> Optional[WorkPool]:
+        response = await self.read_work_pool_raw(work_pool_id)
+        response.raise_for_status()
+
+        pools = pydantic.TypeAdapter(List[WorkPool]).validate_python(response.json())
+        return pools[0] if pools else None
+
+    async def read_work_queue_raw(self, work_queue_id: UUID) -> Response:
+        return await self._http_client.get(f"/work_queues/{work_queue_id}")
+
+    async def read_work_queue_status_raw(self, work_queue_id: UUID) -> Response:
+        return await self._http_client.get(f"/work_queues/{work_queue_id}/status")
+
+    async def pause_work_queue(self, work_queue_id: UUID) -> Response:
+        return await self._http_client.patch(
+            f"/work_queues/{work_queue_id}",
+            json={"is_paused": True},
+        )
+
+    async def resume_work_queue(self, work_queue_id: UUID) -> Response:
+        return await self._http_client.patch(
+            f"/work_queues/{work_queue_id}",
+            json={"is_paused": False},
+        )
+
+    async def read_block_document_raw(
+        self,
+        block_document_id: UUID,
+        include_secrets: bool = True,
+    ) -> Response:
+        return await self._http_client.get(
+            f"/block_documents/{block_document_id}",
+            params=dict(include_secrets=include_secrets),
+        )
+
+    VARIABLE_PAGE_SIZE = 200
+    MAX_VARIABLES_PER_WORKSPACE = 1000
+
+    async def read_workspace_variables(
+        self, names: Optional[List[str]] = None
+    ) -> Dict[str, StrictVariableValue]:
+        variables: Dict[str, StrictVariableValue] = {}
+
+        offset = 0
+
+        filter = VariableFilter()
+
+        if names is not None and not names:
+            return variables
+        elif names is not None:
+            filter.name = VariableFilterName(any_=list(set(names)))
+
+        for offset in range(
+            0, self.MAX_VARIABLES_PER_WORKSPACE, self.VARIABLE_PAGE_SIZE
+        ):
+            response = await self._http_client.post(
+                "/variables/filter",
+                json={
+                    "variables": filter.model_dump(),
+                    "limit": self.VARIABLE_PAGE_SIZE,
+                    "offset": offset,
+                },
+            )
+            if response.status_code >= 300:
+                response.raise_for_status()
+
+            results = response.json()
+            for variable in results:
+                variables[variable["name"]] = variable["value"]
+
+            if len(results) < self.VARIABLE_PAGE_SIZE:
+                break
+
+        return variables
+
+    async def read_concurrency_limit_v2_raw(
+        self, concurrency_limit_id: UUID
+    ) -> Response:
+        return await self._http_client.get(
+            f"/v2/concurrency_limits/{concurrency_limit_id}"
+        )
+
+
+class WorkPoolsOrchestrationClient(BaseClient):
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def read_work_pool(self, work_pool_name: str) -> WorkPool:
+        """
+        Reads information for a given work pool
+        Args:
+            work_pool_name: The name of the work pool to for which to get
+                information.
+        Returns:
+            Information about the requested work pool.
+        """
+        try:
+            response = await self._http_client.get(f"/work_pools/{work_pool_name}")
+            response.raise_for_status()
+            return WorkPool.model_validate(response.json())
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise

prefect/server/api/collections.py

@@ -0,0 +1,75 @@
+import json
+from typing import Any, Dict
+
+import httpx
+from anyio import Path
+from cachetools import TTLCache
+from fastapi import HTTPException, status
+
+from prefect.server.utilities.server import PrefectRouter
+
+router: PrefectRouter = PrefectRouter(prefix="/collections", tags=["Collections"])
+
+GLOBAL_COLLECTIONS_VIEW_CACHE: TTLCache[str, dict[str, Any]] = TTLCache(
+    maxsize=200, ttl=60 * 10
+)
+
+REGISTRY_VIEWS = (
+    "https://raw.githubusercontent.com/PrefectHQ/prefect-collection-registry/main/views"
+)
+KNOWN_VIEWS = {
+    "aggregate-block-metadata": f"{REGISTRY_VIEWS}/aggregate-block-metadata.json",
+    "aggregate-flow-metadata": f"{REGISTRY_VIEWS}/aggregate-flow-metadata.json",
+    "aggregate-worker-metadata": f"{REGISTRY_VIEWS}/aggregate-worker-metadata.json",
+    "demo-flows": f"{REGISTRY_VIEWS}/demo-flows.json",
+}
+
+
+@router.get("/views/{view}")
+async def read_view_content(view: str) -> Dict[str, Any]:
+    """Reads the content of a view from the prefect-collection-registry."""
+    try:
+        return await get_collection_view(view)
+    except KeyError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"View {view} not found in registry",
+        )
+    except httpx.HTTPStatusError as exc:
+        if exc.response.status_code == 404:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail=f"Requested content missing for view {view}",
+            )
+        else:
+            raise
+
+
+async def get_collection_view(view: str) -> dict[str, Any]:
+    try:
+        return GLOBAL_COLLECTIONS_VIEW_CACHE[view]
+    except KeyError:
+        pass
+
+    try:
+        async with httpx.AsyncClient() as client:
+            resp = await client.get(KNOWN_VIEWS[view])
+            resp.raise_for_status()
+
+            data = resp.json()
+            if view == "aggregate-worker-metadata":
+                data.get("prefect", {}).pop("prefect-agent", None)
+
+            GLOBAL_COLLECTIONS_VIEW_CACHE[view] = data
+            return data
+    except Exception:
+        if view not in KNOWN_VIEWS:
+            raise
+        local_file = Path(__file__).parent / Path(f"collections_data/views/{view}.json")
+        if await local_file.exists():
+            raw_data = await local_file.read_text()
+            data = json.loads(raw_data)
+            GLOBAL_COLLECTIONS_VIEW_CACHE[view] = data
+            return data
+        else:
+            raise

prefect/server/api/concurrency_limits.py

@@ -0,0 +1,286 @@
+"""
+Routes for interacting with concurrency limit objects.
+"""
+
+from typing import List, Optional, Sequence
+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.api.concurrency_limits_v2 import MinimalConcurrencyLimitResponse
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.models import concurrency_limits
+from prefect.server.utilities.server import PrefectRouter
+from prefect.settings import PREFECT_TASK_RUN_TAG_CONCURRENCY_SLOT_WAIT_SECONDS
+from prefect.types._datetime import now
+
+router: PrefectRouter = PrefectRouter(
+    prefix="/concurrency_limits", tags=["Concurrency Limits"]
+)
+
+
+@router.post("/")
+async def create_concurrency_limit(
+    concurrency_limit: schemas.actions.ConcurrencyLimitCreate,
+    response: Response,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.ConcurrencyLimit:
+    # hydrate the input model into a full model
+    concurrency_limit_model = schemas.core.ConcurrencyLimit(
+        **concurrency_limit.model_dump()
+    )
+
+    async with db.session_context(begin_transaction=True) as session:
+        model = await models.concurrency_limits.create_concurrency_limit(
+            session=session, concurrency_limit=concurrency_limit_model
+        )
+
+    if model.created >= now("UTC"):
+        response.status_code = status.HTTP_201_CREATED
+
+    return model
+
+
+@router.get("/{id}")
+async def read_concurrency_limit(
+    concurrency_limit_id: UUID = Path(
+        ..., description="The concurrency limit id", alias="id"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.ConcurrencyLimit:
+    """
+    Get a concurrency limit by id.
+
+    The `active slots` field contains a list of TaskRun IDs currently using a
+    concurrency slot for the specified tag.
+    """
+    async with db.session_context() as session:
+        model = await models.concurrency_limits.read_concurrency_limit(
+            session=session, concurrency_limit_id=concurrency_limit_id
+        )
+    if not model:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency limit not found"
+        )
+    return model
+
+
+@router.get("/tag/{tag}")
+async def read_concurrency_limit_by_tag(
+    tag: str = Path(..., description="The tag name", alias="tag"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.ConcurrencyLimit:
+    """
+    Get a concurrency limit by tag.
+
+    The `active slots` field contains a list of TaskRun IDs currently using a
+    concurrency slot for the specified tag.
+    """
+
+    async with db.session_context() as session:
+        model = await models.concurrency_limits.read_concurrency_limit_by_tag(
+            session=session, tag=tag
+        )
+
+    if not model:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND, detail="Concurrency limit not found"
+        )
+    return model
+
+
+@router.post("/filter")
+async def read_concurrency_limits(
+    limit: int = dependencies.LimitBody(),
+    offset: int = Body(0, ge=0),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> Sequence[schemas.core.ConcurrencyLimit]:
+    """
+    Query for concurrency limits.
+
+    For each concurrency limit the `active slots` field contains a list of TaskRun IDs
+    currently using a concurrency slot for the specified tag.
+    """
+    async with db.session_context() as session:
+        return await models.concurrency_limits.read_concurrency_limits(
+            session=session,
+            limit=limit,
+            offset=offset,
+        )
+
+
+@router.post("/tag/{tag}/reset")
+async def reset_concurrency_limit_by_tag(
+    tag: str = Path(..., description="The tag name"),
+    slot_override: Optional[List[UUID]] = Body(
+        None,
+        embed=True,
+        description="Manual override for active concurrency limit slots.",
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    async with db.session_context(begin_transaction=True) as session:
+        model = await models.concurrency_limits.reset_concurrency_limit_by_tag(
+            session=session, tag=tag, slot_override=slot_override
+        )
+    if not model:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency limit not found"
+        )
+
+
+@router.delete("/{id}")
+async def delete_concurrency_limit(
+    concurrency_limit_id: UUID = Path(
+        ..., description="The concurrency limit id", alias="id"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.concurrency_limits.delete_concurrency_limit(
+            session=session, concurrency_limit_id=concurrency_limit_id
+        )
+    if not result:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency limit not found"
+        )
+
+
+@router.delete("/tag/{tag}")
+async def delete_concurrency_limit_by_tag(
+    tag: str = Path(..., description="The tag name"),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    async with db.session_context(begin_transaction=True) as session:
+        result = await models.concurrency_limits.delete_concurrency_limit_by_tag(
+            session=session, tag=tag
+        )
+    if not result:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency limit not found"
+        )
+
+
+class Abort(Exception):
+    def __init__(self, reason: str):
+        self.reason = reason
+
+
+class Delay(Exception):
+    def __init__(self, delay_seconds: float, reason: str):
+        self.delay_seconds = delay_seconds
+        self.reason = reason
+
+
+@router.post("/increment")
+async def increment_concurrency_limits_v1(
+    names: List[str] = Body(..., description="The tags to acquire a slot for"),
+    task_run_id: UUID = Body(
+        ..., description="The ID of the task run acquiring the slot"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[MinimalConcurrencyLimitResponse]:
+    applied_limits = {}
+
+    async with db.session_context(begin_transaction=True) as session:
+        try:
+            applied_limits = {}
+            filtered_limits = (
+                await concurrency_limits.filter_concurrency_limits_for_orchestration(
+                    session, tags=names
+                )
+            )
+            run_limits = {limit.tag: limit for limit in filtered_limits}
+            for tag, cl in run_limits.items():
+                limit = cl.concurrency_limit
+                if limit == 0:
+                    # limits of 0 will deadlock, and the transition needs to abort
+                    for stale_tag in applied_limits.keys():
+                        stale_limit = run_limits.get(stale_tag, None)
+                        active_slots = set(stale_limit.active_slots)
+                        active_slots.discard(str(task_run_id))
+                        stale_limit.active_slots = list(active_slots)
+
+                    raise Abort(
+                        reason=(
+                            f'The concurrency limit on tag "{tag}" is 0 and will '
+                            "deadlock if the task tries to run again."
+                        ),
+                    )
+                elif len(cl.active_slots) >= limit:
+                    # if the limit has already been reached, delay the transition
+                    for stale_tag in applied_limits.keys():
+                        stale_limit = run_limits.get(stale_tag, None)
+                        active_slots = set(stale_limit.active_slots)
+                        active_slots.discard(str(task_run_id))
+                        stale_limit.active_slots = list(active_slots)
+
+                    raise Delay(
+                        delay_seconds=PREFECT_TASK_RUN_TAG_CONCURRENCY_SLOT_WAIT_SECONDS.value(),
+                        reason=f"Concurrency limit for the {tag} tag has been reached",
+                    )
+                else:
+                    # log the TaskRun ID to active_slots
+                    applied_limits[tag] = cl
+                    active_slots = set(cl.active_slots)
+                    active_slots.add(str(task_run_id))
+                    cl.active_slots = list(active_slots)
+        except Exception as e:
+            for tag in applied_limits.keys():
+                cl = await concurrency_limits.read_concurrency_limit_by_tag(
+                    session, tag
+                )
+                active_slots = set(cl.active_slots)
+                active_slots.discard(str(task_run_id))
+                cl.active_slots = list(active_slots)
+
+            if isinstance(e, Delay):
+                raise HTTPException(
+                    status_code=status.HTTP_423_LOCKED,
+                    detail=e.reason,
+                    headers={"Retry-After": str(e.delay_seconds)},
+                )
+            elif isinstance(e, Abort):
+                raise HTTPException(
+                    status_code=status.HTTP_423_LOCKED,
+                    detail=e.reason,
+                )
+            else:
+                raise
+    return [
+        MinimalConcurrencyLimitResponse(
+            name=limit.tag, limit=limit.concurrency_limit, id=limit.id
+        )
+        for limit in applied_limits.values()
+    ]
+
+
+@router.post("/decrement")
+async def decrement_concurrency_limits_v1(
+    names: List[str] = Body(..., description="The tags to release a slot for"),
+    task_run_id: UUID = Body(
+        ..., description="The ID of the task run releasing the slot"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    async with db.session_context(begin_transaction=True) as session:
+        filtered_limits = (
+            await concurrency_limits.filter_concurrency_limits_for_orchestration(
+                session, tags=names
+            )
+        )
+        run_limits = {limit.tag: limit for limit in filtered_limits}
+        for tag, cl in run_limits.items():
+            active_slots = set(cl.active_slots)
+            active_slots.discard(str(task_run_id))
+            cl.active_slots = list(active_slots)
+
+    return [
+        MinimalConcurrencyLimitResponse(
+            name=limit.tag, limit=limit.concurrency_limit, id=limit.id
+        )
+        for limit in run_limits.values()
+    ]