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

prefect/_internal/schemas/validators.py

@@ -13,13 +13,14 @@ import warnings
 from copy import copy
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
+from uuid import UUID
 
 import jsonschema
 import pendulum
 import yaml
-from pydantic_extra_types.pendulum_dt import DateTime
 
 from prefect.exceptions import InvalidRepositoryURLError
+from prefect.types import DateTime
 from prefect.utilities.collections import isiterable
 from prefect.utilities.dockerutils import get_prefect_image_name
 from prefect.utilities.filesystem import relative_path_to_current_platform
@@ -32,6 +33,7 @@ LOWERCASE_LETTERS_NUMBERS_AND_UNDERSCORES_REGEX = "^[a-z0-9_]*$"
 
 if TYPE_CHECKING:
     from prefect.blocks.core import Block
+    from prefect.serializers import Serializer
     from prefect.utilities.callables import ParameterSchema
 
 
@@ -577,7 +579,7 @@ def validate_picklelib_and_modules(values: dict) -> dict:
     return values
 
 
-def validate_dump_kwargs(value: dict) -> dict:
+def validate_dump_kwargs(value: dict[str, Any]) -> dict[str, Any]:
     # `default` is set by `object_encoder`. A user provided callable would make this
     # class unserializable anyway.
     if "default" in value:
@@ -585,7 +587,7 @@ def validate_dump_kwargs(value: dict) -> dict:
     return value
 
 
-def validate_load_kwargs(value: dict) -> dict:
+def validate_load_kwargs(value: dict[str, Any]) -> dict[str, Any]:
     # `object_hook` is set by `object_decoder`. A user provided callable would make
     # this class unserializable anyway.
     if "object_hook" in value:
@@ -595,7 +597,7 @@ def validate_load_kwargs(value: dict) -> dict:
     return value
 
 
-def cast_type_names_to_serializers(value):
+def cast_type_names_to_serializers(value: Union[str, "Serializer"]) -> "Serializer":
     from prefect.serializers import Serializer
 
     if isinstance(value, str):
@@ -653,7 +655,7 @@ def validate_message_template_variables(v: Optional[str]) -> Optional[str]:
     return v
 
 
-def validate_default_queue_id_not_none(v: Optional[str]) -> Optional[str]:
+def validate_default_queue_id_not_none(v: Optional[UUID]) -> UUID:
     if v is None:
         raise ValueError(
             "`default_queue_id` is a required field. If you are "

prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-02T18:57:11-0600",
+ "date": "2024-12-16T10:06:12-0800",
  "dirty": true,
  "error": null,
- "full-revisionid": "3c06654ec8be4f9e9bf5c304814f163b7727d28e",
- "version": "3.1.5"
+ "full-revisionid": "c05ffa6d1816f980d574b42119dffb9f8c8300a3",
+ "version": "3.1.7"
 }
 '''  # END VERSION_JSON
 

prefect/automations.py

@@ -1,10 +1,10 @@
-from typing import Optional
+from typing import Optional, Type
 from uuid import UUID
 
 from pydantic import Field
 from typing_extensions import Self
 
-from prefect.client.utilities import get_or_create_client
+from prefect.client.orchestration import get_client
 from prefect.events.actions import (
     CallWebhook,
     CancelFlowRun,
@@ -99,10 +99,10 @@ class Automation(AutomationCore):
         )
         created_automation = auto_to_create.create()
         """
-        client, _ = get_or_create_client()
-        automation = AutomationCore(**self.model_dump(exclude={"id"}))
-        self.id = await client.create_automation(automation=automation)
-        return self
+        async with get_client() as client:
+            automation = AutomationCore(**self.model_dump(exclude={"id"}))
+            self.id = await client.create_automation(automation=automation)
+            return self
 
     @sync_compatible
     async def update(self: Self):
@@ -112,15 +112,16 @@ class Automation(AutomationCore):
         auto.name = "new name"
         auto.update()
         """
-
-        client, _ = get_or_create_client()
-        automation = AutomationCore(**self.model_dump(exclude={"id", "owner_resource"}))
-        await client.update_automation(automation_id=self.id, automation=automation)
+        async with get_client() as client:
+            automation = AutomationCore(
+                **self.model_dump(exclude={"id", "owner_resource"})
+            )
+            await client.update_automation(automation_id=self.id, automation=automation)
 
     @classmethod
     @sync_compatible
     async def read(
-        cls: Self, id: Optional[UUID] = None, name: Optional[str] = None
+        cls: Type[Self], id: Optional[UUID] = None, name: Optional[str] = None
     ) -> Self:
         """
         Read an automation by ID or name.
@@ -134,20 +135,25 @@ class Automation(AutomationCore):
             raise ValueError("Only one of id or name can be provided")
         if not id and not name:
             raise ValueError("One of id or name must be provided")
-        client, _ = get_or_create_client()
-        if id:
-            try:
-                automation = await client.read_automation(automation_id=id)
-            except PrefectHTTPStatusError as exc:
-                if exc.response.status_code == 404:
+        async with get_client() as client:
+            if id:
+                try:
+                    automation = await client.read_automation(automation_id=id)
+                except PrefectHTTPStatusError as exc:
+                    if exc.response.status_code == 404:
+                        raise ValueError(f"Automation with ID {id!r} not found")
+                    raise
+                if automation is None:
                     raise ValueError(f"Automation with ID {id!r} not found")
-            return Automation(**automation.model_dump())
-        else:
-            automation = await client.read_automations_by_name(name=name)
-            if len(automation) > 0:
-                return Automation(**automation[0].model_dump()) if automation else None
+                return Automation(**automation.model_dump())
             else:
-                raise ValueError(f"Automation with name {name!r} not found")
+                automation = await client.read_automations_by_name(name=name)
+                if len(automation) > 0:
+                    return (
+                        Automation(**automation[0].model_dump()) if automation else None
+                    )
+                else:
+                    raise ValueError(f"Automation with name {name!r} not found")
 
     @sync_compatible
     async def delete(self: Self) -> bool:
@@ -155,14 +161,14 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.delete()
         """
-        try:
-            client, _ = get_or_create_client()
-            await client.delete_automation(self.id)
-            return True
-        except PrefectHTTPStatusError as exc:
-            if exc.response.status_code == 404:
-                return False
-            raise
+        async with get_client() as client:
+            try:
+                await client.delete_automation(self.id)
+                return True
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    return False
+                raise
 
     @sync_compatible
     async def disable(self: Self) -> bool:
@@ -171,14 +177,14 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.disable()
         """
-        try:
-            client, _ = get_or_create_client()
-            await client.pause_automation(self.id)
-            return True
-        except PrefectHTTPStatusError as exc:
-            if exc.response.status_code == 404:
-                return False
-            raise
+        async with get_client() as client:
+            try:
+                await client.pause_automation(self.id)
+                return True
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    return False
+                raise
 
     @sync_compatible
     async def enable(self: Self) -> bool:
@@ -187,11 +193,11 @@ class Automation(AutomationCore):
         auto = Automation.read(id = 123)
         auto.enable()
         """
-        try:
-            client, _ = get_or_create_client()
-            await client.resume_automation("asd")
-            return True
-        except PrefectHTTPStatusError as exc:
-            if exc.response.status_code == 404:
-                return False
-            raise
+        async with get_client() as client:
+            try:
+                await client.resume_automation(self.id)
+                return True
+            except PrefectHTTPStatusError as exc:
+                if exc.response.status_code == 404:
+                    return False
+                raise

prefect/blocks/abstract.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from logging import Logger
+from logging import Logger, LoggerAdapter
 from pathlib import Path
 from typing import (
     Any,
@@ -15,7 +15,7 @@ from typing import (
     Union,
 )
 
-from typing_extensions import Self
+from typing_extensions import Self, TypeAlias
 
 from prefect.blocks.core import Block
 from prefect.exceptions import MissingContextError
@@ -23,6 +23,8 @@ from prefect.logging.loggers import get_logger, get_run_logger
 
 T = TypeVar("T")
 
+LoggerOrAdapter: TypeAlias = Union[Logger, LoggerAdapter]
+
 
 class CredentialsBlock(Block, ABC):
     """
@@ -34,7 +36,7 @@ class CredentialsBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the CredentialsBlock
         is called from within a flow or task run context.
@@ -73,10 +75,10 @@ class NotificationBlock(Block, ABC):
     """
 
     _block_schema_capabilities = ["notify"]
-    _events_excluded_methods = Block._events_excluded_methods.default + ["notify"]
+    _events_excluded_methods = Block._events_excluded_methods + ["notify"]
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the NotificationBlock
         is called from within a flow or task run context.
@@ -123,7 +125,7 @@ class JobRun(ABC, Generic[T]):  # not a block
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the JobRun
         is called from within a flow or task run context.
@@ -158,7 +160,7 @@ class JobBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the JobBlock
         is called from within a flow or task run context.
@@ -202,7 +204,7 @@ class DatabaseBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the DatabaseBlock
         is called from within a flow or task run context.
@@ -337,7 +339,7 @@ class ObjectStorageBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the ObjectStorageBlock
         is called from within a flow or task run context.
@@ -469,7 +471,7 @@ class SecretBlock(Block, ABC):
     """
 
     @property
-    def logger(self) -> Logger:
+    def logger(self) -> LoggerOrAdapter:
         """
         Returns a logger based on whether the SecretBlock
         is called from within a flow or task run context.

prefect/blocks/core.py

@@ -41,6 +41,7 @@ from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal, ParamSpec, Self, get_args
 
 import prefect.exceptions
+from prefect._internal.compatibility.async_dispatch import async_dispatch
 from prefect.client.schemas import (
     DEFAULT_BLOCK_SCHEMA_VERSION,
     BlockDocument,
@@ -53,7 +54,7 @@ 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.asyncutils import run_coro_as_sync, sync_compatible
 from prefect.utilities.collections import listrepr, remove_nested_keys, visit_collection
 from prefect.utilities.dispatch import lookup_type, register_base_type
 from prefect.utilities.hashing import hash_objects
@@ -64,7 +65,7 @@ from prefect.utilities.slugify import slugify
 if TYPE_CHECKING:
     from pydantic.main import IncEx
 
-    from prefect.client.orchestration import PrefectClient
+    from prefect.client.orchestration import PrefectClient, SyncPrefectClient
 
 R = TypeVar("R")
 P = ParamSpec("P")
@@ -280,7 +281,7 @@ class Block(BaseModel, ABC):
         json_schema_extra=schema_extra,
     )
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
         self.block_initialization()
 
@@ -326,7 +327,9 @@ class Block(BaseModel, ABC):
 
     # Exclude `save` as it uses the `sync_compatible` decorator and needs to be
     # decorated directly.
-    _events_excluded_methods = ["block_initialization", "save", "dict"]
+    _events_excluded_methods: ClassVar[List[str]] = PrivateAttr(
+        default=["block_initialization", "save", "dict"]
+    )
 
     @classmethod
     def __dispatch_key__(cls):
@@ -626,7 +629,8 @@ class Block(BaseModel, ABC):
         """Generates a default code example for the current class"""
         qualified_name = to_qualified_name(cls)
         module_str = ".".join(qualified_name.split(".")[:-1])
-        class_name = cls.__name__
+        origin = cls.__pydantic_generic_metadata__.get("origin") or cls
+        class_name = origin.__name__
         block_variable_name = f'{cls.get_block_type_slug().replace("-", "_")}_block'
 
         return dedent(
@@ -775,12 +779,11 @@ class Block(BaseModel, ABC):
                 )
 
     @classmethod
-    @inject_client
-    async def _get_block_document(
+    async def _aget_block_document(
         cls,
         name: str,
-        client: Optional["PrefectClient"] = None,
-    ):
+        client: "PrefectClient",
+    ) -> tuple[BlockDocument, str]:
         if cls.__name__ == "Block":
             block_type_slug, block_document_name = name.split("/", 1)
         else:
@@ -799,6 +802,30 @@ class Block(BaseModel, ABC):
 
         return block_document, block_document_name
 
+    @classmethod
+    def _get_block_document(
+        cls,
+        name: str,
+        client: "SyncPrefectClient",
+    ) -> tuple[BlockDocument, str]:
+        if cls.__name__ == "Block":
+            block_type_slug, block_document_name = name.split("/", 1)
+        else:
+            block_type_slug = cls.get_block_type_slug()
+            block_document_name = name
+
+        try:
+            block_document = client.read_block_document_by_name(
+                name=block_document_name, block_type_slug=block_type_slug
+            )
+        except prefect.exceptions.ObjectNotFound as e:
+            raise ValueError(
+                f"Unable to find block document named {block_document_name} for block"
+                f" type {block_type_slug}"
+            ) from e
+
+        return block_document, block_document_name
+
     @classmethod
     @sync_compatible
     @inject_client
@@ -827,9 +854,97 @@ class Block(BaseModel, ABC):
         return block_document, block_document.name
 
     @classmethod
-    @sync_compatible
     @inject_client
-    async def load(
+    async def aload(
+        cls,
+        name: str,
+        validate: bool = True,
+        client: Optional["PrefectClient"] = None,
+    ) -> "Self":
+        """
+        Retrieves data from the block document with the given name for the block type
+        that corresponds with the current class and returns an instantiated version of
+        the current class with the data stored in the block document.
+
+        If a block document for a given block type is saved with a different schema
+        than the current class calling `aload`, a warning will be raised.
+
+        If the current class schema is a subset of the block document schema, the block
+        can be loaded as normal using the default `validate = True`.
+
+        If the current class schema is a superset of the block document schema, `aload`
+        must be called with `validate` set to False to prevent a validation error. In
+        this case, the block attributes will default to `None` and must be set manually
+        and saved to a new block document before the block can be used as expected.
+
+        Args:
+            name: The name or slug of the block document. A block document slug is a
+                string with the format <block_type_slug>/<block_document_name>
+            validate: If False, the block document will be loaded without Pydantic
+                validating the block schema. This is useful if the block schema has
+                changed client-side since the block document referred to by `name` was saved.
+            client: The client to use to load the block document. If not provided, the
+                default client will be injected.
+
+        Raises:
+            ValueError: If the requested block document is not found.
+
+        Returns:
+            An instance of the current class hydrated with the data stored in the
+            block document with the specified name.
+
+        Examples:
+            Load from a Block subclass with a block document name:
+            ```python
+            class Custom(Block):
+                message: str
+
+            Custom(message="Hello!").save("my-custom-message")
+
+            loaded_block = await Custom.aload("my-custom-message")
+            ```
+
+            Load from Block with a block document slug:
+            ```python
+            class Custom(Block):
+                message: str
+
+            Custom(message="Hello!").save("my-custom-message")
+
+            loaded_block = await Block.aload("custom/my-custom-message")
+            ```
+
+            Migrate a block document to a new schema:
+            ```python
+            # original class
+            class Custom(Block):
+                message: str
+
+            Custom(message="Hello!").save("my-custom-message")
+
+            # Updated class with new required field
+            class Custom(Block):
+                message: str
+                number_of_ducks: int
+
+            loaded_block = await Custom.aload("my-custom-message", validate=False)
+
+            # Prints UserWarning about schema mismatch
+
+            loaded_block.number_of_ducks = 42
+
+            loaded_block.save("my-custom-message", overwrite=True)
+            ```
+        """
+        if TYPE_CHECKING:
+            assert isinstance(client, PrefectClient)
+        block_document, _ = await cls._aget_block_document(name, client=client)
+
+        return cls._load_from_block_document(block_document, validate=validate)
+
+    @classmethod
+    @async_dispatch(aload)
+    def load(
         cls,
         name: str,
         validate: bool = True,
@@ -910,9 +1025,19 @@ class Block(BaseModel, ABC):
             loaded_block.save("my-custom-message", overwrite=True)
             ```
         """
-        block_document, block_document_name = await cls._get_block_document(
-            name, client=client
-        )
+        # Need to use a `PrefectClient` here to ensure `Block.load` and `Block.aload` signatures match
+        # TODO: replace with only sync client once all internal calls are updated to use `Block.aload` and `@async_dispatch` is removed
+        if client is None:
+            # If a client wasn't provided, we get to use a sync client
+            from prefect.client.orchestration import get_client
+
+            with get_client(sync_client=True) as sync_client:
+                block_document, _ = cls._get_block_document(name, client=sync_client)
+        else:
+            # If a client was provided, reuse it, even though it's async, to avoid excessive client creation
+            block_document, _ = run_coro_as_sync(
+                cls._aget_block_document(name, client=client)
+            )
 
         return cls._load_from_block_document(block_document, validate=validate)
 
@@ -966,14 +1091,16 @@ class Block(BaseModel, ABC):
         """
         block_document = None
         if isinstance(ref, (str, UUID)):
-            block_document, _ = await cls._get_block_document_by_id(ref)
+            block_document, _ = await cls._get_block_document_by_id(ref, client=client)
         elif isinstance(ref, dict):
             if block_document_id := ref.get("block_document_id"):
                 block_document, _ = await cls._get_block_document_by_id(
-                    block_document_id
+                    block_document_id, client=client
                 )
             elif block_document_slug := ref.get("block_document_slug"):
-                block_document, _ = await cls._get_block_document(block_document_slug)
+                block_document, _ = await cls._get_block_document(
+                    block_document_slug, client=client
+                )
 
         if not block_document:
             raise ValueError(f"Invalid reference format {ref!r}.")
@@ -1218,7 +1345,9 @@ class Block(BaseModel, ABC):
         name: str,
         client: Optional["PrefectClient"] = None,
     ):
-        block_document, block_document_name = await cls._get_block_document(name)
+        if TYPE_CHECKING:
+            assert isinstance(client, PrefectClient)
+        block_document, _ = await cls._aget_block_document(name, client=client)
 
         await client.delete_block_document(block_document.id)
 
@@ -1229,7 +1358,7 @@ class Block(BaseModel, ABC):
         """
         block_type_slug = kwargs.pop("block_type_slug", None)
         if block_type_slug:
-            subcls = lookup_type(cls, dispatch_key=block_type_slug)
+            subcls = cls.get_block_class_from_key(block_type_slug)
             return super().__new__(subcls)
         else:
             return super().__new__(cls)

prefect/blocks/system.py

@@ -9,10 +9,10 @@ from pydantic import (
     field_validator,
 )
 from pydantic import Secret as PydanticSecret
-from pydantic_extra_types.pendulum_dt import DateTime as PydanticDateTime
 
 from prefect._internal.compatibility.deprecated import deprecated_class
 from prefect.blocks.core import Block
+from prefect.types import DateTime as PydanticDateTime
 
 _SecretValueType = Union[
     Annotated[StrictStr, Field(title="string")],
@@ -130,6 +130,7 @@ class Secret(Block, Generic[T]):
 
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/c6f20e556dd16effda9df16551feecfb5822092b-48x48.png"
     _documentation_url = "https://docs.prefect.io/latest/develop/blocks"
+    _description = "A block that represents a secret value. The value stored in this block will be obfuscated when this block is viewed or edited in the UI."
 
     value: Union[SecretStr, PydanticSecret[T]] = Field(
         default=...,