prefect-client 3.1.9__py3-none-any.whl → 3.1.11__py3-none-any.whl

prefect/client/orchestration/_logs/client.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
+
+from prefect.client.orchestration.base import BaseAsyncClient, BaseClient
+from prefect.client.schemas.sorting import (
+    LogSort,
+)
+
+if TYPE_CHECKING:
+    from prefect.client.schemas.actions import (
+        LogCreate,
+    )
+    from prefect.client.schemas.filters import (
+        LogFilter,
+    )
+    from prefect.client.schemas.objects import (
+        Log,
+    )
+
+
+class LogClient(BaseClient):
+    def create_logs(self, logs: Iterable[Union["LogCreate", dict[str, Any]]]) -> None:
+        """
+        Create logs for a flow or task run
+        """
+        from prefect.client.schemas.actions import LogCreate
+
+        serialized_logs = [
+            log.model_dump(mode="json") if isinstance(log, LogCreate) else log
+            for log in logs
+        ]
+        self.request("POST", "/logs/", json=serialized_logs)
+
+    def read_logs(
+        self,
+        log_filter: Optional["LogFilter"] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        sort: "LogSort" = LogSort.TIMESTAMP_ASC,
+    ) -> list["Log"]:
+        """
+        Read flow and task run logs.
+        """
+        body: dict[str, Any] = {
+            "logs": log_filter.model_dump(mode="json") if log_filter else None,
+            "limit": limit,
+            "offset": offset,
+            "sort": sort,
+        }
+        response = self.request("POST", "/logs/filter", json=body)
+        from prefect.client.schemas.objects import Log
+
+        return Log.model_validate_list(response.json())
+
+
+class LogAsyncClient(BaseAsyncClient):
+    async def create_logs(
+        self, logs: Iterable[Union["LogCreate", dict[str, Any]]]
+    ) -> None:
+        """
+        Create logs for a flow or task run
+
+        Args:
+            logs: An iterable of `LogCreate` objects or already json-compatible dicts
+        """
+        from prefect.client.schemas.actions import LogCreate
+
+        serialized_logs = [
+            log.model_dump(mode="json") if isinstance(log, LogCreate) else log
+            for log in logs
+        ]
+        await self.request("POST", "/logs/", json=serialized_logs)
+
+    async def read_logs(
+        self,
+        log_filter: Optional["LogFilter"] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        sort: "LogSort" = LogSort.TIMESTAMP_ASC,
+    ) -> list[Log]:
+        """
+        Read flow and task run logs.
+        """
+        body: dict[str, Any] = {
+            "logs": log_filter.model_dump(mode="json") if log_filter else None,
+            "limit": limit,
+            "offset": offset,
+            "sort": sort,
+        }
+
+        response = await self.request("POST", "/logs/filter", json=body)
+        from prefect.client.schemas.objects import Log
+
+        return Log.model_validate_list(response.json())

prefect/client/orchestration/_variables/client.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import httpx
+
+from prefect.client.orchestration.base import BaseAsyncClient, BaseClient
+from prefect.exceptions import ObjectNotFound
+
+if TYPE_CHECKING:
+    from prefect.client.schemas.actions import (
+        VariableCreate,
+        VariableUpdate,
+    )
+    from prefect.client.schemas.objects import (
+        Variable,
+    )
+
+
+class VariableClient(BaseClient):
+    def create_variable(self, variable: "VariableCreate") -> "Variable":
+        """
+        Creates an variable with the provided configuration.
+
+        Args:
+            variable: Desired configuration for the new variable.
+        Returns:
+            Information about the newly created variable.
+        """
+        response = self._client.post(
+            "/variables/",
+            json=variable.model_dump(mode="json", exclude_unset=True),
+        )
+        from prefect.client.schemas.objects import Variable
+
+        return Variable.model_validate(response.json())
+
+    def read_variable_by_name(self, name: str) -> "Variable | None":
+        """Reads a variable by name. Returns None if no variable is found."""
+        try:
+            response = self.request(
+                "GET", "/variables/name/{name}", path_params={"name": name}
+            )
+            from prefect.client.schemas.objects import Variable
+
+            return Variable(**response.json())
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return None
+            else:
+                raise
+
+    def read_variables(self, limit: int | None = None) -> list["Variable"]:
+        """Reads all variables."""
+        response = self.request("POST", "/variables/filter", json={"limit": limit})
+        from prefect.client.schemas.objects import Variable
+
+        return Variable.model_validate_list(response.json())
+
+    def update_variable(self, variable: "VariableUpdate") -> None:
+        """
+        Updates a variable with the provided configuration.
+
+        Args:
+            variable: Desired configuration for the updated variable.
+        Returns:
+            Information about the updated variable.
+        """
+        self._client.patch(
+            f"/variables/name/{variable.name}",
+            json=variable.model_dump(mode="json", exclude_unset=True),
+        )
+        return None
+
+    def delete_variable_by_name(self, name: str) -> None:
+        """Deletes a variable by name."""
+        try:
+            self.request(
+                "DELETE",
+                "/variables/name/{name}",
+                path_params={"name": name},
+            )
+            return None
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+
+class VariableAsyncClient(BaseAsyncClient):
+    async def create_variable(self, variable: "VariableCreate") -> "Variable":
+        """Creates a variable with the provided configuration."""
+        response = await self._client.post(
+            "/variables/",
+            json=variable.model_dump(mode="json", exclude_unset=True),
+        )
+        from prefect.client.schemas.objects import Variable
+
+        return Variable.model_validate(response.json())
+
+    async def read_variable_by_name(self, name: str) -> "Variable | None":
+        """Reads a variable by name. Returns None if no variable is found."""
+        try:
+            response = await self.request(
+                "GET",
+                "/variables/name/{name}",
+                path_params={"name": name},
+            )
+            from prefect.client.schemas.objects import Variable
+
+            return Variable.model_validate(response.json())
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                return None
+            else:
+                raise
+
+    async def read_variables(self, limit: int | None = None) -> list["Variable"]:
+        """Reads all variables."""
+        response = await self.request(
+            "POST", "/variables/filter", json={"limit": limit}
+        )
+        from prefect.client.schemas.objects import Variable
+
+        return Variable.model_validate_list(response.json())
+
+    async def update_variable(self, variable: "VariableUpdate") -> None:
+        """
+        Updates a variable with the provided configuration.
+
+        Args:
+            variable: Desired configuration for the updated variable.
+        Returns:
+            Information about the updated variable.
+        """
+        await self.request(
+            "PATCH",
+            "/variables/name/{name}",
+            path_params={"name": variable.name},
+            json=variable.model_dump(mode="json", exclude_unset=True),
+        )
+        return None
+
+    async def delete_variable_by_name(self, name: str) -> None:
+        """Deletes a variable by name."""
+        try:
+            await self.request(
+                "DELETE",
+                "/variables/name/{name}",
+                path_params={"name": name},
+            )
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise

prefect/client/orchestration/base.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    from httpx import AsyncClient, Client, Response
+
+    from prefect.client.orchestration.routes import ServerRoutes
+
+HTTP_METHODS: TypeAlias = Literal["GET", "POST", "PUT", "DELETE", "PATCH"]
+
+
+class BaseClient:
+    def __init__(self, client: "Client"):
+        self._client = client
+
+    def request(
+        self,
+        method: HTTP_METHODS,
+        path: "ServerRoutes",
+        params: dict[str, Any] | None = None,
+        path_params: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> "Response":
+        if path_params:
+            path = path.format(**path_params)  # type: ignore
+        return self._client.request(method, path, params=params, **kwargs)
+
+
+class BaseAsyncClient:
+    def __init__(self, client: "AsyncClient"):
+        self._client = client
+
+    async def request(
+        self,
+        method: HTTP_METHODS,
+        path: "ServerRoutes",
+        params: dict[str, Any] | None = None,
+        path_params: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> "Response":
+        if path_params:
+            path = path.format(**path_params)  # type: ignore
+        return await self._client.request(method, path, params=params, **kwargs)

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/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:
@@ -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,14 +343,11 @@ 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():
+        if isinstance(self.data, ResultRecord) and should_persist_result():
             data = self.data.metadata
         else:
             data = None

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()
+    ]