prefect-client 3.1.6__py3-none-any.whl → 3.1.8__py3-none-any.whl

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()
 
@@ -628,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(
@@ -777,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:
@@ -801,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
@@ -829,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,
@@ -912,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)
 
@@ -968,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}.")
@@ -1220,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)
 

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=...,

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)
@@ -4415,3 +4433,91 @@ class SyncPrefectClient:
             json=labels,
         )
         response.raise_for_status()
+
+    def read_block_document_by_name(
+        self,
+        name: str,
+        block_type_slug: str,
+        include_secrets: bool = True,
+    ) -> BlockDocument:
+        """
+        Read the block document with the specified name that corresponds to a
+        specific block type name.
+
+        Args:
+            name: The block document name.
+            block_type_slug: The block type slug.
+            include_secrets (bool): whether to include secret values
+                on the Block, corresponding to Pydantic's `SecretStr` and
+                `SecretBytes` fields. These fields are automatically obfuscated
+                by Pydantic, but users can additionally choose not to receive
+                their values from the API. Note that any business logic on the
+                Block may not work if this is `False`.
+
+        Raises:
+            httpx.RequestError: if the block document was not found for any reason
+
+        Returns:
+            A block document or None.
+        """
+        try:
+            response = self._client.get(
+                f"/block_types/slug/{block_type_slug}/block_documents/name/{name}",
+                params=dict(include_secrets=include_secrets),
+            )
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
+                raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
+            else:
+                raise
+        return BlockDocument.model_validate(response.json())
+
+    def create_variable(self, variable: VariableCreate) -> Variable:
+        """
+        Creates an variable with the provided configuration.
+
+        Args:
+            variable: Desired configuration for the new variable.
+        Returns:
+            Information about the newly created variable.
+        """
+        response = self._client.post(
+            "/variables/",
+            json=variable.model_dump(mode="json", exclude_unset=True),
+        )
+        return Variable(**response.json())
+
+    def update_variable(self, variable: VariableUpdate) -> None:
+        """
+        Updates a variable with the provided configuration.
+
+        Args:
+            variable: Desired configuration for the updated variable.
+        Returns:
+            Information about the updated variable.
+        """
+        self._client.patch(
+            f"/variables/name/{variable.name}",
+            json=variable.model_dump(mode="json", exclude_unset=True),
+        )
+
+    def read_variable_by_name(self, name: str) -> Optional[Variable]:
+        """Reads a variable by name. Returns None if no variable is found."""
+        try:
+            response = self._client.get(f"/variables/name/{name}")
+            return Variable(**response.json())
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == status.HTTP_404_NOT_FOUND:
+                return None
+            else:
+                raise
+
+    def delete_variable_by_name(self, name: str) -> None:
+        """Deletes a variable by name."""
+        try:
+            self._client.delete(f"/variables/name/{name}")
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 404:
+                raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
+            else:
+                raise

prefect/client/schemas/actions.py

@@ -4,7 +4,6 @@ from uuid import UUID, uuid4
 
 import jsonschema
 from pydantic import Field, field_validator, model_validator
-from pydantic_extra_types.pendulum_dt import DateTime
 
 import prefect.client.schemas.objects as objects
 from prefect._internal.schemas.bases import ActionBaseModel
@@ -27,6 +26,7 @@ from prefect.client.schemas.schedules import SCHEDULE_TYPES
 from prefect.settings import PREFECT_DEPLOYMENT_SCHEDULE_MAX_SCHEDULED_RUNS
 from prefect.types import (
     MAX_VARIABLE_NAME_LENGTH,
+    DateTime,
     KeyValueLabelsField,
     Name,
     NonEmptyishName,
@@ -393,10 +393,10 @@ class DeploymentFlowRunCreate(ActionBaseModel):
         default_factory=objects.FlowRunPolicy
     )
     tags: list[str] = Field(default_factory=list)
-    idempotency_key: Optional[str] = Field(None)
-    parent_task_run_id: Optional[UUID] = Field(None)
-    work_queue_name: Optional[str] = Field(None)
-    job_variables: Optional[dict] = Field(None)
+    idempotency_key: Optional[str] = Field(default=None)
+    parent_task_run_id: Optional[UUID] = Field(default=None)
+    work_queue_name: Optional[str] = Field(default=None)
+    job_variables: Optional[dict[str, Any]] = Field(default=None)
     labels: KeyValueLabelsField = Field(default_factory=dict)
 
 

prefect/client/schemas/filters.py

@@ -6,10 +6,10 @@ from typing import List, Optional
 from uuid import UUID
 
 from pydantic import Field
-from pydantic_extra_types.pendulum_dt import DateTime
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 from prefect.client.schemas.objects import StateType
+from prefect.types import DateTime
 from prefect.utilities.collections import AutoEnum
 
 

prefect/client/schemas/objects.py

@@ -69,7 +69,7 @@ if TYPE_CHECKING:
 
     DateTime = pendulum.DateTime
 else:
-    from pydantic_extra_types.pendulum_dt import DateTime
+    from prefect.types import DateTime
 
 
 R = TypeVar("R", default=Any)
@@ -207,7 +207,7 @@ class State(ObjectBaseModel, Generic[R]):
 
     type: StateType
     name: Optional[str] = Field(default=None)
-    timestamp: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
+    timestamp: DateTime = Field(default_factory=lambda: DateTime.now("UTC"))
     message: Optional[str] = Field(default=None, examples=["Run started"])
     state_details: StateDetails = Field(default_factory=StateDetails)
     data: Annotated[
@@ -429,9 +429,9 @@ class State(ObjectBaseModel, Generic[R]):
         return self.model_copy(
             update={
                 "id": uuid4(),
-                "created": pendulum.now("utc"),
-                "updated": pendulum.now("utc"),
-                "timestamp": pendulum.now("utc"),
+                "created": DateTime.now("utc"),
+                "updated": DateTime.now("utc"),
+                "timestamp": DateTime.now("utc"),
             },
             **kwargs,
         )

prefect/client/schemas/responses.py

@@ -3,13 +3,12 @@ from typing import Any, ClassVar, Generic, Optional, TypeVar, Union
 from uuid import UUID
 
 from pydantic import ConfigDict, Field
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Literal
 
 import prefect.client.schemas.objects as objects
 from prefect._internal.schemas.bases import ObjectBaseModel, PrefectBaseModel
 from prefect._internal.schemas.fields import CreatedBy, UpdatedBy
-from prefect.types import KeyValueLabelsField
+from prefect.types import DateTime, KeyValueLabelsField
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.names import generate_slug
 

prefect/client/schemas/schedules.py

@@ -26,7 +26,7 @@ if TYPE_CHECKING:
     # together.
     DateTime = pendulum.DateTime
 else:
-    from pydantic_extra_types.pendulum_dt import DateTime
+    from prefect.types import DateTime
 
 MAX_ITERATIONS = 1000
 # approx. 1 years worth of RDATEs + buffer

prefect/client/subscriptions.py

@@ -10,6 +10,7 @@ from starlette.status import WS_1008_POLICY_VIOLATION
 from typing_extensions import Self
 
 from prefect._internal.schemas.bases import IDBaseModel
+from prefect.events.clients import websocket_connect
 from prefect.logging import get_logger
 from prefect.settings import PREFECT_API_KEY
 
@@ -34,7 +35,7 @@ class Subscription(Generic[S]):
 
         self.keys: list[str] = list(keys)
 
-        self._connect = websockets.connect(
+        self._connect = websocket_connect(
             self.subscription_url,
             subprotocols=[websockets.Subprotocol("prefect")],
         )

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
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, overload
 
 from typing_extensions import Concatenate, ParamSpec, TypeGuard, TypeVar
 
@@ -71,9 +71,23 @@ 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/context.py

@@ -26,7 +26,6 @@ from typing import (
 )
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
-from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Self
 
 import prefect.logging
@@ -48,6 +47,7 @@ from prefect.settings.legacy import (
 )
 from prefect.states import State
 from prefect.task_runners import TaskRunner
+from prefect.types import DateTime
 from prefect.utilities.services import start_client_metrics_server
 
 T = TypeVar("T")

prefect/deployments/flow_runs.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import TYPE_CHECKING, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
 from uuid import UUID
 
 import anyio
@@ -35,7 +35,7 @@ logger = get_logger(__name__)
 async def run_deployment(
     name: Union[str, UUID],
     client: Optional["PrefectClient"] = None,
-    parameters: Optional[dict] = None,
+    parameters: Optional[dict[str, Any]] = None,
     scheduled_time: Optional[datetime] = None,
     flow_run_name: Optional[str] = None,
     timeout: Optional[float] = None,
@@ -44,7 +44,7 @@ async def run_deployment(
     idempotency_key: Optional[str] = None,
     work_queue_name: Optional[str] = None,
     as_subflow: Optional[bool] = True,
-    job_variables: Optional[dict] = None,
+    job_variables: Optional[dict[str, Any]] = None,
 ) -> "FlowRun":
     """
     Create a flow run for a deployment and return it after completion or a timeout.