prefect-client 3.1.10__py3-none-any.whl → 3.1.12__py3-none-any.whl

prefect/_experimental/lineage.py

@@ -86,18 +86,18 @@ async def emit_lineage_event(
     else:
         upstream_resources.extend(related_resources)
 
+    # We want to consider all resources upstream and downstream of the event as
+    # lineage-related, including flows, flow runs, etc., so we add the label to
+    # all resources involved in the event.
+    for res in upstream_resources + downstream_resources:
+        if "prefect.resource.lineage-group" not in res:
+            res["prefect.resource.lineage-group"] = "global"
+
     # Emit an event for each downstream resource. This is necessary because
     # our event schema allows one primary resource and many related resources,
     # and for the purposes of lineage, related resources can only represent
     # upstream resources.
     for resource in downstream_resources:
-        # Downstream lineage resources need to have the
-        # prefect.resource.lineage-group label. All upstram resources from a
-        # downstream resource with this label will be considered lineage-related
-        # resources.
-        if "prefect.resource.lineage-group" not in resource:
-            resource["prefect.resource.lineage-group"] = "global"
-
         emit_kwargs: Dict[str, Any] = {
             "event": event_name,
             "resource": resource,
@@ -170,7 +170,6 @@ async def emit_result_write_event(
             {
                 "prefect.resource.id": result_resource_uri,
                 "prefect.resource.role": "result",
-                "prefect.resource.lineage-group": "global",
             }
         ]
         await emit_lineage_event(

prefect/_experimental/sla/client.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from prefect.client.orchestration.base import BaseAsyncClient, BaseClient
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from prefect._experimental.sla.objects import SlaTypes
+
+
+class SlaClient(BaseClient):
+    def create_sla(self, sla: "SlaTypes") -> "UUID":
+        """
+        Creates a service level agreement.
+        Args:
+            sla: The SLA to create. Must have a deployment ID set.
+        Raises:
+            httpx.RequestError: if the SLA was not created for any reason
+        Returns:
+            the ID of the SLA in the backend
+        """
+        if not sla.owner_resource:
+            raise ValueError(
+                "Deployment ID is not set. Please set using `set_deployment_id`."
+            )
+
+        response = self.request(
+            "POST",
+            "/slas/",
+            json=sla.model_dump(mode="json", exclude_unset=True),
+        )
+        response.raise_for_status()
+
+        from uuid import UUID
+
+        return UUID(response.json().get("id"))
+
+
+class SlaAsyncClient(BaseAsyncClient):
+    async def create_sla(self, sla: "SlaTypes") -> "UUID":
+        """
+        Creates a service level agreement.
+        Args:
+            sla: The SLA to create. Must have a deployment ID set.
+        Raises:
+            httpx.RequestError: if the SLA was not created for any reason
+        Returns:
+            the ID of the SLA in the backend
+        """
+        if not sla.owner_resource:
+            raise ValueError(
+                "Deployment ID is not set. Please set using `set_deployment_id`."
+            )
+
+        response = await self.request(
+            "POST",
+            "/slas/",
+            json=sla.model_dump(mode="json", exclude_unset=True),
+        )
+        response.raise_for_status()
+
+        from uuid import UUID
+
+        return UUID(response.json().get("id"))

prefect/_experimental/sla/objects.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import abc
+from typing import Literal, Optional, Union
+from uuid import UUID
+
+from pydantic import Field, PrivateAttr, computed_field
+from typing_extensions import TypeAlias
+
+from prefect._internal.schemas.bases import PrefectBaseModel
+
+
+class ServiceLevelAgreement(PrefectBaseModel, abc.ABC):
+    """An ORM representation of a Service Level Agreement."""
+
+    _deployment_id: Optional[UUID] = PrivateAttr(default=None)
+
+    name: str = Field(
+        default=...,
+        description="The name of the SLA. Names must be unique on a per-deployment basis.",
+    )
+    severity: Literal["minor", "low", "moderate", "high", "critical"] = Field(
+        default="moderate",
+        description="The severity of the SLA.",
+    )
+    enabled: Optional[bool] = Field(
+        default=True,
+        description="Whether the SLA is enabled.",
+    )
+
+    def set_deployment_id(self, deployment_id: UUID):
+        self._deployment_id = deployment_id
+        return self
+
+    @computed_field
+    @property
+    def owner_resource(self) -> Union[str, None]:
+        if self._deployment_id:
+            return f"prefect.deployment.{self._deployment_id}"
+        return None
+
+
+class TimeToCompletionSla(ServiceLevelAgreement):
+    """An SLA that triggers when a flow run takes longer than the specified duration."""
+
+    duration: int = Field(
+        default=...,
+        description="The maximum flow run duration allowed before the SLA is violated, expressed in seconds.",
+    )
+
+
+# Concrete SLA types
+SlaTypes: TypeAlias = Union[TimeToCompletionSla]

prefect/_internal/_logging.py

@@ -1,4 +1,14 @@
 import logging
+import sys
+
+from typing_extensions import Self
+
+if sys.version_info < (3, 11):
+
+    def getLevelNamesMapping() -> dict[str, int]:
+        return getattr(logging, "_nameToLevel").copy()
+else:
+    getLevelNamesMapping = logging.getLevelNamesMapping  # novermin
 
 
 class SafeLogger(logging.Logger):
@@ -11,11 +21,13 @@ class SafeLogger(logging.Logger):
         # deadlocks during complex concurrency handling
         from prefect.settings import PREFECT_LOGGING_INTERNAL_LEVEL
 
-        return level >= logging._nameToLevel[PREFECT_LOGGING_INTERNAL_LEVEL.value()]
+        internal_level = getLevelNamesMapping()[PREFECT_LOGGING_INTERNAL_LEVEL.value()]
+
+        return level >= internal_level
 
-    def getChild(self, suffix: str):
+    def getChild(self, suffix: str) -> Self:
         logger = super().getChild(suffix)
-        logger.__class__ = SafeLogger
+        logger.__class__ = self.__class__
         return logger
 
 

prefect/_internal/compatibility/async_dispatch.py

@@ -1,11 +1,13 @@
 import asyncio
 import inspect
+from collections.abc import Coroutine
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union, cast
 
 from typing_extensions import ParamSpec
 
 if TYPE_CHECKING:
+    from prefect.flows import Flow
     from prefect.tasks import Task
 
 R = TypeVar("R")
@@ -41,7 +43,9 @@ def is_in_async_context() -> bool:
 
 
 def _is_acceptable_callable(
-    obj: Union[Callable[P, R], "Task[P, R]", classmethod],
+    obj: Union[
+        Callable[P, R], "Flow[P, R]", "Task[P, R]", "classmethod[type[Any], P, R]"
+    ],
 ) -> bool:
     if inspect.iscoroutinefunction(obj):
         return True
@@ -58,7 +62,10 @@ def _is_acceptable_callable(
 
 
 def async_dispatch(
-    async_impl: Callable[P, Coroutine[Any, Any, R]],
+    async_impl: Union[
+        Callable[P, Coroutine[Any, Any, R]],
+        "classmethod[type[Any], P, Coroutine[Any, Any, R]]",
+    ],
 ) -> Callable[[Callable[P, R]], Callable[P, Union[R, Coroutine[Any, Any, R]]]]:
     """
     Decorator that dispatches to either sync or async implementation based on context.
@@ -66,27 +73,26 @@ def async_dispatch(
     Args:
         async_impl: The async implementation to dispatch to when in async context
     """
+    if not _is_acceptable_callable(async_impl):
+        raise TypeError("async_impl must be an async function")
+    if isinstance(async_impl, classmethod):
+        async_impl = cast(Callable[P, Coroutine[Any, Any, R]], async_impl.__func__)
 
     def decorator(
         sync_fn: Callable[P, R],
     ) -> Callable[P, Union[R, Coroutine[Any, Any, R]]]:
-        if not _is_acceptable_callable(async_impl):
-            raise TypeError("async_impl must be an async function")
-
         @wraps(sync_fn)
         def wrapper(
             *args: P.args,
-            _sync: Optional[bool] = None,  # type: ignore
             **kwargs: P.kwargs,
         ) -> Union[R, Coroutine[Any, Any, R]]:
-            should_run_sync = _sync if _sync is not None else not is_in_async_context()
-
-            if should_run_sync:
-                return sync_fn(*args, **kwargs)
-            if isinstance(async_impl, classmethod):
-                return async_impl.__func__(*args, **kwargs)
-            return async_impl(*args, **kwargs)
-
-        return wrapper  # type: ignore
+            _sync = kwargs.pop("_sync", None)
+            should_run_sync = (
+                bool(_sync) if _sync is not None else not is_in_async_context()
+            )
+            fn = sync_fn if should_run_sync else async_impl
+            return fn(*args, **kwargs)
+
+        return wrapper
 
     return decorator

prefect/_internal/compatibility/deprecated.py

@@ -13,10 +13,11 @@ e.g. Jan 2023.
 import functools
 import sys
 import warnings
-from typing import Any, Callable, List, Optional, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 import pendulum
 from pydantic import BaseModel
+from typing_extensions import ParamSpec, TypeAlias, TypeVar
 
 from prefect.utilities.callables import get_call_parameters
 from prefect.utilities.importtools import (
@@ -25,8 +26,10 @@ from prefect.utilities.importtools import (
     to_qualified_name,
 )
 
-T = TypeVar("T", bound=Callable)
+P = ParamSpec("P")
+R = TypeVar("R", infer_variance=True)
 M = TypeVar("M", bound=BaseModel)
+T = TypeVar("T")
 
 
 DEPRECATED_WARNING = (
@@ -38,7 +41,7 @@ DEPRECATED_MOVED_WARNING = (
     "path after {end_date}. {help}"
 )
 DEPRECATED_DATEFMT = "MMM YYYY"  # e.g. Feb 2023
-DEPRECATED_MODULE_ALIASES: List[AliasedModuleDefinition] = []
+DEPRECATED_MODULE_ALIASES: list[AliasedModuleDefinition] = []
 
 
 class PrefectDeprecationWarning(DeprecationWarning):
@@ -53,7 +56,7 @@ def generate_deprecation_message(
     end_date: Optional[str] = None,
     help: str = "",
     when: str = "",
-):
+) -> str:
     if not start_date and not end_date:
         raise ValueError(
             "A start date is required if an end date is not provided. Suggested start"
@@ -61,6 +64,8 @@ def generate_deprecation_message(
         )
 
     if not end_date:
+        if TYPE_CHECKING:
+            assert start_date is not None
         parsed_start_date = pendulum.from_format(start_date, DEPRECATED_DATEFMT)
         parsed_end_date = parsed_start_date.add(months=6)
         end_date = parsed_end_date.format(DEPRECATED_DATEFMT)
@@ -83,8 +88,8 @@ def deprecated_callable(
     end_date: Optional[str] = None,
     stacklevel: int = 2,
     help: str = "",
-) -> Callable[[T], T]:
-    def decorator(fn: T):
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
         message = generate_deprecation_message(
             name=to_qualified_name(fn),
             start_date=start_date,
@@ -93,7 +98,7 @@ def deprecated_callable(
         )
 
         @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             warnings.warn(message, PrefectDeprecationWarning, stacklevel=stacklevel)
             return fn(*args, **kwargs)
 
@@ -108,8 +113,8 @@ def deprecated_class(
     end_date: Optional[str] = None,
     stacklevel: int = 2,
     help: str = "",
-) -> Callable[[T], T]:
-    def decorator(cls: T):
+) -> Callable[[type[T]], type[T]]:
+    def decorator(cls: type[T]) -> type[T]:
         message = generate_deprecation_message(
             name=to_qualified_name(cls),
             start_date=start_date,
@@ -120,7 +125,7 @@ def deprecated_class(
         original_init = cls.__init__
 
         @functools.wraps(original_init)
-        def new_init(self, *args, **kwargs):
+        def new_init(self: T, *args: Any, **kwargs: Any) -> None:
             warnings.warn(message, PrefectDeprecationWarning, stacklevel=stacklevel)
             original_init(self, *args, **kwargs)
 
@@ -139,7 +144,7 @@ def deprecated_parameter(
     help: str = "",
     when: Optional[Callable[[Any], bool]] = None,
     when_message: str = "",
-) -> Callable[[T], T]:
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """
     Mark a parameter in a callable as deprecated.
 
@@ -155,7 +160,7 @@ def deprecated_parameter(
 
     when = when or (lambda _: True)
 
-    def decorator(fn: T):
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
         message = generate_deprecation_message(
             name=f"The parameter {name!r} for {fn.__name__!r}",
             start_date=start_date,
@@ -165,7 +170,7 @@ def deprecated_parameter(
         )
 
         @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             try:
                 parameters = get_call_parameters(fn, args, kwargs, apply_defaults=False)
             except Exception:
@@ -182,6 +187,10 @@ def deprecated_parameter(
     return decorator
 
 
+JsonValue: TypeAlias = Union[int, float, str, bool, None, list["JsonValue"], "JsonDict"]
+JsonDict: TypeAlias = dict[str, JsonValue]
+
+
 def deprecated_field(
     name: str,
     *,
@@ -191,7 +200,7 @@ def deprecated_field(
     help: str = "",
     when: Optional[Callable[[Any], bool]] = None,
     stacklevel: int = 2,
-):
+) -> Callable[[type[M]], type[M]]:
     """
     Mark a field in a Pydantic model as deprecated.
 
@@ -212,7 +221,7 @@ def deprecated_field(
 
     # Replaces the model's __init__ method with one that performs an additional warning
     # check
-    def decorator(model_cls: Type[M]) -> Type[M]:
+    def decorator(model_cls: type[M]) -> type[M]:
         message = generate_deprecation_message(
             name=f"The field {name!r} in {model_cls.__name__!r}",
             start_date=start_date,
@@ -224,7 +233,7 @@ def deprecated_field(
         cls_init = model_cls.__init__
 
         @functools.wraps(model_cls.__init__)
-        def __init__(__pydantic_self__, **data: Any) -> None:
+        def __init__(__pydantic_self__: M, **data: Any) -> None:
             if name in data.keys() and when(data[name]):
                 warnings.warn(message, PrefectDeprecationWarning, stacklevel=stacklevel)
 
@@ -232,8 +241,23 @@ def deprecated_field(
 
             field = __pydantic_self__.model_fields.get(name)
             if field is not None:
-                field.json_schema_extra = field.json_schema_extra or {}
-                field.json_schema_extra["deprecated"] = True
+                json_schema_extra = field.json_schema_extra or {}
+
+                if not isinstance(json_schema_extra, dict):
+                    # json_schema_extra is a hook function; wrap it to add the deprecated flag.
+                    extra_func = json_schema_extra
+
+                    @functools.wraps(extra_func)
+                    def wrapped(__json_schema: JsonDict) -> None:
+                        extra_func(__json_schema)
+                        __json_schema["deprecated"] = True
+
+                    json_schema_extra = wrapped
+
+                else:
+                    json_schema_extra["deprecated"] = True
+
+                field.json_schema_extra = json_schema_extra
 
         # Patch the model's init method
         model_cls.__init__ = __init__

prefect/_internal/compatibility/migration.py

@@ -43,7 +43,7 @@ Removed objects:
 """
 
 import sys
-from typing import Any, Callable, Dict
+from typing import Any, Callable
 
 from pydantic_core import PydanticCustomError
 
@@ -157,7 +157,7 @@ def getattr_migration(module_name: str) -> Callable[[str], Any]:
                 f"`{import_path}` has been removed. {error_message}"
             )
 
-        globals: Dict[str, Any] = sys.modules[module_name].__dict__
+        globals: dict[str, Any] = sys.modules[module_name].__dict__
         if name in globals:
             return globals[name]
 

prefect/_internal/concurrency/inspection.py

@@ -7,7 +7,6 @@ import linecache
 import sys
 import threading
 from types import FrameType
-from typing import List, Optional
 
 """
 The following functions are derived from dask/distributed which is licensed under the
@@ -72,26 +71,25 @@ def repr_frame(frame: FrameType) -> str:
     return text + "\n\t" + line
 
 
-def call_stack(frame: FrameType) -> List[str]:
+def call_stack(frame: FrameType) -> list[str]:
     """Create a call text stack from a frame"""
-    L = []
-    cur_frame: Optional[FrameType] = frame
+    frames: list[str] = []
+    cur_frame = frame
     while cur_frame:
-        L.append(repr_frame(cur_frame))
+        frames.append(repr_frame(cur_frame))
         cur_frame = cur_frame.f_back
-    return L[::-1]
+    return frames[::-1]
 
 
-def stack_for_threads(*threads: threading.Thread) -> List[str]:
-    frames = sys._current_frames()
+def stack_for_threads(*threads: threading.Thread) -> list[str]:
+    frames = sys._current_frames()  # pyright: ignore[reportPrivateUsage]
     try:
-        lines = []
+        lines: list[str] = []
         for thread in threads:
-            lines.append(
-                f"------ Call stack of {thread.name} ({hex(thread.ident)}) -----"
-            )
-            thread_frames = frames.get(thread.ident)
-            if thread_frames:
+            ident = thread.ident
+            hex_ident = hex(ident) if ident is not None else "<unknown>"
+            lines.append(f"------ Call stack of {thread.name} ({hex_ident}) -----")
+            if ident is not None and (thread_frames := frames.get(ident)):
                 lines.append("".join(call_stack(thread_frames)))
             else:
                 lines.append("No stack frames found")

prefect/_internal/concurrency/primitives.py

@@ -27,7 +27,7 @@ class Event:
     """
 
     def __init__(self) -> None:
-        self._waiters = collections.deque()
+        self._waiters: collections.deque[asyncio.Future[bool]] = collections.deque()
         self._value = False
         self._lock = threading.Lock()
 
@@ -69,7 +69,7 @@ class Event:
             if self._value:
                 return True
 
-            fut = asyncio.get_running_loop().create_future()
+            fut: asyncio.Future[bool] = asyncio.get_running_loop().create_future()
             self._waiters.append(fut)
 
         try: