prefect-client 3.1.15__py3-none-any.whl → 3.2.0__py3-none-any.whl

prefect/client/orchestration/_work_pools/client.py

@@ -0,0 +1,598 @@
+from __future__ import annotations
+
+import warnings
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from httpx import HTTPStatusError
+
+from prefect.client.base import ServerType
+from prefect.client.orchestration.base import BaseAsyncClient, BaseClient
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from prefect.client.schemas.actions import (
+        WorkPoolCreate,
+        WorkPoolUpdate,
+    )
+    from prefect.client.schemas.filters import (
+        WorkerFilter,
+        WorkPoolFilter,
+    )
+    from prefect.client.schemas.objects import (
+        Worker,
+        WorkerMetadata,
+        WorkPool,
+    )
+    from prefect.client.schemas.responses import WorkerFlowRunResponse
+
+from prefect.exceptions import ObjectAlreadyExists, ObjectNotFound
+
+
+class WorkPoolClient(BaseClient):
+    def send_worker_heartbeat(
+        self,
+        work_pool_name: str,
+        worker_name: str,
+        heartbeat_interval_seconds: float | None = None,
+        get_worker_id: bool = False,
+        worker_metadata: "WorkerMetadata | None" = None,
+    ) -> "UUID | None":
+        """
+        Sends a worker heartbeat for a given work pool.
+
+        Args:
+            work_pool_name: The name of the work pool to heartbeat against.
+            worker_name: The name of the worker sending the heartbeat.
+            return_id: Whether to return the worker ID. Note: will return `None` if the connected server does not support returning worker IDs, even if `return_id` is `True`.
+            worker_metadata: Metadata about the worker to send to the server.
+        """
+        from uuid import UUID
+
+        params: dict[str, Any] = {
+            "name": worker_name,
+            "heartbeat_interval_seconds": heartbeat_interval_seconds,
+        }
+        if worker_metadata:
+            params["metadata"] = worker_metadata.model_dump(mode="json")
+        if get_worker_id:
+            params["return_id"] = get_worker_id
+
+        resp = self.request(
+            "POST",
+            "/work_pools/{work_pool_name}/workers/heartbeat",
+            path_params={"work_pool_name": work_pool_name},
+            json=params,
+        )
+        from prefect.settings import get_current_settings
+
+        if (
+            (
+                self.server_type == ServerType.CLOUD
+                or get_current_settings().testing.test_mode
+            )
+            and get_worker_id
+            and resp.status_code == 200
+        ):
+            return UUID(resp.text)
+        else:
+            return None
+
+    def read_workers_for_work_pool(
+        self,
+        work_pool_name: str,
+        worker_filter: "WorkerFilter | None" = None,
+        offset: int | None = None,
+        limit: int | None = None,
+    ) -> list["Worker"]:
+        """
+        Reads workers for a given work pool.
+
+        Args:
+            work_pool_name: The name of the work pool for which to get
+                member workers.
+            worker_filter: Criteria by which to filter workers.
+            limit: Limit for the worker query.
+            offset: Limit for the worker query.
+        """
+        from prefect.client.schemas.objects import Worker
+
+        response = self.request(
+            "POST",
+            "/work_pools/{work_pool_name}/workers/filter",
+            path_params={"work_pool_name": work_pool_name},
+            json={
+                "workers": (
+                    worker_filter.model_dump(mode="json", exclude_unset=True)
+                    if worker_filter
+                    else None
+                ),
+                "offset": offset,
+                "limit": limit,
+            },
+        )
+
+        return Worker.model_validate_list(response.json())
+
+    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.
+        """
+        from prefect.client.schemas.objects import WorkPool
+
+        try:
+            response = self.request(
+                "GET",
+                "/work_pools/{name}",
+                path_params={"name": work_pool_name},
+            )
+            return WorkPool.model_validate(response.json())
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+    def read_work_pools(
+        self,
+        limit: int | None = None,
+        offset: int = 0,
+        work_pool_filter: "WorkPoolFilter | None" = None,
+    ) -> list["WorkPool"]:
+        """
+        Reads work pools.
+
+        Args:
+            limit: Limit for the work pool query.
+            offset: Offset for the work pool query.
+            work_pool_filter: Criteria by which to filter work pools.
+
+        Returns:
+            A list of work pools.
+        """
+        from prefect.client.schemas.objects import WorkPool
+
+        body: dict[str, Any] = {
+            "limit": limit,
+            "offset": offset,
+            "work_pools": (
+                work_pool_filter.model_dump(mode="json") if work_pool_filter else None
+            ),
+        }
+        response = self.request("POST", "/work_pools/filter", json=body)
+        return WorkPool.model_validate_list(response.json())
+
+    def create_work_pool(
+        self,
+        work_pool: "WorkPoolCreate",
+        overwrite: bool = False,
+    ) -> "WorkPool":
+        """
+        Creates a work pool with the provided configuration.
+
+        Args:
+            work_pool: Desired configuration for the new work pool.
+
+        Returns:
+            Information about the newly created work pool.
+        """
+        from prefect.client.schemas.actions import WorkPoolUpdate
+        from prefect.client.schemas.objects import WorkPool
+
+        try:
+            response = self.request(
+                "POST",
+                "/work_pools/",
+                json=work_pool.model_dump(mode="json", exclude_unset=True),
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 409:
+                if overwrite:
+                    existing_work_pool = self.read_work_pool(
+                        work_pool_name=work_pool.name
+                    )
+                    if existing_work_pool.type != work_pool.type:
+                        warnings.warn(
+                            "Overwriting work pool type is not supported. Ignoring provided type.",
+                            category=UserWarning,
+                        )
+                    self.update_work_pool(
+                        work_pool_name=work_pool.name,
+                        work_pool=WorkPoolUpdate.model_validate(
+                            work_pool.model_dump(exclude={"name", "type"})
+                        ),
+                    )
+                    response = self.request(
+                        "GET",
+                        "/work_pools/{name}",
+                        path_params={"name": work_pool.name},
+                    )
+                else:
+                    raise ObjectAlreadyExists(http_exc=e) from e
+            else:
+                raise
+
+        return WorkPool.model_validate(response.json())
+
+    def update_work_pool(
+        self,
+        work_pool_name: str,
+        work_pool: "WorkPoolUpdate",
+    ) -> None:
+        """
+        Updates a work pool.
+
+        Args:
+            work_pool_name: Name of the work pool to update.
+            work_pool: Fields to update in the work pool.
+        """
+        try:
+            self.request(
+                "PATCH",
+                "/work_pools/{name}",
+                path_params={"name": work_pool_name},
+                json=work_pool.model_dump(mode="json", exclude_unset=True),
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+    def delete_work_pool(
+        self,
+        work_pool_name: str,
+    ) -> None:
+        """
+        Deletes a work pool.
+
+        Args:
+            work_pool_name: Name of the work pool to delete.
+        """
+        try:
+            self.request(
+                "DELETE",
+                "/work_pools/{name}",
+                path_params={"name": work_pool_name},
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+    def get_scheduled_flow_runs_for_work_pool(
+        self,
+        work_pool_name: str,
+        work_queue_names: list[str] | None = None,
+        scheduled_before: datetime | None = None,
+    ) -> list["WorkerFlowRunResponse"]:
+        """
+        Retrieves scheduled flow runs for the provided set of work pool queues.
+
+        Args:
+            work_pool_name: The name of the work pool that the work pool
+                queues are associated with.
+            work_queue_names: The names of the work pool queues from which
+                to get scheduled flow runs.
+            scheduled_before: Datetime used to filter returned flow runs. Flow runs
+                scheduled for after the given datetime string will not be returned.
+
+        Returns:
+            A list of worker flow run responses containing information about the
+            retrieved flow runs.
+        """
+        from prefect.client.schemas.responses import WorkerFlowRunResponse
+
+        body: dict[str, Any] = {}
+        if work_queue_names is not None:
+            body["work_queue_names"] = list(work_queue_names)
+        if scheduled_before:
+            body["scheduled_before"] = str(scheduled_before)
+
+        try:
+            response = self.request(
+                "POST",
+                "/work_pools/{name}/get_scheduled_flow_runs",
+                path_params={"name": work_pool_name},
+                json=body,
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+        return WorkerFlowRunResponse.model_validate_list(response.json())
+
+
+class WorkPoolAsyncClient(BaseAsyncClient):
+    async def send_worker_heartbeat(
+        self,
+        work_pool_name: str,
+        worker_name: str,
+        heartbeat_interval_seconds: float | None = None,
+        get_worker_id: bool = False,
+        worker_metadata: "WorkerMetadata | None" = None,
+    ) -> "UUID | None":
+        """
+        Sends a worker heartbeat for a given work pool.
+
+        Args:
+            work_pool_name: The name of the work pool to heartbeat against.
+            worker_name: The name of the worker sending the heartbeat.
+            return_id: Whether to return the worker ID. Note: will return `None` if the connected server does not support returning worker IDs, even if `return_id` is `True`.
+            worker_metadata: Metadata about the worker to send to the server.
+        """
+        from uuid import UUID
+
+        params: dict[str, Any] = {
+            "name": worker_name,
+            "heartbeat_interval_seconds": heartbeat_interval_seconds,
+        }
+        if worker_metadata:
+            params["metadata"] = worker_metadata.model_dump(mode="json")
+        if get_worker_id:
+            params["return_id"] = get_worker_id
+
+        resp = await self.request(
+            "POST",
+            "/work_pools/{work_pool_name}/workers/heartbeat",
+            path_params={"work_pool_name": work_pool_name},
+            json=params,
+        )
+        from prefect.settings import get_current_settings
+
+        if (
+            (
+                self.server_type == ServerType.CLOUD
+                or get_current_settings().testing.test_mode
+            )
+            and get_worker_id
+            and resp.status_code == 200
+        ):
+            return UUID(resp.text)
+        else:
+            return None
+
+    async def read_workers_for_work_pool(
+        self,
+        work_pool_name: str,
+        worker_filter: "WorkerFilter | None" = None,
+        offset: int | None = None,
+        limit: int | None = None,
+    ) -> list["Worker"]:
+        """
+        Reads workers for a given work pool.
+
+        Args:
+            work_pool_name: The name of the work pool for which to get
+                member workers.
+            worker_filter: Criteria by which to filter workers.
+            limit: Limit for the worker query.
+            offset: Limit for the worker query.
+        """
+        from prefect.client.schemas.objects import Worker
+
+        response = await self.request(
+            "POST",
+            "/work_pools/{work_pool_name}/workers/filter",
+            path_params={"work_pool_name": work_pool_name},
+            json={
+                "workers": (
+                    worker_filter.model_dump(mode="json", exclude_unset=True)
+                    if worker_filter
+                    else None
+                ),
+                "offset": offset,
+                "limit": limit,
+            },
+        )
+
+        return Worker.model_validate_list(response.json())
+
+    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.
+        """
+        from prefect.client.schemas.objects import WorkPool
+
+        try:
+            response = await self.request(
+                "GET",
+                "/work_pools/{name}",
+                path_params={"name": work_pool_name},
+            )
+            return WorkPool.model_validate(response.json())
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+    async def read_work_pools(
+        self,
+        limit: int | None = None,
+        offset: int = 0,
+        work_pool_filter: "WorkPoolFilter | None" = None,
+    ) -> list["WorkPool"]:
+        """
+        Reads work pools.
+
+        Args:
+            limit: Limit for the work pool query.
+            offset: Offset for the work pool query.
+            work_pool_filter: Criteria by which to filter work pools.
+
+        Returns:
+            A list of work pools.
+        """
+        from prefect.client.schemas.objects import WorkPool
+
+        body: dict[str, Any] = {
+            "limit": limit,
+            "offset": offset,
+            "work_pools": (
+                work_pool_filter.model_dump(mode="json") if work_pool_filter else None
+            ),
+        }
+        response = await self.request("POST", "/work_pools/filter", json=body)
+        return WorkPool.model_validate_list(response.json())
+
+    async def create_work_pool(
+        self,
+        work_pool: "WorkPoolCreate",
+        overwrite: bool = False,
+    ) -> "WorkPool":
+        """
+        Creates a work pool with the provided configuration.
+
+        Args:
+            work_pool: Desired configuration for the new work pool.
+
+        Returns:
+            Information about the newly created work pool.
+        """
+        from prefect.client.schemas.actions import WorkPoolUpdate
+        from prefect.client.schemas.objects import WorkPool
+
+        try:
+            response = await self.request(
+                "POST",
+                "/work_pools/",
+                json=work_pool.model_dump(mode="json", exclude_unset=True),
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 409:
+                if overwrite:
+                    existing_work_pool = await self.read_work_pool(
+                        work_pool_name=work_pool.name
+                    )
+                    if existing_work_pool.type != work_pool.type:
+                        warnings.warn(
+                            "Overwriting work pool type is not supported. Ignoring provided type.",
+                            category=UserWarning,
+                        )
+                    await self.update_work_pool(
+                        work_pool_name=work_pool.name,
+                        work_pool=WorkPoolUpdate.model_validate(
+                            work_pool.model_dump(exclude={"name", "type"})
+                        ),
+                    )
+                    response = await self.request(
+                        "GET",
+                        "/work_pools/{name}",
+                        path_params={"name": work_pool.name},
+                    )
+                else:
+                    raise ObjectAlreadyExists(http_exc=e) from e
+            else:
+                raise
+
+        return WorkPool.model_validate(response.json())
+
+    async def update_work_pool(
+        self,
+        work_pool_name: str,
+        work_pool: "WorkPoolUpdate",
+    ) -> None:
+        """
+        Updates a work pool.
+
+        Args:
+            work_pool_name: Name of the work pool to update.
+            work_pool: Fields to update in the work pool.
+        """
+        try:
+            await self.request(
+                "PATCH",
+                "/work_pools/{name}",
+                path_params={"name": work_pool_name},
+                json=work_pool.model_dump(mode="json", exclude_unset=True),
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+    async def delete_work_pool(
+        self,
+        work_pool_name: str,
+    ) -> None:
+        """
+        Deletes a work pool.
+
+        Args:
+            work_pool_name: Name of the work pool to delete.
+        """
+        try:
+            await self.request(
+                "DELETE",
+                "/work_pools/{name}",
+                path_params={"name": work_pool_name},
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+    async def get_scheduled_flow_runs_for_work_pool(
+        self,
+        work_pool_name: str,
+        work_queue_names: list[str] | None = None,
+        scheduled_before: datetime | None = None,
+    ) -> list["WorkerFlowRunResponse"]:
+        """
+        Retrieves scheduled flow runs for the provided set of work pool queues.
+
+        Args:
+            work_pool_name: The name of the work pool that the work pool
+                queues are associated with.
+            work_queue_names: The names of the work pool queues from which
+                to get scheduled flow runs.
+            scheduled_before: Datetime used to filter returned flow runs. Flow runs
+                scheduled for after the given datetime string will not be returned.
+
+        Returns:
+            A list of worker flow run responses containing information about the
+            retrieved flow runs.
+        """
+        from prefect.client.schemas.responses import WorkerFlowRunResponse
+
+        body: dict[str, Any] = {}
+        if work_queue_names is not None:
+            body["work_queue_names"] = list(work_queue_names)
+        if scheduled_before:
+            body["scheduled_before"] = str(scheduled_before)
+
+        try:
+            response = await self.request(
+                "POST",
+                "/work_pools/{name}/get_scheduled_flow_runs",
+                path_params={"name": work_pool_name},
+                json=body,
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+
+        return WorkerFlowRunResponse.model_validate_list(response.json())

prefect/client/orchestration/base.py

@@ -7,12 +7,15 @@ from typing_extensions import TypeAlias
 if TYPE_CHECKING:
     from httpx import AsyncClient, Client, Response
 
+    from prefect.client.base import ServerType
     from prefect.client.orchestration.routes import ServerRoutes
 
 HTTP_METHODS: TypeAlias = Literal["GET", "POST", "PUT", "DELETE", "PATCH"]
 
 
 class BaseClient:
+    server_type: "ServerType"
+
     def __init__(self, client: "Client"):
         self._client = client
 
@@ -26,10 +29,13 @@ class BaseClient:
     ) -> "Response":
         if path_params:
             path = path.format(**path_params)  # type: ignore
-        return self._client.request(method, path, params=params, **kwargs)
+        request = self._client.build_request(method, path, params=params, **kwargs)
+        return self._client.send(request)
 
 
 class BaseAsyncClient:
+    server_type: "ServerType"
+
     def __init__(self, client: "AsyncClient"):
         self._client = client
 
@@ -43,4 +49,5 @@ class BaseAsyncClient:
     ) -> "Response":
         if path_params:
             path = path.format(**path_params)  # type: ignore
-        return await self._client.request(method, path, params=params, **kwargs)
+        request = self._client.build_request(method, path, params=params, **kwargs)
+        return await self._client.send(request)