prefect-client 3.2.1__py3-none-any.whl → 3.2.3__py3-none-any.whl

prefect/server/api/concurrency_limits_v2.py

@@ -0,0 +1,269 @@
+from typing import List, Literal, Optional, Union
+from uuid import UUID
+
+from fastapi import Body, Depends, HTTPException, Path, status
+
+import prefect.server.models as models
+import prefect.server.schemas as schemas
+from prefect.server.api.dependencies import LimitBody
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.schemas import actions
+from prefect.server.utilities.schemas import PrefectBaseModel
+from prefect.server.utilities.server import PrefectRouter
+
+router: PrefectRouter = PrefectRouter(
+    prefix="/v2/concurrency_limits", tags=["Concurrency Limits V2"]
+)
+
+
+@router.post("/", status_code=status.HTTP_201_CREATED)
+async def create_concurrency_limit_v2(
+    concurrency_limit: actions.ConcurrencyLimitV2Create,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.core.ConcurrencyLimitV2:
+    async with db.session_context(begin_transaction=True) as session:
+        model = await models.concurrency_limits_v2.create_concurrency_limit(
+            session=session, concurrency_limit=concurrency_limit
+        )
+
+    return schemas.core.ConcurrencyLimitV2.model_validate(model)
+
+
+@router.get("/{id_or_name}")
+async def read_concurrency_limit_v2(
+    id_or_name: Union[UUID, str] = Path(
+        ..., description="The ID or name of the concurrency limit", alias="id_or_name"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> schemas.responses.GlobalConcurrencyLimitResponse:
+    if isinstance(id_or_name, str):  # TODO: this seems like it shouldn't be necessary
+        try:
+            id_or_name = UUID(id_or_name)
+        except ValueError:
+            pass
+    async with db.session_context() as session:
+        if isinstance(id_or_name, UUID):
+            model = await models.concurrency_limits_v2.read_concurrency_limit(
+                session, concurrency_limit_id=id_or_name
+            )
+        else:
+            model = await models.concurrency_limits_v2.read_concurrency_limit(
+                session, name=id_or_name
+            )
+
+    if not model:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency Limit not found"
+        )
+
+    return schemas.responses.GlobalConcurrencyLimitResponse.model_validate(model)
+
+
+@router.post("/filter")
+async def read_all_concurrency_limits_v2(
+    limit: int = LimitBody(),
+    offset: int = Body(0, ge=0),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[schemas.responses.GlobalConcurrencyLimitResponse]:
+    async with db.session_context() as session:
+        concurrency_limits = (
+            await models.concurrency_limits_v2.read_all_concurrency_limits(
+                session=session,
+                limit=limit,
+                offset=offset,
+            )
+        )
+
+    return [
+        schemas.responses.GlobalConcurrencyLimitResponse.model_validate(limit)
+        for limit in concurrency_limits
+    ]
+
+
+@router.patch("/{id_or_name}", status_code=status.HTTP_204_NO_CONTENT)
+async def update_concurrency_limit_v2(
+    concurrency_limit: actions.ConcurrencyLimitV2Update,
+    id_or_name: Union[UUID, str] = Path(
+        ..., description="The ID or name of the concurrency limit", alias="id_or_name"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    if isinstance(id_or_name, str):  # TODO: this seems like it shouldn't be necessary
+        try:
+            id_or_name = UUID(id_or_name)
+        except ValueError:
+            pass
+    async with db.session_context(begin_transaction=True) as session:
+        if isinstance(id_or_name, UUID):
+            updated = await models.concurrency_limits_v2.update_concurrency_limit(
+                session,
+                concurrency_limit_id=id_or_name,
+                concurrency_limit=concurrency_limit,
+            )
+        else:
+            updated = await models.concurrency_limits_v2.update_concurrency_limit(
+                session, name=id_or_name, concurrency_limit=concurrency_limit
+            )
+
+    if not updated:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency Limit not found"
+        )
+
+
+@router.delete("/{id_or_name}", status_code=status.HTTP_204_NO_CONTENT)
+async def delete_concurrency_limit_v2(
+    id_or_name: Union[UUID, str] = Path(
+        ..., description="The ID or name of the concurrency limit", alias="id_or_name"
+    ),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> None:
+    if isinstance(id_or_name, str):  # TODO: this seems like it shouldn't be necessary
+        try:
+            id_or_name = UUID(id_or_name)
+        except ValueError:
+            pass
+    async with db.session_context(begin_transaction=True) as session:
+        deleted = False
+        if isinstance(id_or_name, UUID):
+            deleted = await models.concurrency_limits_v2.delete_concurrency_limit(
+                session, concurrency_limit_id=id_or_name
+            )
+        else:
+            deleted = await models.concurrency_limits_v2.delete_concurrency_limit(
+                session, name=id_or_name
+            )
+
+    if not deleted:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency Limit not found"
+        )
+
+
+class MinimalConcurrencyLimitResponse(PrefectBaseModel):
+    id: UUID
+    name: str
+    limit: int
+
+
+@router.post("/increment", status_code=status.HTTP_200_OK)
+async def bulk_increment_active_slots(
+    slots: int = Body(..., gt=0),
+    names: List[str] = Body(..., min_items=1),
+    mode: Literal["concurrency", "rate_limit"] = Body("concurrency"),
+    create_if_missing: Optional[bool] = Body(None),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[MinimalConcurrencyLimitResponse]:
+    async with db.session_context(begin_transaction=True) as session:
+        limits = [
+            schemas.core.ConcurrencyLimitV2.model_validate(limit)
+            for limit in (
+                await models.concurrency_limits_v2.bulk_read_or_create_concurrency_limits(
+                    session=session, names=names, create_if_missing=create_if_missing
+                )
+            )
+        ]
+
+        active_limits = [limit for limit in limits if bool(limit.active)]
+
+        if any(limit.limit < slots for limit in active_limits):
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail="Slots requested is greater than the limit",
+            )
+
+        non_decaying = [
+            str(limit.name)
+            for limit in active_limits
+            if limit.slot_decay_per_second == 0.0
+        ]
+
+        if mode == "rate_limit" and non_decaying:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    "Only concurrency limits with slot decay can be used for "
+                    "rate limiting. The following limits do not have a decay "
+                    f"configured: {','.join(non_decaying)!r}"
+                ),
+            )
+        acquired = await models.concurrency_limits_v2.bulk_increment_active_slots(
+            session=session,
+            concurrency_limit_ids=[limit.id for limit in active_limits],
+            slots=slots,
+        )
+
+        if not acquired:
+            await session.rollback()
+
+    if acquired:
+        return [
+            MinimalConcurrencyLimitResponse(
+                id=limit.id, name=str(limit.name), limit=limit.limit
+            )
+            for limit in limits
+        ]
+    else:
+        async with db.session_context(begin_transaction=True) as session:
+            await models.concurrency_limits_v2.bulk_update_denied_slots(
+                session=session,
+                concurrency_limit_ids=[limit.id for limit in active_limits],
+                slots=slots,
+            )
+
+        def num_blocking_slots(limit: schemas.core.ConcurrencyLimitV2) -> float:
+            if limit.slot_decay_per_second > 0.0:
+                return slots + limit.denied_slots
+            else:
+                return (slots + limit.denied_slots) / limit.limit
+
+        blocking_limit = max((limit for limit in active_limits), key=num_blocking_slots)
+        blocking_slots = num_blocking_slots(blocking_limit)
+
+        wait_time_per_slot = (
+            blocking_limit.avg_slot_occupancy_seconds
+            if blocking_limit.slot_decay_per_second == 0.0
+            else (1.0 / blocking_limit.slot_decay_per_second)
+        )
+
+        retry_after = wait_time_per_slot * blocking_slots
+
+        raise HTTPException(
+            status_code=status.HTTP_423_LOCKED,
+            headers={
+                "Retry-After": str(retry_after),
+            },
+        )
+
+
+@router.post("/decrement", status_code=status.HTTP_200_OK)
+async def bulk_decrement_active_slots(
+    slots: int = Body(..., gt=0),
+    names: List[str] = Body(..., min_items=1),
+    occupancy_seconds: Optional[float] = Body(None, gt=0.0),
+    create_if_missing: bool = Body(True),
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[MinimalConcurrencyLimitResponse]:
+    async with db.session_context(begin_transaction=True) as session:
+        limits = (
+            await models.concurrency_limits_v2.bulk_read_or_create_concurrency_limits(
+                session=session, names=names, create_if_missing=create_if_missing
+            )
+        )
+
+        if not limits:
+            return []
+
+        await models.concurrency_limits_v2.bulk_decrement_active_slots(
+            session=session,
+            concurrency_limit_ids=[limit.id for limit in limits if bool(limit.active)],
+            slots=slots,
+            occupancy_seconds=occupancy_seconds,
+        )
+
+    return [
+        MinimalConcurrencyLimitResponse(
+            id=limit.id, name=str(limit.name), limit=limit.limit
+        )
+        for limit in limits
+    ]

prefect/server/api/csrf_token.py

@@ -0,0 +1,38 @@
+from typing import TYPE_CHECKING
+
+from fastapi import Depends, Query, status
+from starlette.exceptions import HTTPException
+
+from prefect.logging import get_logger
+from prefect.server import models, schemas
+from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.utilities.server import PrefectRouter
+from prefect.settings import PREFECT_SERVER_CSRF_PROTECTION_ENABLED
+
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("server.api")
+
+router: PrefectRouter = PrefectRouter(prefix="/csrf-token")
+
+
+@router.get("")
+async def create_csrf_token(
+    db: PrefectDBInterface = Depends(provide_database_interface),
+    client: str = Query(..., description="The client to create a CSRF token for"),
+) -> schemas.core.CsrfToken:
+    """Create or update a CSRF token for a client"""
+    if PREFECT_SERVER_CSRF_PROTECTION_ENABLED.value() is False:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="CSRF protection is disabled.",
+        )
+
+    async with db.session_context(begin_transaction=True) as session:
+        token = await models.csrf_token.create_or_update_csrf_token(
+            session=session, client=client
+        )
+        await models.csrf_token.delete_expired_tokens(session=session)
+
+    return token

prefect/server/api/dependencies.py

@@ -0,0 +1,196 @@
+"""
+Utilities for injecting FastAPI dependencies.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from base64 import b64decode
+from typing import Annotated, Any, Optional
+from uuid import UUID
+
+from fastapi import Body, Depends, Header, HTTPException, status
+from packaging.version import Version
+from starlette.requests import Request
+
+from prefect.server import schemas
+from prefect.settings import PREFECT_API_DEFAULT_LIMIT
+
+
+def provide_request_api_version(
+    x_prefect_api_version: str = Header(None),
+) -> Version | None:
+    if not x_prefect_api_version:
+        return
+
+    # parse version
+    try:
+        _, _, _ = [int(v) for v in x_prefect_api_version.split(".")]
+    except ValueError:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=(
+                "Invalid X-PREFECT-API-VERSION header format.Expected header in format"
+                f" 'x.y.z' but received {x_prefect_api_version}"
+            ),
+        )
+    return Version(x_prefect_api_version)
+
+
+class EnforceMinimumAPIVersion:
+    """
+    FastAPI Dependency used to check compatibility between the version of the api
+    and a given request.
+
+    Looks for the header 'X-PREFECT-API-VERSION' in the request and compares it
+    to the api's version. Rejects requests that are lower than the minimum version.
+    """
+
+    def __init__(self, minimum_api_version: str, logger: logging.Logger):
+        self.minimum_api_version = minimum_api_version
+        versions = [int(v) for v in minimum_api_version.split(".")]
+        self.api_major: int = versions[0]
+        self.api_minor: int = versions[1]
+        self.api_patch: int = versions[2]
+        self.logger = logger
+
+    async def __call__(
+        self,
+        x_prefect_api_version: str = Header(None),
+    ) -> None:
+        request_version = x_prefect_api_version
+
+        # if no version header, assume latest and continue
+        if not request_version:
+            return
+
+        # parse version
+        try:
+            major, minor, patch = [int(v) for v in request_version.split(".")]
+        except ValueError:
+            await self._notify_of_invalid_value(request_version)
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=(
+                    "Invalid X-PREFECT-API-VERSION header format."
+                    f"Expected header in format 'x.y.z' but received {request_version}"
+                ),
+            )
+
+        if (major, minor, patch) < (self.api_major, self.api_minor, self.api_patch):
+            await self._notify_of_outdated_version(request_version)
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=(
+                    f"The request specified API version {request_version} but this "
+                    f"server requires version {self.minimum_api_version} or higher."
+                ),
+            )
+
+    async def _notify_of_invalid_value(self, request_version: str):
+        self.logger.error(
+            f"Invalid X-PREFECT-API-VERSION header format: '{request_version}'"
+        )
+
+    async def _notify_of_outdated_version(self, request_version: str):
+        self.logger.error(
+            f"X-PREFECT-API-VERSION header specifies version '{request_version}' "
+            f"but minimum allowed version is '{self.minimum_api_version}'"
+        )
+
+
+def LimitBody() -> Any:
+    """
+    A `fastapi.Depends` factory for pulling a `limit: int` parameter from the
+    request body while determining the default from the current settings.
+    """
+
+    def get_limit(
+        limit: int = Body(
+            None,
+            description="Defaults to PREFECT_API_DEFAULT_LIMIT if not provided.",
+        ),
+    ):
+        default_limit = PREFECT_API_DEFAULT_LIMIT.value()
+        limit = limit if limit is not None else default_limit
+        if not limit >= 0:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail="Invalid limit: must be greater than or equal to 0.",
+            )
+        if limit > default_limit:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=f"Invalid limit: must be less than or equal to {default_limit}.",
+            )
+        return limit
+
+    return Depends(get_limit)
+
+
+def get_created_by(
+    prefect_automation_id: Optional[UUID] = Header(None, include_in_schema=False),
+    prefect_automation_name: Optional[str] = Header(None, include_in_schema=False),
+) -> Optional[schemas.core.CreatedBy]:
+    """A dependency that returns the provenance information to use when creating objects
+    during this API call."""
+    if prefect_automation_id and prefect_automation_name:
+        try:
+            display_value = b64decode(prefect_automation_name.encode()).decode()
+        except Exception:
+            display_value = None
+
+        if display_value:
+            return schemas.core.CreatedBy(
+                id=prefect_automation_id,
+                type="AUTOMATION",
+                display_value=display_value,
+            )
+
+    return None
+
+
+def get_updated_by(
+    prefect_automation_id: Optional[UUID] = Header(None, include_in_schema=False),
+    prefect_automation_name: Optional[str] = Header(None, include_in_schema=False),
+) -> Optional[schemas.core.UpdatedBy]:
+    """A dependency that returns the provenance information to use when updating objects
+    during this API call."""
+    if prefect_automation_id and prefect_automation_name:
+        return schemas.core.UpdatedBy(
+            id=prefect_automation_id,
+            type="AUTOMATION",
+            display_value=prefect_automation_name,
+        )
+
+    return None
+
+
+def is_ephemeral_request(request: Request) -> bool:
+    """
+    A dependency that returns whether the request is to an ephemeral server.
+    """
+    return "ephemeral-prefect" in str(request.base_url)
+
+
+PREFECT_CLIENT_USER_AGENT_PATTERN = re.compile(
+    r"^prefect/(\d+\.\d+\.\d+(?:[a-z.+0-9]+)?) \(API \S+\)$"
+)
+
+
+def get_prefect_client_version(
+    user_agent: Annotated[Optional[str], Header(include_in_schema=False)] = None,
+) -> Optional[str]:
+    """
+    Attempts to parse out the Prefect client version from the User-Agent header.
+
+    The Prefect client sets the User-Agent header like so:
+      f"prefect/{prefect.__version__} (API {constants.SERVER_API_VERSION})"
+    """
+    if not user_agent:
+        return None
+
+    if client_version := PREFECT_CLIENT_USER_AGENT_PATTERN.match(user_agent):
+        return client_version.group(1)
+    return None