prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/cache_policies.py

@@ -0,0 +1,239 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional
+
+from prefect.context import TaskRunContext
+from prefect.utilities.hashing import hash_objects
+
+
+@dataclass
+class CachePolicy:
+    """
+    Base class for all cache policies.
+    """
+
+    @classmethod
+    def from_cache_key_fn(
+        cls, cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ) -> "CacheKeyFnPolicy":
+        """
+        Given a function generates a key policy.
+        """
+        return CacheKeyFnPolicy(cache_key_fn=cache_key_fn)
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        raise NotImplementedError
+
+    def __sub__(self, other: str) -> "CompoundCachePolicy":
+        if not isinstance(other, str):
+            raise TypeError("Can only subtract strings from key policies.")
+        if isinstance(self, Inputs):
+            exclude = self.exclude or []
+            return Inputs(exclude=exclude + [other])
+        elif isinstance(self, CompoundCachePolicy):
+            new = Inputs(exclude=[other])
+            policies = self.policies or []
+            return CompoundCachePolicy(policies=policies + [new])
+        else:
+            new = Inputs(exclude=[other])
+            return CompoundCachePolicy(policies=[self, new])
+
+    def __add__(self, other: "CachePolicy") -> "CompoundCachePolicy":
+        # adding _None is a no-op
+        if isinstance(other, _None):
+            return self
+        elif isinstance(self, _None):
+            return other
+
+        if isinstance(self, CompoundCachePolicy):
+            policies = self.policies or []
+            return CompoundCachePolicy(policies=policies + [other])
+        elif isinstance(other, CompoundCachePolicy):
+            policies = other.policies or []
+            return CompoundCachePolicy(policies=policies + [self])
+        else:
+            return CompoundCachePolicy(policies=[self, other])
+
+
+@dataclass
+class CacheKeyFnPolicy(CachePolicy):
+    """
+    This policy accepts a custom function with signature f(task_run_context, task_parameters, flow_parameters) -> str
+    and uses it to compute a task run cache key.
+    """
+
+    # making it optional for tests
+    cache_key_fn: Optional[
+        Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ] = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        if self.cache_key_fn:
+            return self.cache_key_fn(task_ctx, inputs)
+
+
+@dataclass
+class CompoundCachePolicy(CachePolicy):
+    """
+    This policy is constructed from two or more other cache policies and works by computing the keys
+    for each policy individually, and then hashing a sorted tuple of all computed keys.
+
+    Any keys that return `None` will be ignored.
+    """
+
+    policies: Optional[list] = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        keys = []
+        for policy in self.policies or []:
+            policy_key = policy.compute_key(
+                task_ctx=task_ctx,
+                inputs=inputs,
+                flow_parameters=flow_parameters,
+                **kwargs,
+            )
+            if policy_key is not None:
+                keys.append(policy_key)
+        if not keys:
+            return None
+        return hash_objects(*keys)
+
+
+@dataclass
+class _None(CachePolicy):
+    """
+    Policy that always returns `None` for the computed cache key.
+    This policy prevents persistence.
+    """
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        return None
+
+
+@dataclass
+class TaskSource(CachePolicy):
+    """
+    Policy for computing a cache key based on the source code of the task.
+    """
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Optional[Dict[str, Any]],
+        flow_parameters: Optional[Dict[str, Any]],
+        **kwargs,
+    ) -> Optional[str]:
+        if not task_ctx:
+            return None
+        try:
+            lines = inspect.getsource(task_ctx.task)
+        except TypeError:
+            lines = inspect.getsource(task_ctx.task.fn.__class__)
+        except OSError as exc:
+            if "could not get source code" in str(exc):
+                lines = task_ctx.task.fn.__code__.co_code
+            else:
+                raise
+
+        return hash_objects(lines)
+
+
+@dataclass
+class FlowParameters(CachePolicy):
+    """
+    Policy that computes the cache key based on a hash of the flow parameters.
+    """
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        if not flow_parameters:
+            return None
+        return hash_objects(flow_parameters)
+
+
+@dataclass
+class RunId(CachePolicy):
+    """
+    Returns either the prevailing flow run ID, or if not found, the prevailing task
+    run ID.
+    """
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        if not task_ctx:
+            return None
+        run_id = task_ctx.task_run.flow_run_id
+        if run_id is None:
+            run_id = task_ctx.task_run.id
+        return str(run_id)
+
+
+@dataclass
+class Inputs(CachePolicy):
+    """
+    Policy that computes a cache key based on a hash of the runtime inputs provided to the task..
+    """
+
+    exclude: Optional[list] = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        hashed_inputs = {}
+        inputs = inputs or {}
+        exclude = self.exclude or []
+
+        if not inputs:
+            return None
+
+        for key, val in inputs.items():
+            if key not in exclude:
+                hashed_inputs[key] = val
+
+        return hash_objects(hashed_inputs)
+
+
+INPUTS = Inputs()
+NONE = _None()
+TASK_SOURCE = TaskSource()
+FLOW_PARAMETERS = FlowParameters()
+RUN_ID = RunId()
+DEFAULT = INPUTS + TASK_SOURCE + RUN_ID

prefect/client/__init__.py

@@ -15,3 +15,7 @@ $ python -m asyncio
 ```
 </div>
 """
+
+from prefect._internal.compatibility.migration import getattr_migration
+
+__getattr__ = getattr_migration(__name__)

prefect/client/base.py

@@ -25,8 +25,7 @@ import anyio
 import httpx
 from asgi_lifespan import LifespanManager
 from httpx import HTTPStatusError, Request, Response
-from prefect._vendor.starlette import status
-from prefect._vendor.starlette.testclient import TestClient
+from starlette import status
 from typing_extensions import Self
 
 import prefect
@@ -35,10 +34,14 @@ from prefect.client.schemas.objects import CsrfToken
 from prefect.exceptions import PrefectHTTPStatusError
 from prefect.logging import get_logger
 from prefect.settings import (
+    PREFECT_API_URL,
     PREFECT_CLIENT_MAX_RETRIES,
     PREFECT_CLIENT_RETRY_EXTRA_CODES,
     PREFECT_CLIENT_RETRY_JITTER_FACTOR,
+    PREFECT_CLOUD_API_URL,
+    PREFECT_SERVER_ALLOW_EPHEMERAL_MODE,
 )
+from prefect.utilities.collections import AutoEnum
 from prefect.utilities.math import bounded_poisson_interval, clamped_poisson_interval
 
 # Datastores for lifespan management, keys should be a tuple of thread and app
@@ -390,7 +393,7 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
 
                 raise
 
-            token: CsrfToken = CsrfToken.parse_obj(token_response.json())
+            token: CsrfToken = CsrfToken.model_validate(token_response.json())
             self.csrf_token = token.token
             self.csrf_token_expiration = token.expiration
 
@@ -604,7 +607,7 @@ class PrefectHttpxSyncClient(httpx.Client):
 
                 raise
 
-            token: CsrfToken = CsrfToken.parse_obj(token_response.json())
+            token: CsrfToken = CsrfToken.model_validate(token_response.json())
             self.csrf_token = token.token
             self.csrf_token_expiration = token.expiration
 
@@ -612,28 +615,31 @@ class PrefectHttpxSyncClient(httpx.Client):
         request.headers["Prefect-Csrf-Client"] = str(self.csrf_client_id)
 
 
-class PrefectHttpxSyncEphemeralClient(TestClient, PrefectHttpxSyncClient):
-    """
-    This client is a synchronous httpx client that can be used to talk directly
-    to an ASGI app, such as an ephemeral Prefect API.
-
-    It is a subclass of both Starlette's `TestClient` and Prefect's
-    `PrefectHttpxSyncClient`, so it combines the synchronous testing
-    capabilities of `TestClient` with the Prefect-specific behaviors of
-    `PrefectHttpxSyncClient`.
-    """
+class ServerType(AutoEnum):
+    EPHEMERAL = AutoEnum.auto()
+    SERVER = AutoEnum.auto()
+    CLOUD = AutoEnum.auto()
+    UNCONFIGURED = AutoEnum.auto()
 
-    def __init__(
-        self,
-        *args,
-        # override TestClient default
-        raise_server_exceptions=False,
-        **kwargs,
-    ):
-        super().__init__(
-            *args,
-            raise_server_exceptions=raise_server_exceptions,
-            **kwargs,
-        )
 
-    pass
+def determine_server_type() -> ServerType:
+    """
+    Determine the server type based on the current settings.
+
+    Returns:
+        - `ServerType.EPHEMERAL` if the ephemeral server is enabled
+        - `ServerType.SERVER` if a API URL is configured and it is not a cloud URL
+        - `ServerType.CLOUD` if an API URL is configured and it is a cloud URL
+        - `ServerType.UNCONFIGURED` if no API URL is configured and ephemeral mode is
+            not enabled
+    """
+    api_url = PREFECT_API_URL.value()
+    if api_url is None:
+        if PREFECT_SERVER_ALLOW_EPHEMERAL_MODE.value():
+            return ServerType.EPHEMERAL
+        else:
+            return ServerType.UNCONFIGURED
+    if api_url.startswith(PREFECT_CLOUD_API_URL.value()):
+        return ServerType.CLOUD
+    else:
+        return ServerType.SERVER

prefect/client/cloud.py

@@ -1,23 +1,20 @@
 import re
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 
 import anyio
 import httpx
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic
-else:
-    import pydantic
-
-from prefect._vendor.starlette import status
+import pydantic
+from starlette import status
 
 import prefect.context
 import prefect.settings
 from prefect.client.base import PrefectHttpxAsyncClient
-from prefect.client.schemas import Workspace
-from prefect.exceptions import PrefectException
+from prefect.client.schemas.objects import (
+    IPAllowlist,
+    IPAllowlistMyAccessResponse,
+    Workspace,
+)
+from prefect.exceptions import ObjectNotFound, PrefectException
 from prefect.settings import (
     PREFECT_API_KEY,
     PREFECT_CLOUD_API_URL,
@@ -76,6 +73,27 @@ class CloudClient:
             **httpx_settings, enable_csrf_support=False
         )
 
+        api_url = prefect.settings.PREFECT_API_URL.value() or ""
+        if match := (
+            re.search(PARSE_API_URL_REGEX, host)
+            or re.search(PARSE_API_URL_REGEX, api_url)
+        ):
+            self.account_id, self.workspace_id = match.groups()
+
+    @property
+    def account_base_url(self) -> str:
+        if not self.account_id:
+            raise ValueError("Account ID not set")
+
+        return f"accounts/{self.account_id}"
+
+    @property
+    def workspace_base_url(self) -> str:
+        if not self.workspace_id:
+            raise ValueError("Workspace ID not set")
+
+        return f"{self.account_base_url}/workspaces/{self.workspace_id}"
+
     async def api_healthcheck(self):
         """
         Attempts to connect to the Cloud API and raises the encountered exception if not
@@ -87,18 +105,43 @@ class CloudClient:
             await self.read_workspaces()
 
     async def read_workspaces(self) -> List[Workspace]:
-        workspaces = pydantic.parse_obj_as(
-            List[Workspace], await self.get("/me/workspaces")
+        workspaces = pydantic.TypeAdapter(List[Workspace]).validate_python(
+            await self.get("/me/workspaces")
         )
         return workspaces
 
     async def read_worker_metadata(self) -> Dict[str, Any]:
-        configured_url = prefect.settings.PREFECT_API_URL.value()
-        account_id, workspace_id = re.findall(PARSE_API_URL_REGEX, configured_url)[0]
-        return await self.get(
-            f"accounts/{account_id}/workspaces/{workspace_id}/collections/work_pool_types"
+        response = await self.get(
+            f"{self.workspace_base_url}/collections/work_pool_types"
+        )
+        return cast(Dict[str, Any], response)
+
+    async def read_account_settings(self) -> Dict[str, Any]:
+        response = await self.get(f"{self.account_base_url}/settings")
+        return cast(Dict[str, Any], response)
+
+    async def update_account_settings(self, settings: Dict[str, Any]):
+        await self.request(
+            "PATCH",
+            f"{self.account_base_url}/settings",
+            json=settings,
         )
 
+    async def read_account_ip_allowlist(self) -> IPAllowlist:
+        response = await self.get(f"{self.account_base_url}/ip_allowlist")
+        return IPAllowlist.model_validate(response)
+
+    async def update_account_ip_allowlist(self, updated_allowlist: IPAllowlist):
+        await self.request(
+            "PUT",
+            f"{self.account_base_url}/ip_allowlist",
+            json=updated_allowlist.model_dump(mode="json"),
+        )
+
+    async def check_ip_allowlist_access(self) -> IPAllowlistMyAccessResponse:
+        response = await self.get(f"{self.account_base_url}/ip_allowlist/my_access")
+        return IPAllowlistMyAccessResponse.model_validate(response)
+
     async def __aenter__(self):
         await self._client.__aenter__()
         return self
@@ -127,9 +170,11 @@ class CloudClient:
                 status.HTTP_401_UNAUTHORIZED,
                 status.HTTP_403_FORBIDDEN,
             ):
-                raise CloudUnauthorizedError
+                raise CloudUnauthorizedError(str(exc)) from exc
+            elif exc.response.status_code == status.HTTP_404_NOT_FOUND:
+                raise ObjectNotFound(http_exc=exc) from exc
             else:
-                raise exc
+                raise
 
         if res.status_code == status.HTTP_204_NO_CONTENT:
             return

prefect/client/collections.py

@@ -29,6 +29,6 @@ def get_collections_metadata_client(
     """
     orchestration_client = get_client(httpx_settings=httpx_settings)
     if orchestration_client.server_type == ServerType.CLOUD:
-        return get_cloud_client(httpx_settings=httpx_settings)
+        return get_cloud_client(httpx_settings=httpx_settings, infer_cloud_url=True)
     else:
         return orchestration_client