prefect-client 2.18.0__py3-none-any.whl → 2.18.1__py3-none-any.whl

prefect/_internal/schemas/fields.py

@@ -3,6 +3,7 @@ from typing import Optional
 from uuid import UUID
 
 import pendulum
+from typing_extensions import TypeAlias
 
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
 
@@ -12,19 +13,37 @@ else:
     from pydantic import BaseModel, Field
 
 
-class DateTimeTZ(pendulum.DateTime):
-    @classmethod
-    def __get_validators__(cls):
-        yield cls.validate
+# Rather than subclassing pendulum.DateTime to add our pydantic-specific validation,
+# which will lead to a lot of funky typing issues, we'll just monkeypatch the pydantic
+# validators onto the class.  Retaining this type alias means that we can still use it
+# as we have been in class definitions, also guaranteeing that we'll be applying these
+# validators by importing this module.
 
-    @classmethod
-    def validate(cls, v):
-        if isinstance(v, str):
-            return pendulum.parse(v)
-        elif isinstance(v, datetime.datetime):
-            return pendulum.instance(v)
-        else:
-            raise ValueError("Unrecognized datetime.")
+DateTimeTZ: TypeAlias = pendulum.DateTime
+
+
+def _datetime_patched_classmethod(function):
+    if hasattr(DateTimeTZ, function.__name__):
+        return function
+    setattr(DateTimeTZ, function.__name__, classmethod(function))
+    return function
+
+
+@_datetime_patched_classmethod
+def __get_validators__(cls):
+    yield getattr(cls, "validate")
+
+
+@_datetime_patched_classmethod
+def validate(cls, v) -> pendulum.DateTime:
+    if isinstance(v, str):
+        parsed = pendulum.parse(v)
+        assert isinstance(parsed, pendulum.DateTime)
+        return parsed
+    elif isinstance(v, datetime.datetime):
+        return pendulum.instance(v)
+    else:
+        raise ValueError("Unrecognized datetime.")
 
 
 class CreatedBy(BaseModel):

prefect/blocks/core.py

@@ -996,7 +996,7 @@ class Block(BaseModel, ABC):
         return self._block_document_id
 
     @sync_compatible
-    @instrument_instance_method_call()
+    @instrument_instance_method_call
     async def save(
         self,
         name: Optional[str] = None,

prefect/blocks/notifications.py

@@ -62,7 +62,7 @@ class AbstractAppriseNotificationBlock(NotificationBlock, ABC):
         self._start_apprise_client(self.url)
 
     @sync_compatible
-    @instrument_instance_method_call()
+    @instrument_instance_method_call
     async def notify(
         self,
         body: str,
@@ -717,7 +717,7 @@ class CustomWebhookNotificationBlock(NotificationBlock):
                     raise KeyError(f"{name}/{placeholder}")
 
     @sync_compatible
-    @instrument_instance_method_call()
+    @instrument_instance_method_call
     async def notify(self, body: str, subject: Optional[str] = None):
         import httpx
 

prefect/blocks/system.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-import pendulum
-
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
 
 if HAS_PYDANTIC_V2:
@@ -9,6 +7,7 @@ if HAS_PYDANTIC_V2:
 else:
     from pydantic import Field, SecretStr
 
+from prefect._internal.schemas.fields import DateTimeTZ
 from prefect.blocks.core import Block
 
 
@@ -76,7 +75,7 @@ class DateTime(Block):
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/8b3da9a6621e92108b8e6a75b82e15374e170ff7-48x48.png"
     _documentation_url = "https://docs.prefect.io/api-ref/prefect/blocks/system/#prefect.blocks.system.DateTime"
 
-    value: pendulum.DateTime = Field(
+    value: DateTimeTZ = Field(
         default=...,
         description="An ISO 8601-compatible datetime value.",
     )

prefect/client/orchestration.py

@@ -12,6 +12,7 @@ from typing import (
     Optional,
     Set,
     Tuple,
+    TypeVar,
     Union,
 )
 from uuid import UUID, uuid4
@@ -20,15 +21,19 @@ import certifi
 import httpcore
 import httpx
 import pendulum
+from typing_extensions import ParamSpec
 
 from prefect._internal.compatibility.deprecated import (
     handle_deprecated_infra_overrides_parameter,
 )
 from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect.client.schemas import sorting
+from prefect.events import filters
 from prefect.settings import (
     PREFECT_API_SERVICES_TRIGGERS_ENABLED,
     PREFECT_EXPERIMENTAL_EVENTS,
 )
+from prefect.utilities.asyncutils import run_sync
 
 if HAS_PYDANTIC_V2:
     import pydantic.v1 as pydantic
@@ -151,6 +156,9 @@ if TYPE_CHECKING:
 
 from prefect.client.base import ASGIApp, PrefectHttpxClient, app_lifespan_context
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
 class ServerType(AutoEnum):
     EPHEMERAL = AutoEnum.auto()
@@ -164,7 +172,9 @@ class ServerType(AutoEnum):
         return PREFECT_EXPERIMENTAL_EVENTS and PREFECT_API_SERVICES_TRIGGERS_ENABLED
 
 
-def get_client(httpx_settings: Optional[dict] = None) -> "PrefectClient":
+def get_client(
+    httpx_settings: Optional[Dict[str, Any]] = None, sync_client: bool = False
+) -> "PrefectClient":
     """
     Retrieve a HTTP client for communicating with the Prefect REST API.
 
@@ -174,6 +184,13 @@ def get_client(httpx_settings: Optional[dict] = None) -> "PrefectClient":
     async with get_client() as client:
         await client.hello()
     ```
+
+    To return a synchronous client, pass sync_client=True:
+
+    ```python
+    with get_client(sync_client=True) as client:
+        client.hello()
+    ```
     """
     ctx = prefect.context.get_settings_context()
     api = PREFECT_API_URL.value()
@@ -184,11 +201,18 @@ def get_client(httpx_settings: Optional[dict] = None) -> "PrefectClient":
 
         api = create_app(ctx.settings, ephemeral=True)
 
-    return PrefectClient(
-        api,
-        api_key=PREFECT_API_KEY.value(),
-        httpx_settings=httpx_settings,
-    )
+    if sync_client:
+        return SyncPrefectClient(
+            api,
+            api_key=PREFECT_API_KEY.value(),
+            httpx_settings=httpx_settings,
+        )
+    else:
+        return PrefectClient(
+            api,
+            api_key=PREFECT_API_KEY.value(),
+            httpx_settings=httpx_settings,
+        )
 
 
 class PrefectClient:
@@ -2136,21 +2160,23 @@ class PrefectClient:
 
     async def create_task_run(
         self,
-        task: "TaskObject",
+        task: "TaskObject[P, R]",
         flow_run_id: Optional[UUID],
         dynamic_key: str,
-        name: str = None,
-        extra_tags: Iterable[str] = None,
-        state: prefect.states.State = None,
-        task_inputs: Dict[
-            str,
-            List[
-                Union[
-                    TaskRunResult,
-                    Parameter,
-                    Constant,
-                ]
-            ],
+        name: Optional[str] = None,
+        extra_tags: Optional[Iterable[str]] = None,
+        state: Optional[prefect.states.State[R]] = None,
+        task_inputs: Optional[
+            Dict[
+                str,
+                List[
+                    Union[
+                        TaskRunResult,
+                        Parameter,
+                        Constant,
+                    ]
+                ],
+            ]
         ] = None,
     ) -> TaskRun:
         """
@@ -3143,9 +3169,7 @@ class PrefectClient:
         response.raise_for_status()
         return pydantic.parse_obj_as(List[Automation], response.json())
 
-    async def find_automation(
-        self, id_or_name: str, exit_if_not_found: bool = True
-    ) -> Optional[Automation]:
+    async def find_automation(self, id_or_name: str) -> Optional[Automation]:
         try:
             id = UUID(id_or_name)
         except ValueError:
@@ -3180,6 +3204,34 @@ class PrefectClient:
         response.raise_for_status()
         return Automation.parse_obj(response.json())
 
+    async def read_automations_by_name(self, name: str) -> List[Automation]:
+        """
+        Query the Prefect API for an automation by name. Only automations matching the provided name will be returned.
+
+        Args:
+            name: the name of the automation to query
+
+        Returns:
+            a list of Automation model representations of the automations
+        """
+        if not self.server_type.supports_automations():
+            self._raise_for_unsupported_automations()
+        automation_filter = filters.AutomationFilter(name=dict(any_=[name]))
+
+        response = await self._client.post(
+            "/automations/filter",
+            json={
+                "sort": sorting.AutomationSort.UPDATED_DESC,
+                "automations": automation_filter.dict(json_compatible=True)
+                if automation_filter
+                else None,
+            },
+        )
+
+        response.raise_for_status()
+
+        return pydantic.parse_obj_as(List[Automation], response.json())
+
     async def pause_automation(self, automation_id: UUID):
         if not self.server_type.supports_automations():
             self._raise_for_unsupported_automations()
@@ -3284,3 +3336,212 @@ class PrefectClient:
 
     def __exit__(self, *_):
         assert False, "This should never be called but must be defined for __enter__"
+
+
+class SyncPrefectClient:
+    """
+    A synchronous client for interacting with the [Prefect REST API](/api-ref/rest-api/).
+
+    Args:
+        api: the REST API URL or FastAPI application to connect to
+        api_key: An optional API key for authentication.
+        api_version: The API version this client is compatible with.
+        httpx_settings: An optional dictionary of settings to pass to the underlying
+            `httpx.AsyncClient`
+
+    Examples:
+
+        Say hello to a Prefect REST API
+
+        <div class="terminal">
+        ```
+        >>> with get_client(sync_client=True) as client:
+        >>>     response = client.hello()
+        >>>
+        >>> print(response.json())
+        👋
+        ```
+        </div>
+    """
+
+    def __init__(
+        self,
+        api: Union[str, ASGIApp],
+        *,
+        api_key: Optional[str] = None,
+        api_version: Optional[str] = None,
+        httpx_settings: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._prefect_client = PrefectClient(
+            api=api,
+            api_key=api_key,
+            api_version=api_version,
+            httpx_settings=httpx_settings,
+        )
+
+    def __enter__(self):
+        run_sync(self._prefect_client.__aenter__())
+        return self
+
+    def __exit__(self, *exc_info):
+        return run_sync(self._prefect_client.__aexit__(*exc_info))
+
+    async def __aenter__(self):
+        raise RuntimeError(
+            "The `SyncPrefectClient` must be entered with a sync context. Use '"
+            "with SyncPrefectClient(...)' not 'async with SyncPrefectClient(...)'"
+        )
+
+    async def __aexit__(self, *_):
+        assert False, "This should never be called but must be defined for __aenter__"
+
+    def hello(self) -> httpx.Response:
+        """
+        Send a GET request to /hello for testing purposes.
+        """
+        return run_sync(self._prefect_client.hello())
+
+    def create_task_run(
+        self,
+        task: "TaskObject[P, R]",
+        flow_run_id: Optional[UUID],
+        dynamic_key: str,
+        name: Optional[str] = None,
+        extra_tags: Optional[Iterable[str]] = None,
+        state: Optional[prefect.states.State[R]] = None,
+        task_inputs: Optional[
+            Dict[
+                str,
+                List[
+                    Union[
+                        TaskRunResult,
+                        Parameter,
+                        Constant,
+                    ]
+                ],
+            ]
+        ] = None,
+    ) -> TaskRun:
+        """
+        Create a task run
+
+        Args:
+            task: The Task to run
+            flow_run_id: The flow run id with which to associate the task run
+            dynamic_key: A key unique to this particular run of a Task within the flow
+            name: An optional name for the task run
+            extra_tags: an optional list of extra tags to apply to the task run in
+                addition to `task.tags`
+            state: The initial state for the run. If not provided, defaults to
+                `Pending` for now. Should always be a `Scheduled` type.
+            task_inputs: the set of inputs passed to the task
+
+        Returns:
+            The created task run.
+        """
+        return run_sync(
+            self._prefect_client.create_task_run(
+                task=task,
+                flow_run_id=flow_run_id,
+                dynamic_key=dynamic_key,
+                name=name,
+                extra_tags=extra_tags,
+                state=state,
+                task_inputs=task_inputs,
+            )
+        )
+
+    def set_task_run_state(
+        self,
+        task_run_id: UUID,
+        state: prefect.states.State,
+        force: bool = False,
+    ) -> OrchestrationResult:
+        """
+        Set the state of a task run.
+
+        Args:
+            task_run_id: the id of the task run
+            state: the state to set
+            force: if True, disregard orchestration logic when setting the state,
+                forcing the Prefect API to accept the state
+
+        Returns:
+            an OrchestrationResult model representation of state orchestration output
+        """
+        return run_sync(
+            self._prefect_client.set_task_run_state(
+                task_run_id=task_run_id,
+                state=state,
+                force=force,
+            )
+        )
+
+    def create_flow_run(
+        self,
+        flow_id: UUID,
+        parameters: Optional[Dict[str, Any]] = None,
+        context: Optional[Dict[str, Any]] = None,
+        scheduled_start_time: Optional[datetime.datetime] = None,
+        run_name: Optional[str] = None,
+        labels: Optional[List[str]] = None,
+        parameters_json: Optional[str] = None,
+        run_config: Optional[Dict[str, Any]] = None,
+        idempotency_key: Optional[str] = None,
+    ) -> FlowRunResponse:
+        """
+        Create a new flow run.
+
+        Args:
+            - flow_id (UUID): the ID of the flow to create a run for
+            - parameters (Optional[Dict[str, Any]]): a dictionary of parameter values to pass to the flow
+            - context (Optional[Dict[str, Any]]): a dictionary of context values to pass to the flow
+            - scheduled_start_time (Optional[datetime.datetime]): the scheduled start time for the flow run
+            - run_name (Optional[str]): a name to assign to the flow run
+            - labels (Optional[List[str]]): a list of labels to assign to the flow run
+            - parameters_json (Optional[str]): a JSON string of parameter values to pass to the flow
+            - run_config (Optional[Dict[str, Any]]): a dictionary of run configuration options
+            - idempotency_key (Optional[str]): a key to ensure idempotency when creating the flow run
+
+        Returns:
+            - FlowRunResponse: the created flow run
+        """
+        return run_sync(
+            self._prefect_client.create_flow_run(
+                flow_id=flow_id,
+                parameters=parameters,
+                context=context,
+                scheduled_start_time=scheduled_start_time,
+                run_name=run_name,
+                labels=labels,
+                parameters_json=parameters_json,
+                run_config=run_config,
+                idempotency_key=idempotency_key,
+            )
+        )
+
+    async def set_flow_run_state(
+        self,
+        flow_run_id: UUID,
+        state: "prefect.states.State",
+        force: bool = False,
+    ) -> OrchestrationResult:
+        """
+        Set the state of a flow run.
+
+        Args:
+            flow_run_id: the id of the flow run
+            state: the state to set
+            force: if True, disregard orchestration logic when setting the state,
+                forcing the Prefect API to accept the state
+
+        Returns:
+            an OrchestrationResult model representation of state orchestration output
+        """
+        return run_sync(
+            self._prefect_client.set_flow_run_state(
+                flow_run_id=flow_run_id,
+                state=state,
+                force=force,
+            )
+        )

prefect/client/schemas/sorting.py

@@ -27,6 +27,15 @@ class TaskRunSort(AutoEnum):
     END_TIME_DESC = AutoEnum.auto()
 
 
+class AutomationSort(AutoEnum):
+    """Defines automation sorting options."""
+
+    CREATED_DESC = AutoEnum.auto()
+    UPDATED_DESC = AutoEnum.auto()
+    NAME_ASC = AutoEnum.auto()
+    NAME_DESC = AutoEnum.auto()
+
+
 class LogSort(AutoEnum):
     """Defines log sorting options."""
 

prefect/client/utilities.py

@@ -6,14 +6,25 @@ Utilities for working with clients.
 # circular imports for decorators such as `inject_client` which are widely used.
 
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, Tuple, cast
-
-from typing_extensions import ParamSpec
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Optional,
+    Tuple,
+    TypeVar,
+    cast,
+)
+
+from typing_extensions import Concatenate, ParamSpec
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
 P = ParamSpec("P")
+R = TypeVar("R")
 
 
 def get_or_create_client(
@@ -52,6 +63,17 @@ def get_or_create_client(
         return get_httpx_client(), False
 
 
+def client_injector(
+    func: Callable[Concatenate["PrefectClient", P], Awaitable[R]],
+) -> Callable[P, Awaitable[R]]:
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        client, _ = get_or_create_client()
+        return await func(client, *args, **kwargs)
+
+    return wrapper
+
+
 def inject_client(
     fn: Callable[P, Coroutine[Any, Any, Any]],
 ) -> Callable[P, Coroutine[Any, Any, Any]]:

prefect/concurrency/asyncio.py

@@ -1,10 +1,16 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import List, Literal, Union
+from typing import List, Literal, Union, cast
 
 import httpx
 import pendulum
 
+try:
+    from pendulum import Interval
+except ImportError:
+    # pendulum < 3
+    from pendulum.period import Period as Interval  # type: ignore
+
 from prefect import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
 
@@ -30,10 +36,10 @@ async def concurrency(names: Union[str, List[str]], occupy: int = 1):
     try:
         yield
     finally:
-        occupancy_seconds: float = (
-            pendulum.now("UTC") - acquisition_time
-        ).total_seconds()
-        await _release_concurrency_slots(names, occupy, occupancy_seconds)
+        occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
+        await _release_concurrency_slots(
+            names, occupy, occupancy_period.total_seconds()
+        )
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 

prefect/concurrency/events.py

@@ -12,11 +12,11 @@ def _emit_concurrency_event(
     slots: int,
     follows: Union[Event, None] = None,
 ) -> Union[Event, None]:
-    resource = {
+    resource: Dict[str, str] = {
         "prefect.resource.id": f"prefect.concurrency-limit.{primary_limit.id}",
         "prefect.resource.name": primary_limit.name,
-        "slots-acquired": slots,
-        "limit": primary_limit.limit,
+        "slots-acquired": str(slots),
+        "limit": str(primary_limit.limit),
     }
 
     related = [

prefect/concurrency/services.py

@@ -64,7 +64,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
                 raise RuntimeError("Cannot put items in a stopped service instance.")
 
             logger.debug("Service %r enqueuing item %r", self, item)
-            future = concurrent.futures.Future()
+            future: concurrent.futures.Future = concurrent.futures.Future()
 
             occupy, mode = item
             self._queue.put_nowait((occupy, mode, future))

prefect/concurrency/sync.py

@@ -1,8 +1,14 @@
 from contextlib import contextmanager
-from typing import List, Union
+from typing import List, Union, cast
 
 import pendulum
 
+try:
+    from pendulum import Interval
+except ImportError:
+    # pendulum < 3
+    from pendulum.period import Period as Interval  # type: ignore
+
 from prefect._internal.concurrency.api import create_call, from_sync
 from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
@@ -30,11 +36,9 @@ def concurrency(names: Union[str, List[str]], occupy: int = 1):
     try:
         yield
     finally:
-        occupancy_seconds: float = (
-            pendulum.now("UTC") - acquisition_time
-        ).total_seconds()
+        occupancy_period = cast(Interval, pendulum.now("UTC") - acquisition_time)
         _call_async_function_from_sync(
-            _release_concurrency_slots, names, occupy, occupancy_seconds
+            _release_concurrency_slots, names, occupy, occupancy_period.total_seconds()
         )
         _emit_concurrency_release_events(limits, occupy, emitted_events)