prefect-client 3.1.7__py3-none-any.whl → 3.1.9__py3-none-any.whl

prefect/__init__.py

@@ -2,29 +2,14 @@
 
 # Setup version and path constants
 
+import sys
 from . import _version
 import importlib
 import pathlib
-from typing import TYPE_CHECKING, Any
-
-__version_info__ = _version.get_versions()
-__version__ = __version_info__["version"]
-
-# The absolute path to this module
-__module_path__ = pathlib.Path(__file__).parent
-# The absolute path to the root of the repository, only valid for use during development
-__development_base_path__ = __module_path__.parents[1]
-
-# The absolute path to the built UI within the Python module, used by
-# `prefect server start` to serve a dynamic build of the UI
-__ui_static_subpath__ = __module_path__ / "server" / "ui_build"
-
-# The absolute path to the built UI within the Python module
-__ui_static_path__ = __module_path__ / "server" / "ui"
-
-del _version, pathlib
+from typing import TYPE_CHECKING, Any, Optional, TypedDict, cast
 
 if TYPE_CHECKING:
+    from importlib.machinery import ModuleSpec
     from .main import (
         allow_failure,
         flow,
@@ -45,80 +30,89 @@ if TYPE_CHECKING:
         suspend_flow_run,
     )
 
-_slots: dict[str, Any] = {
-    "__version_info__": __version_info__,
-    "__version__": __version__,
-    "__module_path__": __module_path__,
-    "__development_base_path__": __development_base_path__,
-    "__ui_static_subpath__": __ui_static_subpath__,
-    "__ui_static_path__": __ui_static_path__,
-}
+    __spec__: ModuleSpec
+
+    # Versioneer provides version information as dictionaries
+    # with these keys
+    class VersionInfo(TypedDict("_FullRevisionId", {"full-revisionid": str})):
+        version: str
+        dirty: Optional[bool]
+        error: Optional[str]
+        date: Optional[str]
 
-_public_api: dict[str, tuple[str, str]] = {
+
+__version_info__: "VersionInfo" = cast("VersionInfo", _version.get_versions())
+__version__ = __version_info__["version"]
+
+# The absolute path to this module
+__module_path__: pathlib.Path = pathlib.Path(__file__).parent
+# The absolute path to the root of the repository, only valid for use during development
+__development_base_path__: pathlib.Path = __module_path__.parents[1]
+
+# The absolute path to the built UI within the Python module, used by
+# `prefect server start` to serve a dynamic build of the UI
+__ui_static_subpath__: pathlib.Path = __module_path__ / "server" / "ui_build"
+
+# The absolute path to the built UI within the Python module
+__ui_static_path__: pathlib.Path = __module_path__ / "server" / "ui"
+
+del _version, pathlib
+
+_public_api: dict[str, tuple[Optional[str], str]] = {
     "allow_failure": (__spec__.parent, ".main"),
+    "aserve": (__spec__.parent, ".main"),
+    "deploy": (__spec__.parent, ".main"),
     "flow": (__spec__.parent, ".main"),
     "Flow": (__spec__.parent, ".main"),
     "get_client": (__spec__.parent, ".main"),
     "get_run_logger": (__spec__.parent, ".main"),
+    "pause_flow_run": (__spec__.parent, ".main"),
+    "resume_flow_run": (__spec__.parent, ".main"),
+    "serve": (__spec__.parent, ".main"),
     "State": (__spec__.parent, ".main"),
+    "suspend_flow_run": (__spec__.parent, ".main"),
     "tags": (__spec__.parent, ".main"),
     "task": (__spec__.parent, ".main"),
     "Task": (__spec__.parent, ".main"),
     "Transaction": (__spec__.parent, ".main"),
     "unmapped": (__spec__.parent, ".main"),
-    "serve": (__spec__.parent, ".main"),
-    "aserve": (__spec__.parent, ".main"),
-    "deploy": (__spec__.parent, ".main"),
-    "pause_flow_run": (__spec__.parent, ".main"),
-    "resume_flow_run": (__spec__.parent, ".main"),
-    "suspend_flow_run": (__spec__.parent, ".main"),
 }
 
 # Declare API for type-checkers
 __all__ = [
+    "__version__",
     "allow_failure",
+    "aserve",
+    "deploy",
     "flow",
     "Flow",
     "get_client",
     "get_run_logger",
+    "pause_flow_run",
+    "resume_flow_run",
+    "serve",
     "State",
+    "suspend_flow_run",
     "tags",
     "task",
     "Task",
     "Transaction",
     "unmapped",
-    "serve",
-    "aserve",
-    "deploy",
-    "pause_flow_run",
-    "resume_flow_run",
-    "suspend_flow_run",
-    "__version_info__",
-    "__version__",
-    "__module_path__",
-    "__development_base_path__",
-    "__ui_static_subpath__",
-    "__ui_static_path__",
 ]
 
 
-def __getattr__(attr_name: str) -> object:
-    if attr_name in _slots:
-        return _slots[attr_name]
+def __getattr__(attr_name: str) -> Any:
     try:
-        dynamic_attr = _public_api.get(attr_name)
-        if dynamic_attr is None:
+        if (dynamic_attr := _public_api.get(attr_name)) is None:
             return importlib.import_module(f".{attr_name}", package=__name__)
 
-        package, module_name = dynamic_attr
-
-        from importlib import import_module
-
-        if module_name == "__module__":
-            return import_module(f".{attr_name}", package=package)
+        package, mname = dynamic_attr
+        if mname == "__module__":
+            return importlib.import_module(f".{attr_name}", package=package)
         else:
-            module = import_module(module_name, package=package)
+            module = importlib.import_module(mname, package=package)
             return getattr(module, attr_name)
     except ModuleNotFoundError as ex:
-        module, _, attribute = ex.name.rpartition(".")
-        raise AttributeError(f"module {module} has no attribute {attribute}") from ex
+        mname, _, attr = (ex.name or "").rpartition(".")
+        ctx = {"name": mname, "obj": attr} if sys.version_info >= (3, 10) else {}
+        raise AttributeError(f"module {mname} has no attribute {attr}", **ctx) from ex

prefect/_internal/concurrency/services.py

@@ -225,7 +225,9 @@ class QueueService(abc.ABC, Generic[T]):
             return future.result()
 
     @classmethod
-    def drain_all(cls, timeout: Optional[float] = None) -> Union[Awaitable, None]:
+    def drain_all(
+        cls, timeout: Optional[float] = None, at_exit=True
+    ) -> Union[Awaitable, None]:
         """
         Stop all instances of the service and wait for all remaining work to be
         completed.
@@ -237,7 +239,7 @@ class QueueService(abc.ABC, Generic[T]):
             instances = tuple(cls._instances.values())
 
             for instance in instances:
-                futures.append(instance._drain())
+                futures.append(instance._drain(at_exit=at_exit))
 
         if get_running_loop() is not None:
             return (
@@ -376,10 +378,10 @@ class BatchedQueueService(QueueService[T]):
 @contextlib.contextmanager
 def drain_on_exit(service: QueueService):
     yield
-    service.drain_all()
+    service.drain_all(at_exit=True)
 
 
 @contextlib.asynccontextmanager
 async def drain_on_exit_async(service: QueueService):
     yield
-    await service.drain_all()
+    await service.drain_all(at_exit=True)

prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-16T10:06:12-0800",
+ "date": "2024-12-20T16:33:15-0500",
  "dirty": true,
  "error": null,
- "full-revisionid": "c05ffa6d1816f980d574b42119dffb9f8c8300a3",
- "version": "3.1.7"
+ "full-revisionid": "e1fe79439b319ce5478aa5f8bd4e836ed4f8a0f5",
+ "version": "3.1.9"
 }
 '''  # END VERSION_JSON
 

prefect/agent.py

@@ -1,6 +1,8 @@
 """
 2024-06-27: This surfaces an actionable error message for moved or removed objects in Prefect 3.0 upgrade.
 """
+from typing import Any, Callable
+
 from prefect._internal.compatibility.migration import getattr_migration
 
-__getattr__ = getattr_migration(__name__)
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)

prefect/client/cloud.py

@@ -1,6 +1,5 @@
 import re
 from typing import Any, NoReturn, Optional, cast
-from uuid import UUID
 
 import anyio
 import httpx
@@ -23,7 +22,6 @@ from prefect.settings import (
     PREFECT_CLOUD_API_URL,
     PREFECT_TESTING_UNIT_TEST_MODE,
 )
-from prefect.types import KeyValueLabels
 
 PARSE_API_URL_REGEX = re.compile(r"accounts/(.{36})/workspaces/(.{36})")
 
@@ -160,25 +158,6 @@ class CloudClient:
         response = await self.get(f"{self.account_base_url}/ip_allowlist/my_access")
         return IPAllowlistMyAccessResponse.model_validate(response)
 
-    async def update_flow_run_labels(
-        self, flow_run_id: UUID, labels: KeyValueLabels
-    ) -> httpx.Response:
-        """
-        Update the labels for a flow run.
-
-        Args:
-            flow_run_id: The identifier for the flow run to update.
-            labels: A dictionary of labels to update for the flow run.
-
-        Returns:
-            an `httpx.Response` object from the PATCH request
-        """
-
-        return await self._client.patch(
-            f"{self.workspace_base_url}/flow_runs/{flow_run_id}/labels",
-            json=labels,
-        )
-
     async def __aenter__(self) -> Self:
         await self._client.__aenter__()
         return self

prefect/client/orchestration.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import datetime
 import ssl
 import warnings
@@ -114,6 +115,7 @@ from prefect.events import filters
 from prefect.events.schemas.automations import Automation, AutomationCore
 from prefect.logging import get_logger
 from prefect.settings import (
+    PREFECT_API_AUTH_STRING,
     PREFECT_API_DATABASE_CONNECTION_URL,
     PREFECT_API_ENABLE_HTTP2,
     PREFECT_API_KEY,
@@ -228,6 +230,7 @@ def get_client(
     if sync_client:
         return SyncPrefectClient(
             api,
+            auth_string=PREFECT_API_AUTH_STRING.value(),
             api_key=PREFECT_API_KEY.value(),
             httpx_settings=httpx_settings,
             server_type=server_type,
@@ -235,6 +238,7 @@ def get_client(
     else:
         return PrefectClient(
             api,
+            auth_string=PREFECT_API_AUTH_STRING.value(),
             api_key=PREFECT_API_KEY.value(),
             httpx_settings=httpx_settings,
             server_type=server_type,
@@ -271,6 +275,7 @@ class PrefectClient:
         self,
         api: Union[str, ASGIApp],
         *,
+        auth_string: Optional[str] = None,
         api_key: Optional[str] = None,
         api_version: Optional[str] = None,
         httpx_settings: Optional[dict[str, Any]] = None,
@@ -299,6 +304,10 @@ class PrefectClient:
         if api_key:
             httpx_settings["headers"].setdefault("Authorization", f"Bearer {api_key}")
 
+        if auth_string:
+            token = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")
+            httpx_settings["headers"].setdefault("Authorization", f"Basic {token}")
+
         # Context management
         self._context_stack: int = 0
         self._exit_stack = AsyncExitStack()
@@ -3469,6 +3478,8 @@ class PrefectClient:
         try:
             api_version = await self.api_version()
         except Exception as e:
+            if "Unauthorized" in str(e):
+                raise e
             raise RuntimeError(f"Failed to reach API at {self.api_url}") from e
 
         api_version = version.parse(api_version)
@@ -3590,6 +3601,7 @@ class SyncPrefectClient:
         self,
         api: Union[str, ASGIApp],
         *,
+        auth_string: Optional[str] = None,
         api_key: Optional[str] = None,
         api_version: Optional[str] = None,
         httpx_settings: Optional[dict[str, Any]] = None,
@@ -3618,6 +3630,10 @@ class SyncPrefectClient:
         if api_key:
             httpx_settings["headers"].setdefault("Authorization", f"Bearer {api_key}")
 
+        if auth_string:
+            token = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")
+            httpx_settings["headers"].setdefault("Authorization", f"Basic {token}")
+
         # Context management
         self._context_stack: int = 0
         self._ephemeral_app: Optional[ASGIApp] = None
@@ -3800,6 +3816,8 @@ class SyncPrefectClient:
         try:
             api_version = self.api_version()
         except Exception as e:
+            if "Unauthorized" in str(e):
+                raise e
             raise RuntimeError(f"Failed to reach API at {self.api_url}") from e
 
         api_version = version.parse(api_version)

prefect/client/schemas/objects.py

@@ -190,6 +190,8 @@ class StateDetails(PrefectBaseModel):
     retriable: Optional[bool] = None
     transition_id: Optional[UUID] = None
     task_parameters_id: Optional[UUID] = None
+    # Captures the trace_id and span_id of the span where this state was created
+    traceparent: Optional[str] = None
 
 
 def data_discriminator(x: Any) -> str:
@@ -237,6 +239,15 @@ class State(ObjectBaseModel, Generic[R]):
     ) -> Union[R, Exception]:
         ...
 
+    @overload
+    def result(
+        self: "State[R]",
+        raise_on_failure: bool = ...,
+        fetch: bool = ...,
+        retry_result_failure: bool = ...,
+    ) -> Union[R, Exception]:
+        ...
+
     @deprecated.deprecated_parameter(
         "fetch",
         when=lambda fetch: fetch is not True,

prefect/client/utilities.py

@@ -7,7 +7,7 @@ Utilities for working with clients.
 
 from collections.abc import Awaitable, Coroutine
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 from typing_extensions import Concatenate, ParamSpec, TypeGuard, TypeVar
 
@@ -71,23 +71,9 @@ def client_injector(
     return wrapper
 
 
-@overload
 def inject_client(
     fn: Callable[P, Coroutine[Any, Any, R]],
 ) -> Callable[P, Coroutine[Any, Any, R]]:
-    ...
-
-
-@overload
-def inject_client(
-    fn: Callable[P, R],
-) -> Callable[P, R]:
-    ...
-
-
-def inject_client(
-    fn: Callable[P, Union[Coroutine[Any, Any, R], R]],
-) -> Callable[P, Union[Coroutine[Any, Any, R], R]]:
     """
     Simple helper to provide a context managed client to an asynchronous function.
 

prefect/deployments/deployments.py

@@ -1,3 +1,5 @@
-from .._internal.compatibility.migration import getattr_migration
+from typing import Any, Callable
 
-__getattr__ = getattr_migration(__name__)
+from prefect._internal.compatibility.migration import getattr_migration
+
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)

prefect/deployments/runner.py

@@ -549,7 +549,9 @@ class RunnerDeployment(BaseModel):
                         module = importlib.import_module(mod_name)
                         flow_file = getattr(module, "__file__", None)
                     except ModuleNotFoundError as exc:
-                        if "__prefect_loader__" in str(exc):
+                        # 16458 adds an identifier to the module name, so checking
+                        # for "__prefect_loader__" (2 underscores) will not match
+                        if "__prefect_loader_" in str(exc):
                             raise ValueError(
                                 "Cannot create a RunnerDeployment from a flow that has been"
                                 " loaded from an entrypoint. To deploy a flow via"

prefect/engine.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from typing import Any, Callable
 from uuid import UUID
 
 from prefect._internal.compatibility.migration import getattr_migration
@@ -72,4 +73,4 @@ if __name__ == "__main__":
         # Let the exit code be determined by the base exception type
         raise
 
-__getattr__ = getattr_migration(__name__)
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)

prefect/events/filters.py

@@ -2,7 +2,7 @@ from typing import List, Optional, Tuple, cast
 from uuid import UUID
 
 import pendulum
-from pydantic import Field, PrivateAttr
+from pydantic import Field
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 from prefect.types import DateTime
@@ -41,18 +41,12 @@ class AutomationFilter(PrefectBaseModel):
 class EventDataFilter(PrefectBaseModel, extra="forbid"):  # type: ignore[call-arg]
     """A base class for filtering event data."""
 
-    _top_level_filter: Optional["EventFilter"] = PrivateAttr(None)
-
     def get_filters(self) -> List["EventDataFilter"]:
         filters: List["EventDataFilter"] = [
             filter
-            for filter in [
-                getattr(self, name) for name, field in self.model_fields.items()
-            ]
+            for filter in [getattr(self, name) for name in self.model_fields]
             if isinstance(filter, EventDataFilter)
         ]
-        for filter in filters:
-            filter._top_level_filter = self._top_level_filter
         return filters
 
     def includes(self, event: Event) -> bool:

prefect/exceptions.py

@@ -4,19 +4,20 @@ Prefect-specific exceptions.
 
 import inspect
 import traceback
+from collections.abc import Iterable
 from types import ModuleType, TracebackType
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from httpx._exceptions import HTTPStatusError
 from pydantic import ValidationError
-from rich.traceback import Traceback
 from typing_extensions import Self
 
-import prefect
+if TYPE_CHECKING:
+    from prefect.states import State
 
 
 def _trim_traceback(
-    tb: TracebackType, remove_modules: Iterable[ModuleType]
+    tb: Optional[TracebackType], remove_modules: Iterable[ModuleType]
 ) -> Optional[TracebackType]:
     """
     Utility to remove frames from specific modules from a traceback.
@@ -32,8 +33,9 @@ def _trim_traceback(
         A traceback, or `None` if all traceback frames originate from an excluded module
 
     """
-    strip_paths = [module.__file__ for module in remove_modules]
-
+    strip_paths = [
+        module.__file__ for module in remove_modules if module.__file__ is not None
+    ]
     while tb and any(
         module_path in str(tb.tb_frame.f_globals.get("__file__", ""))
         for module_path in strip_paths
@@ -91,7 +93,9 @@ class PausedRun(PrefectException):
     Raised when the result from a paused run is retrieved.
     """
 
-    def __init__(self, *args, state=None, **kwargs):
+    def __init__(
+        self, *args: Any, state: Optional["State[Any]"] = None, **kwargs: Any
+    ) -> None:
         super().__init__(*args, **kwargs)
         self.state = state
 
@@ -133,6 +137,8 @@ class ScriptError(PrefectException):
         user_exc: Exception,
         path: str,
     ) -> None:
+        import prefect.utilities.importtools
+
         message = f"Script at {str(path)!r} encountered an exception: {user_exc!r}"
         super().__init__(message)
         self.user_exc = user_exc
@@ -144,30 +150,6 @@ class ScriptError(PrefectException):
         )
 
 
-class FlowScriptError(PrefectException):
-    """
-    Raised when a script errors during evaluation while attempting to load a flow.
-    """
-
-    def __init__(
-        self,
-        user_exc: Exception,
-        script_path: str,
-    ) -> None:
-        message = f"Flow script at {script_path!r} encountered an exception"
-        super().__init__(message)
-
-        self.user_exc = user_exc
-
-    def rich_user_traceback(self, **kwargs):
-        trace = Traceback.extract(
-            type(self.user_exc),
-            self.user_exc,
-            self.user_exc.__traceback__.tb_next.tb_next.tb_next.tb_next,
-        )
-        return Traceback(trace, **kwargs)
-
-
 class ParameterTypeError(PrefectException):
     """
     Raised when a parameter does not pass Pydantic type validation.
@@ -196,7 +178,11 @@ class ParameterBindError(TypeError, PrefectException):
 
     @classmethod
     def from_bind_failure(
-        cls, fn: Callable, exc: TypeError, call_args: List, call_kwargs: Dict
+        cls,
+        fn: Callable[..., Any],
+        exc: TypeError,
+        call_args: tuple[Any, ...],
+        call_kwargs: dict[str, Any],
     ) -> Self:
         fn_signature = str(inspect.signature(fn)).strip("()")
 
@@ -214,7 +200,9 @@ class SignatureMismatchError(PrefectException, TypeError):
         super().__init__(msg)
 
     @classmethod
-    def from_bad_params(cls, expected_params: List[str], provided_params: List[str]):
+    def from_bad_params(
+        cls, expected_params: list[str], provided_params: list[str]
+    ) -> Self:
         msg = (
             f"Function expects parameters {expected_params} but was provided with"
             f" parameters {provided_params}"
@@ -231,14 +219,14 @@ class ObjectNotFound(PrefectException):
         self,
         http_exc: Exception,
         help_message: Optional[str] = None,
-        *args,
-        **kwargs,
-    ):
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         self.http_exc = http_exc
         self.help_message = help_message
         super().__init__(help_message, *args, **kwargs)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.help_message or super().__str__()
 
 
@@ -247,7 +235,7 @@ class ObjectAlreadyExists(PrefectException):
     Raised when the client receives a 409 (conflict) from the API.
     """
 
-    def __init__(self, http_exc: Exception, *args, **kwargs):
+    def __init__(self, http_exc: Exception, *args: Any, **kwargs: Any) -> None:
         self.http_exc = http_exc
         super().__init__(*args, **kwargs)
 
@@ -304,7 +292,9 @@ class Pause(PrefectSignal):
     Raised when a flow run is PAUSED and needs to exit for resubmission.
     """
 
-    def __init__(self, *args, state=None, **kwargs):
+    def __init__(
+        self, *args: Any, state: Optional["State[Any]"] = None, **kwargs: Any
+    ) -> None:
         super().__init__(*args, **kwargs)
         self.state = state
 
@@ -332,7 +322,7 @@ class PrefectHTTPStatusError(HTTPStatusError):
     """
 
     @classmethod
-    def from_httpx_error(cls: Type[Self], httpx_error: HTTPStatusError) -> Self:
+    def from_httpx_error(cls: type[Self], httpx_error: HTTPStatusError) -> Self:
         """
         Generate a `PrefectHTTPStatusError` from an `httpx.HTTPStatusError`.
         """
@@ -441,7 +431,7 @@ class ProfileSettingsValidationError(PrefectException):
     Raised when a profile settings are invalid.
     """
 
-    def __init__(self, errors: List[Tuple[Any, ValidationError]]) -> None:
+    def __init__(self, errors: list[tuple[Any, ValidationError]]) -> None:
         self.errors = errors
 
 

prefect/filesystems.py

@@ -2,7 +2,7 @@ import abc
 import urllib.parse
 from pathlib import Path
 from shutil import copytree
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional
 
 import anyio
 import fsspec
@@ -544,4 +544,4 @@ class NullFileSystem(BaseModel):
         pass
 
 
-__getattr__ = getattr_migration(__name__)
+__getattr__: Callable[[str], Any] = getattr_migration(__name__)

prefect/flow_engine.py

@@ -221,7 +221,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
                     return_data=True,
                     max_depth=-1,
                     remove_annotations=True,
-                    context={},
+                    context={"parameter_name": parameter},
                 )
             except UpstreamTaskError:
                 raise
@@ -784,7 +784,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
                     return_data=True,
                     max_depth=-1,
                     remove_annotations=True,
-                    context={},
+                    context={"parameter_name": parameter},
                 )
             except UpstreamTaskError:
                 raise