prefect-client 3.1.10__py3-none-any.whl → 3.1.12__py3-none-any.whl

prefect/client/orchestration/routes.py

@@ -0,0 +1,145 @@
+from typing import Literal
+
+ServerRoutes = Literal[
+    "/admin/database/clear",
+    "/admin/database/create",
+    "/admin/database/drop",
+    "/admin/settings",
+    "/admin/version",
+    "/artifacts/",
+    "/artifacts/{id}",
+    "/artifacts/{key}/latest",
+    "/artifacts/count",
+    "/artifacts/filter",
+    "/artifacts/latest/count",
+    "/artifacts/latest/filter",
+    "/automations/",
+    "/automations/{id}",
+    "/automations/count",
+    "/automations/filter",
+    "/automations/owned-by/{resource_id}",
+    "/automations/related-to/{resource_id}",
+    "/block_capabilities/",
+    "/block_documents/",
+    "/block_documents/{id}",
+    "/block_documents/count",
+    "/block_documents/filter",
+    "/block_schemas/",
+    "/block_schemas/{id}",
+    "/block_schemas/checksum/{checksum}",
+    "/block_schemas/filter",
+    "/block_types/",
+    "/block_types/{id}",
+    "/block_types/filter",
+    "/block_types/install_system_block_types",
+    "/block_types/slug/{slug}",
+    "/block_types/slug/{slug}/block_documents",
+    "/block_types/slug/{slug}/block_documents/name/{block_document_name}",
+    "/collections/views/{view}",
+    "/concurrency_limits/",
+    "/concurrency_limits/{id}",
+    "/concurrency_limits/decrement",
+    "/concurrency_limits/filter",
+    "/concurrency_limits/increment",
+    "/concurrency_limits/tag/{tag}",
+    "/concurrency_limits/tag/{tag}/reset",
+    "/csrf-token",
+    "/deployments/",
+    "/deployments/{id}",
+    "/deployments/{id}/create_flow_run",
+    "/deployments/{id}/pause_deployment",
+    "/deployments/{id}/resume_deployment",
+    "/deployments/{id}/schedule",
+    "/deployments/{id}/schedules",
+    "/deployments/{id}/schedules/{schedule_id}",
+    "/deployments/{id}/work_queue_check",
+    "/deployments/count",
+    "/deployments/filter",
+    "/deployments/get_scheduled_flow_runs",
+    "/deployments/name/{flow_name}/{deployment_name}",
+    "/deployments/paginate",
+    "/events",
+    "/events/count-by/{countable}",
+    "/events/filter",
+    "/events/filter/next",
+    "/flow_run_notification_policies/",
+    "/flow_run_notification_policies/{id}",
+    "/flow_run_notification_policies/filter",
+    "/flow_run_states/",
+    "/flow_run_states/{id}",
+    "/flow_runs/",
+    "/flow_runs/{id}",
+    "/flow_runs/{id}/graph",
+    "/flow_runs/{id}/graph-v2",
+    "/flow_runs/{id}/input",
+    "/flow_runs/{id}/input/{key}",
+    "/flow_runs/{id}/input/filter",
+    "/flow_runs/{id}/labels",
+    "/flow_runs/{id}/logs/download",
+    "/flow_runs/{id}/resume",
+    "/flow_runs/{id}/set_state",
+    "/flow_runs/count",
+    "/flow_runs/filter",
+    "/flow_runs/history",
+    "/flow_runs/lateness",
+    "/flow_runs/paginate",
+    "/flows/",
+    "/flows/{id}",
+    "/flows/count",
+    "/flows/filter",
+    "/flows/name/{name}",
+    "/flows/paginate",
+    "/health",
+    "/hello",
+    "/logs/",
+    "/logs/filter",
+    "/ready",
+    "/saved_searches/",
+    "/saved_searches/{id}",
+    "/saved_searches/filter",
+    "/task_run_states/",
+    "/task_run_states/{id}",
+    "/task_runs/",
+    "/task_runs/{id}",
+    "/task_runs/{id}/set_state",
+    "/task_runs/count",
+    "/task_runs/filter",
+    "/task_runs/history",
+    "/task_workers/filter",
+    "/templates/validate",
+    "/ui/flow_runs/count-task-runs",
+    "/ui/flow_runs/history",
+    "/ui/flows/count-deployments",
+    "/ui/flows/next-runs",
+    "/ui/schemas/validate",
+    "/ui/task_runs/count",
+    "/ui/task_runs/dashboard/counts",
+    "/v2/concurrency_limits/",
+    "/v2/concurrency_limits/{id_or_name}",
+    "/v2/concurrency_limits/decrement",
+    "/v2/concurrency_limits/filter",
+    "/v2/concurrency_limits/increment",
+    "/variables/",
+    "/variables/{id}",
+    "/variables/count",
+    "/variables/filter",
+    "/variables/name/{name}",
+    "/version",
+    "/work_pools/",
+    "/work_pools/{name}",
+    "/work_pools/{name}/get_scheduled_flow_runs",
+    "/work_pools/{work_pool_name}/queues",
+    "/work_pools/{work_pool_name}/queues/{name}",
+    "/work_pools/{work_pool_name}/queues/filter",
+    "/work_pools/{work_pool_name}/workers/{name}",
+    "/work_pools/{work_pool_name}/workers/filter",
+    "/work_pools/{work_pool_name}/workers/heartbeat",
+    "/work_pools/count",
+    "/work_pools/filter",
+    "/work_queues/",
+    "/work_queues/{id}",
+    "/work_queues/{id}/get_runs",
+    "/work_queues/{id}/status",
+    "/work_queues/filter",
+    "/work_queues/name/{name}",
+]

prefect/client/schemas/__init__.py

@@ -1,32 +1,58 @@
-# Some objects are exported here for backwards compatibility.
-# In general, it is recommended to import schemas from their respective modules.
+import importlib
+import sys
+from typing import Any, TYPE_CHECKING
 
-from .actions import BlockTypeUpdate, StateCreate
-from .objects import (
-    DEFAULT_BLOCK_SCHEMA_VERSION,
-    BlockDocument,
-    BlockSchema,
-    BlockType,
-    FlowRun,
-    FlowRunPolicy,
-    State,
-    StateDetails,
-    StateType,
-    TaskRun,
-    TaskRunInput,
-    TaskRunPolicy,
-    TaskRunResult,
-    Workspace,
-)
-from .responses import (
-    OrchestrationResult,
-    SetStateStatus,
-    StateAbortDetails,
-    StateAcceptDetails,
-    StateRejectDetails,
-)
+if TYPE_CHECKING:
+    from .actions import BlockTypeUpdate, StateCreate
+    from .objects import (
+        DEFAULT_BLOCK_SCHEMA_VERSION,
+        BlockDocument,
+        BlockSchema,
+        BlockType,
+        FlowRun,
+        FlowRunPolicy,
+        State,
+        StateDetails,
+        StateType,
+        TaskRun,
+        TaskRunInput,
+        TaskRunPolicy,
+        TaskRunResult,
+        Workspace,
+    )
+    from .responses import (
+        OrchestrationResult,
+        SetStateStatus,
+        StateAbortDetails,
+        StateAcceptDetails,
+        StateRejectDetails,
+    )
 
-__all__ = (
+_public_api = {
+    "BlockDocument": (__package__, ".objects"),
+    "BlockSchema": (__package__, ".objects"),
+    "BlockType": (__package__, ".objects"),
+    "BlockTypeUpdate": (__package__, ".actions"),
+    "DEFAULT_BLOCK_SCHEMA_VERSION": (__package__, ".objects"),
+    "FlowRun": (__package__, ".objects"),
+    "FlowRunPolicy": (__package__, ".objects"),
+    "OrchestrationResult": (__package__, ".responses"),
+    "SetStateStatus": (__package__, ".responses"),
+    "State": (__package__, ".objects"),
+    "StateAbortDetails": (__package__, ".responses"),
+    "StateAcceptDetails": (__package__, ".responses"),
+    "StateCreate": (__package__, ".actions"),
+    "StateDetails": (__package__, ".objects"),
+    "StateRejectDetails": (__package__, ".responses"),
+    "StateType": (__package__, ".objects"),
+    "TaskRun": (__package__, ".objects"),
+    "TaskRunInput": (__package__, ".objects"),
+    "TaskRunPolicy": (__package__, ".objects"),
+    "TaskRunResult": (__package__, ".objects"),
+    "Workspace": (__package__, ".objects"),
+}
+
+__all__ = [
     "BlockDocument",
     "BlockSchema",
     "BlockType",
@@ -48,4 +74,18 @@ __all__ = (
     "TaskRunPolicy",
     "TaskRunResult",
     "Workspace",
-)
+]
+
+
+def __getattr__(attr_name: str) -> Any:
+    try:
+        if (dynamic_attr := _public_api.get(attr_name)) is None:
+            raise AttributeError(f"module {__name__} has no attribute {attr_name}")
+
+        package, mname = dynamic_attr
+        module = importlib.import_module(mname, package=package)
+        return getattr(module, attr_name)
+    except ModuleNotFoundError as ex:
+        mname, _, attr = (ex.name or "").rpartition(".")
+        ctx = {"name": mname, "obj": attr} if sys.version_info >= (3, 10) else {}
+        raise AttributeError(f"module {mname} has no attribute {attr}", **ctx) from ex

prefect/client/schemas/actions.py

@@ -39,7 +39,7 @@ from prefect.utilities.collections import listrepr
 from prefect.utilities.pydantic import get_class_fields_only
 
 if TYPE_CHECKING:
-    from prefect.results import BaseResult, ResultRecordMetadata
+    from prefect.results import ResultRecordMetadata
 
 R = TypeVar("R")
 
@@ -51,7 +51,7 @@ class StateCreate(ActionBaseModel):
     name: Optional[str] = Field(default=None)
     message: Optional[str] = Field(default=None, examples=["Run started"])
     state_details: StateDetails = Field(default_factory=StateDetails)
-    data: Union["BaseResult[Any]", "ResultRecordMetadata", Any] = Field(
+    data: Union["ResultRecordMetadata", Any] = Field(
         default=None,
     )
 

prefect/client/schemas/filters.py

@@ -505,6 +505,11 @@ class DeploymentFilterTags(PrefectBaseModel, OperatorMixin):
             " superset of the list"
         ),
     )
+    any_: Optional[list[str]] = Field(
+        default=None,
+        examples=[["tag-1", "tag-2"]],
+        description="A list of tags to include",
+    )
     is_null_: Optional[bool] = Field(
         default=None, description="If true, only include deployments without tags"
     )

prefect/client/schemas/objects.py

@@ -10,7 +10,6 @@ from typing import (
     Generic,
     Optional,
     Union,
-    cast,
     overload,
 )
 from uuid import UUID, uuid4
@@ -65,7 +64,7 @@ from prefect.utilities.pydantic import handle_secret_render
 
 if TYPE_CHECKING:
     from prefect.client.schemas.actions import StateCreate
-    from prefect.results import BaseResult, ResultRecordMetadata
+    from prefect.results import ResultRecordMetadata
 
     DateTime = pendulum.DateTime
 else:
@@ -122,7 +121,7 @@ class WorkPoolStatus(AutoEnum):
     PAUSED = AutoEnum.auto()
 
     @property
-    def display_name(self):
+    def display_name(self) -> str:
         return self.name.replace("_", " ").capitalize()
 
 
@@ -195,9 +194,7 @@ class StateDetails(PrefectBaseModel):
 
 
 def data_discriminator(x: Any) -> str:
-    if isinstance(x, dict) and "type" in x and x["type"] != "unpersisted":
-        return "BaseResult"
-    elif isinstance(x, dict) and "storage_key" in x:
+    if isinstance(x, dict) and "storage_key" in x:
         return "ResultRecordMetadata"
     return "Any"
 
@@ -214,7 +211,6 @@ class State(ObjectBaseModel, Generic[R]):
     state_details: StateDetails = Field(default_factory=StateDetails)
     data: Annotated[
         Union[
-            Annotated["BaseResult[R]", Tag("BaseResult")],
             Annotated["ResultRecordMetadata", Tag("ResultRecordMetadata")],
             Annotated[Any, Tag("Any")],
         ],
@@ -347,15 +343,12 @@ class State(ObjectBaseModel, Generic[R]):
         """
         from prefect.client.schemas.actions import StateCreate
         from prefect.results import (
-            BaseResult,
             ResultRecord,
             should_persist_result,
         )
 
-        if isinstance(self.data, BaseResult):
-            data = cast(BaseResult[R], self.data)
-        elif isinstance(self.data, ResultRecord) and should_persist_result():
-            data = self.data.metadata
+        if isinstance(self.data, ResultRecord) and should_persist_result():
+            data = self.data.metadata  # pyright: ignore[reportUnknownMemberType] unable to narrow ResultRecord type
         else:
             data = None
 
@@ -386,7 +379,7 @@ class State(ObjectBaseModel, Generic[R]):
 
     @model_validator(mode="after")
     def set_unpersisted_results_to_none(self) -> Self:
-        if isinstance(self.data, dict) and self.data.get("type") == "unpersisted":
+        if isinstance(self.data, dict) and self.data.get("type") == "unpersisted":  # pyright: ignore[reportUnknownMemberType] unable to narrow dict type
             self.data = None
         return self
 
@@ -531,7 +524,7 @@ class FlowRunPolicy(PrefectBaseModel):
     @classmethod
     def populate_deprecated_fields(cls, values: Any) -> Any:
         if isinstance(values, dict):
-            return set_run_policy_deprecated_fields(values)
+            return set_run_policy_deprecated_fields(values)  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] unable to narrow dict type
         return values
 
 
@@ -1262,7 +1255,7 @@ class BlockDocumentReference(ObjectBaseModel):
     @classmethod
     def validate_parent_and_ref_are_different(cls, values: Any) -> Any:
         if isinstance(values, dict):
-            return validate_parent_and_ref_diff(values)
+            return validate_parent_and_ref_diff(values)  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] unable to narrow dict type
         return values
 
 

prefect/client/schemas/schedules.py

@@ -33,6 +33,23 @@ MAX_ITERATIONS = 1000
 MAX_RRULE_LENGTH = 6500
 
 
+def is_valid_timezone(v: str) -> bool:
+    """
+    Validate that the provided timezone is a valid IANA timezone.
+
+    Unfortunately this list is slightly different from the list of valid
+    timezones in pendulum that we use for cron and interval timezone validation.
+    """
+    from prefect._internal.pytz import HAS_PYTZ
+
+    if HAS_PYTZ:
+        import pytz
+    else:
+        from prefect._internal import pytz
+
+    return v in pytz.all_timezones_set
+
+
 class IntervalSchedule(PrefectBaseModel):
     """
     A schedule formed by adding `interval` increments to an `anchor_date`. If no
@@ -305,18 +322,13 @@ class RRuleSchedule(PrefectBaseModel):
         Unfortunately this list is slightly different from the list of valid
         timezones in pendulum that we use for cron and interval timezone validation.
         """
-        from prefect._internal.pytz import HAS_PYTZ
+        if v is None:
+            return "UTC"
 
-        if HAS_PYTZ:
-            import pytz
-        else:
-            from prefect._internal import pytz
+        if is_valid_timezone(v):
+            return v
 
-        if v and v not in pytz.all_timezones_set:
-            raise ValueError(f'Invalid timezone: "{v}"')
-        elif v is None:
-            return "UTC"
-        return v
+        raise ValueError(f'Invalid timezone: "{v}"')
 
 
 class NoSchedule(PrefectBaseModel):

prefect/concurrency/_asyncio.py

@@ -0,0 +1,87 @@
+import asyncio
+from typing import Literal, Optional
+
+import httpx
+
+from prefect._internal.compatibility.deprecated import deprecated_parameter
+from prefect.client.orchestration import get_client
+from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
+from prefect.logging.loggers import get_run_logger
+
+from .services import ConcurrencySlotAcquisitionService
+
+
+class ConcurrencySlotAcquisitionError(Exception):
+    """Raised when an unhandlable occurs while acquiring concurrency slots."""
+
+
+class AcquireConcurrencySlotTimeoutError(TimeoutError):
+    """Raised when acquiring a concurrency slot times out."""
+
+
+@deprecated_parameter(
+    name="create_if_missing",
+    start_date="Sep 2024",
+    end_date="Oct 2024",
+    when=lambda x: x is not None,
+    help="Limits must be explicitly created before acquiring concurrency slots; see `strict` if you want to enforce this behavior.",
+)
+async def aacquire_concurrency_slots(
+    names: list[str],
+    slots: int,
+    mode: Literal["concurrency", "rate_limit"] = "concurrency",
+    timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = None,
+    max_retries: Optional[int] = None,
+    strict: bool = False,
+) -> list[MinimalConcurrencyLimitResponse]:
+    service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
+    future = service.send(
+        (slots, mode, timeout_seconds, create_if_missing, max_retries)
+    )
+    try:
+        response = await asyncio.wrap_future(future)
+    except TimeoutError as timeout:
+        raise AcquireConcurrencySlotTimeoutError(
+            f"Attempt to acquire concurrency slots timed out after {timeout_seconds} second(s)"
+        ) from timeout
+    except Exception as exc:
+        raise ConcurrencySlotAcquisitionError(
+            f"Unable to acquire concurrency slots on {names!r}"
+        ) from exc
+
+    retval = _response_to_minimal_concurrency_limit_response(response)
+
+    if not retval:
+        if strict:
+            raise ConcurrencySlotAcquisitionError(
+                f"Concurrency limits {names!r} must be created before acquiring slots"
+            )
+        try:
+            logger = get_run_logger()
+        except Exception:
+            pass
+        else:
+            logger.warning(
+                f"Concurrency limits {names!r} do not exist - skipping acquisition."
+            )
+
+    return retval
+
+
+async def arelease_concurrency_slots(
+    names: list[str], slots: int, occupancy_seconds: float
+) -> list[MinimalConcurrencyLimitResponse]:
+    async with get_client() as client:
+        response = await client.release_concurrency_slots(
+            names=names, slots=slots, occupancy_seconds=occupancy_seconds
+        )
+        return _response_to_minimal_concurrency_limit_response(response)
+
+
+def _response_to_minimal_concurrency_limit_response(
+    response: httpx.Response,
+) -> list[MinimalConcurrencyLimitResponse]:
+    return [
+        MinimalConcurrencyLimitResponse.model_validate(obj_) for obj_ in response.json()
+    ]

prefect/concurrency/{events.py → _events.py}

@@ -1,4 +1,4 @@
-from typing import Dict, List, Literal, Optional, Union
+from typing import Literal, Optional, Union
 from uuid import UUID
 
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
@@ -8,11 +8,11 @@ from prefect.events import Event, RelatedResource, emit_event
 def _emit_concurrency_event(
     phase: Union[Literal["acquired"], Literal["released"]],
     primary_limit: MinimalConcurrencyLimitResponse,
-    related_limits: List[MinimalConcurrencyLimitResponse],
+    related_limits: list[MinimalConcurrencyLimitResponse],
     slots: int,
     follows: Union[Event, None] = None,
 ) -> Union[Event, None]:
-    resource: Dict[str, str] = {
+    resource: dict[str, str] = {
         "prefect.resource.id": f"prefect.concurrency-limit.{primary_limit.id}",
         "prefect.resource.name": primary_limit.name,
         "slots-acquired": str(slots),
@@ -38,11 +38,11 @@ def _emit_concurrency_event(
     )
 
 
-def _emit_concurrency_acquisition_events(
-    limits: List[MinimalConcurrencyLimitResponse],
+def emit_concurrency_acquisition_events(
+    limits: list[MinimalConcurrencyLimitResponse],
     occupy: int,
-) -> Dict[UUID, Optional[Event]]:
-    events = {}
+) -> dict[UUID, Optional[Event]]:
+    events: dict[UUID, Optional[Event]] = {}
     for limit in limits:
         event = _emit_concurrency_event("acquired", limit, limits, occupy)
         events[limit.id] = event
@@ -50,10 +50,10 @@ def _emit_concurrency_acquisition_events(
     return events
 
 
-def _emit_concurrency_release_events(
-    limits: List[MinimalConcurrencyLimitResponse],
+def emit_concurrency_release_events(
+    limits: list[MinimalConcurrencyLimitResponse],
     occupy: int,
-    events: Dict[UUID, Optional[Event]],
+    events: dict[UUID, Optional[Event]],
 ) -> None:
     for limit in limits:
         _emit_concurrency_event("released", limit, limits, occupy, events[limit.id])