prefect-client 3.0.1__py3-none-any.whl → 3.0.3__py3-none-any.whl

prefect/_internal/compatibility/deprecated.py

@@ -30,7 +30,7 @@ M = TypeVar("M", bound=BaseModel)
 
 
 DEPRECATED_WARNING = (
-    "{name} has been deprecated{when}. It will not be available after {end_date}."
+    "{name} has been deprecated{when}. It will not be available in new releases after {end_date}."
     " {help}"
 )
 DEPRECATED_MOVED_WARNING = (

prefect/blocks/core.py

@@ -40,7 +40,6 @@ from pydantic import (
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal, ParamSpec, Self, get_args
 
-import prefect
 import prefect.exceptions
 from prefect.client.schemas import (
     DEFAULT_BLOCK_SCHEMA_VERSION,
@@ -52,6 +51,7 @@ from prefect.client.schemas import (
 from prefect.client.utilities import inject_client
 from prefect.events import emit_event
 from prefect.logging.loggers import disable_logger
+from prefect.plugins import load_prefect_collections
 from prefect.types import SecretDict
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.collections import listrepr, remove_nested_keys, visit_collection
@@ -86,7 +86,7 @@ class InvalidBlockRegistration(Exception):
     """
 
 
-def _collect_nested_reference_strings(obj: Dict):
+def _collect_nested_reference_strings(obj: Dict) -> List[str]:
     """
     Collects all nested reference strings (e.g. #/definitions/Model) from a given object.
     """
@@ -739,9 +739,10 @@ class Block(BaseModel, ABC):
         """
         Retrieve the block class implementation given a key.
         """
+
         # Ensure collections are imported and have the opportunity to register types
-        # before looking up the block class
-        prefect.plugins.load_prefect_collections()
+        # before looking up the block class, but only do this once
+        load_prefect_collections()
 
         return lookup_type(cls, key)
 

prefect/blocks/notifications.py

@@ -10,6 +10,7 @@ from prefect.logging import LogEavesdropper
 from prefect.types import SecretDict
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.templating import apply_values, find_placeholders
+from prefect.utilities.urls import validate_restricted_url
 
 PREFECT_NOTIFY_TYPE_DEFAULT = "prefect_default"
 
@@ -80,6 +81,26 @@ class AppriseNotificationBlock(AbstractAppriseNotificationBlock, ABC):
         description="Incoming webhook URL used to send notifications.",
         examples=["https://hooks.example.com/XXX"],
     )
+    allow_private_urls: bool = Field(
+        default=True,
+        description="Whether to allow notifications to private URLs. Defaults to True.",
+    )
+
+    @sync_compatible
+    async def notify(
+        self,
+        body: str,
+        subject: Optional[str] = None,
+    ):
+        if not self.allow_private_urls:
+            try:
+                validate_restricted_url(self.url.get_secret_value())
+            except ValueError as exc:
+                if self._raise_on_failure:
+                    raise NotificationError(str(exc))
+                raise
+
+        await super().notify(body, subject)
 
 
 # TODO: Move to prefect-slack once collection block auto-registration is

prefect/blocks/webhook.py

@@ -6,10 +6,12 @@ from typing_extensions import Literal
 
 from prefect.blocks.core import Block
 from prefect.types import SecretDict
+from prefect.utilities.urls import validate_restricted_url
 
 # Use a global HTTP transport to maintain a process-wide connection pool for
 # interservice requests
 _http_transport = AsyncHTTPTransport()
+_insecure_http_transport = AsyncHTTPTransport(verify=False)
 
 
 class Webhook(Block):
@@ -39,9 +41,20 @@ class Webhook(Block):
         title="Webhook Headers",
         description="A dictionary of headers to send with the webhook request.",
     )
+    allow_private_urls: bool = Field(
+        default=True,
+        description="Whether to allow notifications to private URLs. Defaults to True.",
+    )
+    verify: bool = Field(
+        default=True,
+        description="Whether or not to enforce a secure connection to the webhook.",
+    )
 
     def block_initialization(self):
-        self._client = AsyncClient(transport=_http_transport)
+        if self.verify:
+            self._client = AsyncClient(transport=_http_transport)
+        else:
+            self._client = AsyncClient(transport=_insecure_http_transport)
 
     async def call(self, payload: Optional[dict] = None) -> Response:
         """
@@ -50,6 +63,9 @@ class Webhook(Block):
         Args:
             payload: an optional payload to send when calling the webhook.
         """
+        if not self.allow_private_urls:
+            validate_restricted_url(self.url.get_secret_value())
+
         async with self._client:
             return await self._client.request(
                 method=self.method,

prefect/cache_policies.py

@@ -1,10 +1,19 @@
 import inspect
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, Optional
+from copy import deepcopy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
+
+from typing_extensions import Self
 
 from prefect.context import TaskRunContext
 from prefect.utilities.hashing import hash_objects
 
+if TYPE_CHECKING:
+    from prefect.filesystems import WritableFileSystem
+    from prefect.locking.protocol import LockManager
+    from prefect.transactions import IsolationLevel
+
 
 @dataclass
 class CachePolicy:
@@ -12,6 +21,14 @@ class CachePolicy:
     Base class for all cache policies.
     """
 
+    key_storage: Union["WritableFileSystem", str, Path, None] = None
+    isolation_level: Union[
+        Literal["READ_COMMITTED", "SERIALIZABLE"],
+        "IsolationLevel",
+        None,
+    ] = None
+    lock_manager: Optional["LockManager"] = None
+
     @classmethod
     def from_cache_key_fn(
         cls, cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
@@ -21,6 +38,37 @@ class CachePolicy:
         """
         return CacheKeyFnPolicy(cache_key_fn=cache_key_fn)
 
+    def configure(
+        self,
+        key_storage: Union["WritableFileSystem", str, Path, None] = None,
+        lock_manager: Optional["LockManager"] = None,
+        isolation_level: Union[
+            Literal["READ_COMMITTED", "SERIALIZABLE"], "IsolationLevel", None
+        ] = None,
+    ) -> Self:
+        """
+        Configure the cache policy with the given key storage, lock manager, and isolation level.
+
+        Args:
+            key_storage: The storage to use for cache keys. If not provided,
+                the current key storage will be used.
+            lock_manager: The lock manager to use for the cache policy. If not provided,
+                the current lock manager will be used.
+            isolation_level: The isolation level to use for the cache policy. If not provided,
+                the current isolation level will be used.
+
+        Returns:
+            A new cache policy with the given key storage, lock manager, and isolation level.
+        """
+        new = deepcopy(self)
+        if key_storage is not None:
+            new.key_storage = key_storage
+        if lock_manager is not None:
+            new.lock_manager = lock_manager
+        if isolation_level is not None:
+            new.isolation_level = isolation_level
+        return new
+
     def compute_key(
         self,
         task_ctx: TaskRunContext,
@@ -30,35 +78,48 @@ class CachePolicy:
     ) -> Optional[str]:
         raise NotImplementedError
 
-    def __sub__(self, other: str) -> "CompoundCachePolicy":
+    def __sub__(self, other: str) -> "CachePolicy":
         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":
+        new = Inputs(exclude=[other])
+        return CompoundCachePolicy(policies=[self, new])
+
+    def __add__(self, other: "CachePolicy") -> "CachePolicy":
         # 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])
+        if (
+            other.key_storage is not None
+            and self.key_storage is not None
+            and other.key_storage != self.key_storage
+        ):
+            raise ValueError(
+                "Cannot add CachePolicies with different storage locations."
+            )
+        if (
+            other.isolation_level is not None
+            and self.isolation_level is not None
+            and other.isolation_level != self.isolation_level
+        ):
+            raise ValueError(
+                "Cannot add CachePolicies with different isolation levels."
+            )
+        if (
+            other.lock_manager is not None
+            and self.lock_manager is not None
+            and other.lock_manager != self.lock_manager
+        ):
+            raise ValueError(
+                "Cannot add CachePolicies with different lock implementations."
+            )
+
+        return CompoundCachePolicy(
+            policies=[self, other],
+            key_storage=self.key_storage or other.key_storage,
+            isolation_level=self.isolation_level or other.isolation_level,
+            lock_manager=self.lock_manager or other.lock_manager,
+        )
 
 
 @dataclass
@@ -93,7 +154,7 @@ class CompoundCachePolicy(CachePolicy):
     Any keys that return `None` will be ignored.
     """
 
-    policies: Optional[list] = None
+    policies: List[CachePolicy] = field(default_factory=list)
 
     def compute_key(
         self,
@@ -103,7 +164,7 @@ class CompoundCachePolicy(CachePolicy):
         **kwargs,
     ) -> Optional[str]:
         keys = []
-        for policy in self.policies or []:
+        for policy in self.policies:
             policy_key = policy.compute_key(
                 task_ctx=task_ctx,
                 inputs=inputs,
@@ -133,6 +194,10 @@ class _None(CachePolicy):
     ) -> Optional[str]:
         return None
 
+    def __add__(self, other: "CachePolicy") -> "CachePolicy":
+        # adding _None is a no-op
+        return other
+
 
 @dataclass
 class TaskSource(CachePolicy):
@@ -208,7 +273,7 @@ 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
+    exclude: List[str] = field(default_factory=list)
 
     def compute_key(
         self,
@@ -230,6 +295,11 @@ class Inputs(CachePolicy):
 
         return hash_objects(hashed_inputs)
 
+    def __sub__(self, other: str) -> "CachePolicy":
+        if not isinstance(other, str):
+            raise TypeError("Can only subtract strings from key policies.")
+        return Inputs(exclude=self.exclude + [other])
+
 
 INPUTS = Inputs()
 NONE = _None()

prefect/client/orchestration.py

@@ -86,6 +86,7 @@ from prefect.client.schemas.objects import (
     BlockSchema,
     BlockType,
     ConcurrencyLimit,
+    ConcurrencyOptions,
     Constant,
     DeploymentSchedule,
     Flow,
@@ -94,7 +95,6 @@ from prefect.client.schemas.objects import (
     FlowRunPolicy,
     Log,
     Parameter,
-    QueueFilter,
     TaskRunPolicy,
     TaskRunResult,
     Variable,
@@ -994,7 +994,6 @@ class PrefectClient:
     async def create_work_queue(
         self,
         name: str,
-        tags: Optional[List[str]] = None,
         description: Optional[str] = None,
         is_paused: Optional[bool] = None,
         concurrency_limit: Optional[int] = None,
@@ -1006,8 +1005,6 @@ class PrefectClient:
 
         Args:
             name: a unique name for the work queue
-            tags: DEPRECATED: an optional list of tags to filter on; only work scheduled with these tags
-                will be included in the queue. This option will be removed on 2023-02-23.
             description: An optional description for the work queue.
             is_paused: Whether or not the work queue is paused.
             concurrency_limit: An optional concurrency limit for the work queue.
@@ -1021,18 +1018,7 @@ class PrefectClient:
         Returns:
             The created work queue
         """
-        if tags:
-            warnings.warn(
-                (
-                    "The use of tags for creating work queue filters is deprecated."
-                    " This option will be removed on 2023-02-23."
-                ),
-                DeprecationWarning,
-            )
-            filter = QueueFilter(tags=tags)
-        else:
-            filter = None
-        create_model = WorkQueueCreate(name=name, filter=filter)
+        create_model = WorkQueueCreate(name=name, filter=None)
         if description is not None:
             create_model.description = description
         if is_paused is not None:
@@ -1654,6 +1640,7 @@ class PrefectClient:
         version: Optional[str] = None,
         schedules: Optional[List[DeploymentScheduleCreate]] = None,
         concurrency_limit: Optional[int] = None,
+        concurrency_options: Optional[ConcurrencyOptions] = None,
         parameters: Optional[Dict[str, Any]] = None,
         description: Optional[str] = None,
         work_queue_name: Optional[str] = None,
@@ -1712,6 +1699,7 @@ class PrefectClient:
             paused=paused,
             schedules=schedules or [],
             concurrency_limit=concurrency_limit,
+            concurrency_options=concurrency_options,
             pull_steps=pull_steps,
             enforce_parameter_schema=enforce_parameter_schema,
         )
@@ -2158,7 +2146,10 @@ class PrefectClient:
         try:
             response = await self._client.post(
                 f"/flow_runs/{flow_run_id}/set_state",
-                json=dict(state=state_create.model_dump(mode="json"), force=force),
+                json=dict(
+                    state=state_create.model_dump(mode="json", serialize_as_any=True),
+                    force=force,
+                ),
             )
         except httpx.HTTPStatusError as e:
             if e.response.status_code == status.HTTP_404_NOT_FOUND:
@@ -3055,7 +3046,11 @@ class PrefectClient:
         return response.json()
 
     async def increment_concurrency_slots(
-        self, names: List[str], slots: int, mode: str, create_if_missing: Optional[bool]
+        self,
+        names: List[str],
+        slots: int,
+        mode: str,
+        create_if_missing: Optional[bool] = None,
     ) -> httpx.Response:
         return await self._client.post(
             "/v2/concurrency_limits/increment",
@@ -3063,7 +3058,7 @@ class PrefectClient:
                 "names": names,
                 "slots": slots,
                 "mode": mode,
-                "create_if_missing": create_if_missing,
+                "create_if_missing": create_if_missing if create_if_missing else False,
             },
         )
 
@@ -3140,6 +3135,30 @@ class PrefectClient:
             else:
                 raise
 
+    async def upsert_global_concurrency_limit_by_name(self, name: str, limit: int):
+        """Creates a global concurrency limit with the given name and limit if one does not already exist.
+
+        If one does already exist matching the name then update it's limit if it is different.
+
+        Note: This is not done atomically.
+        """
+        try:
+            existing_limit = await self.read_global_concurrency_limit_by_name(name)
+        except prefect.exceptions.ObjectNotFound:
+            existing_limit = None
+
+        if not existing_limit:
+            await self.create_global_concurrency_limit(
+                GlobalConcurrencyLimitCreate(
+                    name=name,
+                    limit=limit,
+                )
+            )
+        elif existing_limit.limit != limit:
+            await self.update_global_concurrency_limit(
+                name, GlobalConcurrencyLimitUpdate(limit=limit)
+            )
+
     async def read_global_concurrency_limits(
         self, limit: int = 10, offset: int = 0
     ) -> List[GlobalConcurrencyLimitResponse]:
@@ -3934,7 +3953,10 @@ class SyncPrefectClient:
         try:
             response = self._client.post(
                 f"/flow_runs/{flow_run_id}/set_state",
-                json=dict(state=state_create.model_dump(mode="json"), force=force),
+                json=dict(
+                    state=state_create.model_dump(mode="json", serialize_as_any=True),
+                    force=force,
+                ),
             )
         except httpx.HTTPStatusError as e:
             if e.response.status_code == status.HTTP_404_NOT_FOUND:

prefect/client/schemas/actions.py

@@ -38,7 +38,7 @@ from prefect.utilities.collections import listrepr
 from prefect.utilities.pydantic import get_class_fields_only
 
 if TYPE_CHECKING:
-    from prefect.results import BaseResult
+    from prefect.results import BaseResult, ResultRecordMetadata
 
 R = TypeVar("R")
 
@@ -50,7 +50,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[R]", Any] = Field(
+    data: Union["BaseResult[R]", "ResultRecordMetadata", Any] = Field(
         default=None,
     )
 
@@ -161,6 +161,10 @@ class DeploymentCreate(ActionBaseModel):
         default=None,
         description="The concurrency limit for the deployment.",
     )
+    concurrency_options: Optional[objects.ConcurrencyOptions] = Field(
+        default=None,
+        description="The concurrency options for the deployment.",
+    )
     enforce_parameter_schema: Optional[bool] = Field(
         default=None,
         description=(
@@ -237,6 +241,10 @@ class DeploymentUpdate(ActionBaseModel):
         default=None,
         description="The concurrency limit for the deployment.",
     )
+    concurrency_options: Optional[objects.ConcurrencyOptions] = Field(
+        default=None,
+        description="The concurrency options for the deployment.",
+    )
     tags: List[str] = Field(default_factory=list)
     work_queue_name: Optional[str] = Field(None)
     work_pool_name: Optional[str] = Field(

prefect/client/schemas/filters.py

@@ -506,7 +506,7 @@ class DeploymentFilterTags(PrefectBaseModel, OperatorMixin):
 
 
 class DeploymentFilterConcurrencyLimit(PrefectBaseModel):
-    """Filter by `Deployment.concurrency_limit`."""
+    """DEPRECATED: Prefer `Deployment.concurrency_limit_id` over `Deployment.concurrency_limit`."""
 
     ge_: Optional[int] = Field(
         default=None,
@@ -538,7 +538,9 @@ class DeploymentFilter(PrefectBaseModel, OperatorMixin):
         default=None, description="Filter criteria for `Deployment.work_queue_name`"
     )
     concurrency_limit: Optional[DeploymentFilterConcurrencyLimit] = Field(
-        default=None, description="Filter criteria for `Deployment.concurrency_limit`"
+        default=None,
+        description="DEPRECATED: Prefer `Deployment.concurrency_limit_id` over `Deployment.concurrency_limit`. If provided, will be ignored for backwards-compatibility. Will be removed after December 2024.",
+        deprecated=True,
     )
 
 

prefect/client/schemas/objects.py

@@ -3,6 +3,7 @@ import warnings
 from functools import partial
 from typing import (
     TYPE_CHECKING,
+    Annotated,
     Any,
     Dict,
     Generic,
@@ -17,10 +18,12 @@ import orjson
 import pendulum
 from pydantic import (
     ConfigDict,
+    Discriminator,
     Field,
     HttpUrl,
     IPvAnyNetwork,
     SerializationInfo,
+    Tag,
     field_validator,
     model_serializer,
     model_validator,
@@ -59,7 +62,7 @@ from prefect.utilities.names import generate_slug
 from prefect.utilities.pydantic import handle_secret_render
 
 if TYPE_CHECKING:
-    from prefect.results import BaseResult
+    from prefect.results import BaseResult, ResultRecordMetadata
 
 
 R = TypeVar("R", default=Any)
@@ -138,6 +141,30 @@ class WorkQueueStatus(AutoEnum):
     PAUSED = AutoEnum.auto()
 
 
+class ConcurrencyLimitStrategy(AutoEnum):
+    """Enumeration of concurrency limit strategies."""
+
+    ENQUEUE = AutoEnum.auto()
+    CANCEL_NEW = AutoEnum.auto()
+
+
+class ConcurrencyOptions(PrefectBaseModel):
+    """
+    Class for storing the concurrency config in database.
+    """
+
+    collision_strategy: ConcurrencyLimitStrategy
+
+
+class ConcurrencyLimitConfig(PrefectBaseModel):
+    """
+    Class for storing the concurrency limit config in database.
+    """
+
+    limit: int
+    collision_strategy: ConcurrencyLimitStrategy = ConcurrencyLimitStrategy.ENQUEUE
+
+
 class StateDetails(PrefectBaseModel):
     flow_run_id: Optional[UUID] = None
     task_run_id: Optional[UUID] = None
@@ -158,6 +185,14 @@ class StateDetails(PrefectBaseModel):
     task_parameters_id: Optional[UUID] = None
 
 
+def data_discriminator(x: Any) -> str:
+    if isinstance(x, dict) and "type" in x:
+        return "BaseResult"
+    elif isinstance(x, dict) and "storage_key" in x:
+        return "ResultRecordMetadata"
+    return "Any"
+
+
 class State(ObjectBaseModel, Generic[R]):
     """
     The state of a run.
@@ -168,9 +203,14 @@ class State(ObjectBaseModel, Generic[R]):
     timestamp: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
     message: Optional[str] = Field(default=None, examples=["Run started"])
     state_details: StateDetails = Field(default_factory=StateDetails)
-    data: Union["BaseResult[R]", Any] = Field(
-        default=None,
-    )
+    data: Annotated[
+        Union[
+            Annotated["BaseResult[R]", Tag("BaseResult")],
+            Annotated["ResultRecordMetadata", Tag("ResultRecordMetadata")],
+            Annotated[Any, Tag("Any")],
+        ],
+        Discriminator(data_discriminator),
+    ] = Field(default=None)
 
     @overload
     def result(self: "State[R]", raise_on_failure: bool = True) -> R:
@@ -276,10 +316,12 @@ class State(ObjectBaseModel, Generic[R]):
         results should be sent to the API. Other data is only available locally.
         """
         from prefect.client.schemas.actions import StateCreate
-        from prefect.results import BaseResult
+        from prefect.results import BaseResult, ResultRecord, should_persist_result
 
-        if isinstance(self.data, BaseResult) and self.data.serialize_to_none is False:
+        if isinstance(self.data, BaseResult):
             data = self.data
+        elif isinstance(self.data, ResultRecord) and should_persist_result():
+            data = self.data.metadata
         else:
             data = None
 

prefect/client/schemas/responses.py

@@ -314,11 +314,25 @@ class DeploymentResponse(ObjectBaseModel):
         default=..., description="The flow id associated with the deployment."
     )
     concurrency_limit: Optional[int] = Field(
-        default=None, description="The concurrency limit for the deployment."
+        default=None,
+        description="DEPRECATED: Prefer `global_concurrency_limit`. Will always be None for backwards compatibility. Will be removed after December 2024.",
+        deprecated=True,
+    )
+    global_concurrency_limit: Optional["GlobalConcurrencyLimitResponse"] = Field(
+        default=None,
+        description="The global concurrency limit object for enforcing the maximum number of flow runs that can be active at once.",
+    )
+    concurrency_options: Optional[objects.ConcurrencyOptions] = Field(
+        default=None,
+        description="The concurrency options for the deployment.",
     )
     paused: bool = Field(
         default=False, description="Whether or not the deployment is paused."
     )
+    concurrency_options: Optional[objects.ConcurrencyOptions] = Field(
+        default=None,
+        description="The concurrency options for the deployment.",
+    )
     schedules: List[objects.DeploymentSchedule] = Field(
         default_factory=list, description="A list of schedules for the deployment."
     )

prefect/client/types/flexible_schedule_list.py

@@ -7,5 +7,5 @@ if TYPE_CHECKING:
     from prefect.client.schemas.schedules import SCHEDULE_TYPES
 
 FlexibleScheduleList: TypeAlias = Sequence[
-    Union[DeploymentScheduleCreate, dict[str, Any], "SCHEDULE_TYPES"]
+    Union["DeploymentScheduleCreate", dict[str, Any], "SCHEDULE_TYPES"]
 ]