prefect-client 2.20.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/utilities/pydantic.py

@@ -1,25 +1,38 @@
 from functools import partial
-from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar, cast, overload
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Optional,
+    Type,
+    TypeVar,
+    cast,
+    get_origin,
+    overload,
+)
 
 from jsonpatch import JsonPatch as JsonPatchBase
-from pydantic_core import to_jsonable_python
-from typing_extensions import Self
+from pydantic import (
+    BaseModel,
+    GetJsonSchemaHandler,
+    Secret,
+    TypeAdapter,
+    ValidationError,
+)
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema, to_jsonable_python
+from typing_extensions import Literal
 
-from prefect._internal.pydantic.utilities.model_dump import model_dump
-from prefect.pydantic import HAS_PYDANTIC_V2
 from prefect.utilities.dispatch import get_dispatch_key, lookup_type, register_base_type
 from prefect.utilities.importtools import from_qualified_name, to_qualified_name
 
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic_v1
-else:
-    import pydantic as pydantic_v1
-
 D = TypeVar("D", bound=Any)
-M = TypeVar("M", bound=pydantic_v1.BaseModel)
+M = TypeVar("M", bound=BaseModel)
+T = TypeVar("T", bound=Any)
 
 
-def _reduce_model(model: pydantic_v1.BaseModel):
+def _reduce_model(model: BaseModel):
     """
     Helper for serializing a cythonized model with cloudpickle.
 
@@ -30,7 +43,7 @@ def _reduce_model(model: pydantic_v1.BaseModel):
         _unreduce_model,
         (
             to_qualified_name(type(model)),
-            model.json(**getattr(model, "__reduce_kwargs__", {})),
+            model.model_dump_json(**getattr(model, "__reduce_kwargs__", {})),
         ),
     )
 
@@ -38,7 +51,7 @@ def _reduce_model(model: pydantic_v1.BaseModel):
 def _unreduce_model(model_name, json):
     """Helper for restoring model after serialization"""
     model = from_qualified_name(model_name)
-    return model.parse_raw(json)
+    return model.model_validate_json(json)
 
 
 @overload
@@ -53,7 +66,7 @@ def add_cloudpickle_reduction(
     ...
 
 
-def add_cloudpickle_reduction(__model_cls: Type[M] = None, **kwargs: Any):
+def add_cloudpickle_reduction(__model_cls: Optional[Type[M]] = None, **kwargs: Any):
     """
     Adds a `__reducer__` to the given class that ensures it is cloudpickle compatible.
 
@@ -83,7 +96,7 @@ def add_cloudpickle_reduction(__model_cls: Type[M] = None, **kwargs: Any):
         )
 
 
-def get_class_fields_only(model: Type[pydantic_v1.BaseModel]) -> set:
+def get_class_fields_only(model: Type[BaseModel]) -> set:
     """
     Gets all the field names defined on the model class but not any parent classes.
     Any fields that are on the parent but redefined on the subclass are included.
@@ -92,7 +105,7 @@ def get_class_fields_only(model: Type[pydantic_v1.BaseModel]) -> set:
     parent_class_fields = set()
 
     for base in model.__class__.__bases__:
-        if issubclass(base, pydantic_v1.BaseModel):
+        if issubclass(base, BaseModel):
             parent_class_fields.update(base.__annotations__.keys())
 
     return (subclass_class_fields - parent_class_fields) | (
@@ -125,7 +138,7 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
         model_cls, "__dispatch_key__"
     ) or "__dispatch_key__" in getattr(model_cls, "__annotations__", {})
 
-    defines_type_field = "type" in model_cls.__fields__
+    defines_type_field = "type" in model_cls.model_fields
 
     if not defines_dispatch_key and not defines_type_field:
         raise ValueError(
@@ -133,18 +146,8 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
             "or a type field. One of these is required for dispatch."
         )
 
-    elif defines_dispatch_key and not defines_type_field:
-        # Add a type field to store the value of the dispatch key
-        model_cls.__fields__["type"] = pydantic_v1.fields.ModelField(
-            name="type",
-            type_=str,
-            required=True,
-            class_validators=None,
-            model_config=model_cls.__config__,
-        )
-
     elif not defines_dispatch_key and defines_type_field:
-        field_type_annotation = model_cls.__fields__["type"].type_
+        field_type_annotation = model_cls.model_fields["type"].annotation
         if field_type_annotation != str:
             raise TypeError(
                 f"Model class {model_cls.__name__!r} defines a 'type' field with "
@@ -154,7 +157,7 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
         # Set the dispatch key to retrieve the value from the type field
         @classmethod
         def dispatch_key_from_type_field(cls):
-            return cls.__fields__["type"].default
+            return cls.model_fields["type"].default
 
         model_cls.__dispatch_key__ = dispatch_key_from_type_field
 
@@ -176,12 +179,12 @@ def add_type_dispatch(model_cls: Type[M]) -> Type[M]:
         data.setdefault("type", type_string)
         cls_init(__pydantic_self__, **data)
 
-    def __new__(cls: Type[Self], **kwargs) -> Self:
+    def __new__(cls: Type[M], **kwargs: Any) -> M:
         if "type" in kwargs:
             try:
                 subcls = lookup_type(cls, dispatch_key=kwargs["type"])
             except KeyError as exc:
-                raise pydantic_v1.ValidationError(errors=[exc], model=cls)
+                raise ValidationError(errors=[exc], model=cls)
             return cls_new(subcls)
         else:
             return cls_new(cls)
@@ -207,7 +210,7 @@ class PartialModel(Generic[M]):
     a field already has a value.
 
     Example:
-        >>> class MyModel(pydantic_v1.BaseModel):
+        >>> class MyModel(BaseModel):
         >>>     x: int
         >>>     y: str
         >>>     z: float
@@ -237,7 +240,7 @@ class PartialModel(Generic[M]):
             raise ValueError(f"Field {name!r} has already been set.")
 
     def raise_if_not_in_model(self, name):
-        if name not in self.model_cls.__fields__:
+        if name not in self.model_cls.model_fields:
             raise ValueError(f"Field {name!r} is not present in the model.")
 
     def __setattr__(self, __name: str, __value: Any) -> None:
@@ -257,8 +260,22 @@ class PartialModel(Generic[M]):
 
 class JsonPatch(JsonPatchBase):
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetJsonSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.typed_dict_schema(
+            {"patch": core_schema.typed_dict_field(core_schema.dict_schema())}
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        json_schema = handler(core_schema)
+        json_schema = handler.resolve_ref_schema(json_schema)
+        json_schema.pop("required", None)
+        json_schema.pop("properties", None)
+        json_schema.update(
             {
                 "type": "array",
                 "format": "rfc6902",
@@ -268,6 +285,7 @@ class JsonPatch(JsonPatchBase):
                 },
             }
         )
+        return json_schema
 
 
 def custom_pydantic_encoder(
@@ -282,7 +300,79 @@ def custom_pydantic_encoder(
 
         return encoder(obj)
     else:  # We have exited the for loop without finding a suitable encoder
-        if isinstance(obj, pydantic_v1.BaseModel):
-            return model_dump(obj, mode="json")
+        if isinstance(obj, BaseModel):
+            return obj.model_dump(mode="json")
         else:
             return to_jsonable_python(obj)
+
+
+def parse_obj_as(
+    type_: type[T],
+    data: Any,
+    mode: Literal["python", "json", "strings"] = "python",
+) -> T:
+    """Parse a given data structure as a Pydantic model via `TypeAdapter`.
+
+    Read more about `TypeAdapter` [here](https://docs.pydantic.dev/latest/concepts/type_adapter/).
+
+    Args:
+        type_: The type to parse the data as.
+        data: The data to be parsed.
+        mode: The mode to use for parsing, either `python`, `json`, or `strings`.
+            Defaults to `python`, where `data` should be a Python object (e.g. `dict`).
+
+    Returns:
+        The parsed `data` as the given `type_`.
+
+
+    Example:
+        Basic Usage of `parse_as`
+        ```python
+        from prefect.utilities.pydantic import parse_as
+        from pydantic import BaseModel
+
+        class ExampleModel(BaseModel):
+            name: str
+
+        # parsing python objects
+        parsed = parse_as(ExampleModel, {"name": "Marvin"})
+        assert isinstance(parsed, ExampleModel)
+        assert parsed.name == "Marvin"
+
+        # parsing json strings
+        parsed = parse_as(
+            list[ExampleModel],
+            '[{"name": "Marvin"}, {"name": "Arthur"}]',
+            mode="json"
+        )
+        assert all(isinstance(item, ExampleModel) for item in parsed)
+        assert parsed[0].name == "Marvin"
+        assert parsed[1].name == "Arthur"
+
+        # parsing raw strings
+        parsed = parse_as(int, '123', mode="strings")
+        assert isinstance(parsed, int)
+        assert parsed == 123
+        ```
+
+    """
+    adapter = TypeAdapter(type_)
+
+    if get_origin(type_) is list and isinstance(data, dict):
+        data = next(iter(data.values()))
+
+    parser: Callable[[Any], T] = getattr(adapter, f"validate_{mode}")
+
+    return parser(data)
+
+
+def handle_secret_render(value: object, context: dict[str, Any]) -> object:
+    if hasattr(value, "get_secret_value"):
+        return (
+            cast(Secret[object], value).get_secret_value()
+            if context.get("include_secrets", False)
+            else "**********"
+        )
+    elif isinstance(value, BaseModel):
+        return value.model_dump(context=context)
+    return value

prefect/utilities/schema_tools/hydration.py

@@ -11,10 +11,14 @@ from prefect.server.utilities.user_templates import (
     render_user_template_sync,
     validate_user_template,
 )
+from prefect.types import StrictVariableValue
 
 
 class HydrationContext(BaseModel):
-    workspace_variables: Dict[str, str] = Field(default_factory=dict)
+    workspace_variables: Dict[
+        str,
+        StrictVariableValue,
+    ] = Field(default_factory=dict)
     render_workspace_variables: bool = Field(default=False)
     raise_on_error: bool = Field(default=False)
     render_jinja: bool = Field(default=False)
@@ -198,6 +202,11 @@ def json_handler(obj: dict, ctx: HydrationContext):
             dehydrated_json = _hydrate(obj["value"], ctx)
         else:
             dehydrated_json = obj["value"]
+
+        # If the result is a Placeholder, we should return it as is
+        if isinstance(dehydrated_json, Placeholder):
+            return dehydrated_json
+
         try:
             return json.loads(dehydrated_json)
         except (json.decoder.JSONDecodeError, TypeError) as e:
@@ -220,6 +229,10 @@ def jinja_handler(obj: dict, ctx: HydrationContext):
         else:
             dehydrated_jinja = obj["template"]
 
+        # If the result is a Placeholder, we should return it as is
+        if isinstance(dehydrated_jinja, Placeholder):
+            return dehydrated_jinja
+
         try:
             validate_user_template(dehydrated_jinja)
         except (jinja2.exceptions.TemplateSyntaxError, TemplateSecurityError) as exc:
@@ -241,6 +254,10 @@ def workspace_variable_handler(obj: dict, ctx: HydrationContext):
         else:
             dehydrated_variable = obj["variable_name"]
 
+        # If the result is a Placeholder, we should return it as is
+        if isinstance(dehydrated_variable, Placeholder):
+            return dehydrated_variable
+
         if not ctx.render_workspace_variables:
             return WorkspaceVariable(variable_name=obj["variable_name"])
 

prefect/utilities/schema_tools/validation.py

@@ -253,5 +253,35 @@ def preprocess_schema(
                 process_properties(
                     definition["properties"], required_fields, allow_none_with_default
                 )
+            # Allow block types to be referenced by their id
+            if "block_type_slug" in definition:
+                schema["definitions"][definition["title"]] = {
+                    "oneOf": [
+                        definition,
+                        {
+                            "type": "object",
+                            "properties": {
+                                "$ref": {
+                                    "oneOf": [
+                                        {
+                                            "type": "string",
+                                            "format": "uuid",
+                                        },
+                                        {
+                                            "type": "object",
+                                            "additionalProperties": {
+                                                "type": "string",
+                                            },
+                                            "minProperties": 1,
+                                        },
+                                    ]
+                                }
+                            },
+                            "required": [
+                                "$ref",
+                            ],
+                        },
+                    ]
+                }
 
     return schema

prefect/utilities/services.py

@@ -1,14 +1,16 @@
 import sys
-from asyncio import CancelledError
+import threading
 from collections import deque
 from traceback import format_exception
 from types import TracebackType
-from typing import Callable, Coroutine, Deque, Tuple
+from typing import Callable, Coroutine, Deque, Optional, Tuple
+from wsgiref.simple_server import WSGIServer
 
 import anyio
 import httpx
 
 from prefect.logging.loggers import get_logger
+from prefect.settings import PREFECT_CLIENT_ENABLE_METRICS, PREFECT_CLIENT_METRICS_PORT
 from prefect.utilities.collections import distinct
 from prefect.utilities.math import clamped_poisson_interval
 
@@ -23,7 +25,7 @@ async def critical_service_loop(
     backoff: int = 1,
     printer: Callable[..., None] = print,
     run_once: bool = False,
-    jitter_range: float = None,
+    jitter_range: Optional[float] = None,
 ):
     """
     Runs the given `workload` function on the specified `interval`, while being
@@ -68,11 +70,6 @@ async def critical_service_loop(
                 backoff_count = 0
 
             track_record.append(True)
-        except CancelledError as exc:
-            # Exit immediately because the task was cancelled, possibly due
-            # to a signal or timeout.
-            logger.debug(f"Run of {workload!r} cancelled", exc_info=exc)
-            return
         except httpx.TransportError as exc:
             # httpx.TransportError is the base class for any kind of communications
             # error, like timeouts, connection failures, etc.  This does _not_ cover
@@ -144,7 +141,7 @@ async def critical_service_loop(
             failures.clear()
             printer(
                 "Backing off due to consecutive errors, using increased interval of "
-                f" {interval * 2 ** backoff_count}s."
+                f" {interval * 2**backoff_count}s."
             )
 
         if run_once:
@@ -156,3 +153,32 @@ async def critical_service_loop(
             sleep = interval * 2**backoff_count
 
         await anyio.sleep(sleep)
+
+
+_metrics_server: Optional[Tuple[WSGIServer, threading.Thread]] = None
+
+
+def start_client_metrics_server():
+    """Start the process-wide Prometheus metrics server for client metrics (if enabled
+    with `PREFECT_CLIENT_ENABLE_METRICS`) on the port `PREFECT_CLIENT_METRICS_PORT`."""
+    if not PREFECT_CLIENT_ENABLE_METRICS:
+        return
+
+    global _metrics_server
+    if _metrics_server:
+        return
+
+    from prometheus_client import start_http_server
+
+    _metrics_server = start_http_server(port=PREFECT_CLIENT_METRICS_PORT.value())
+
+
+def stop_client_metrics_server():
+    """Start the process-wide Prometheus metrics server for client metrics, if it has
+    previously been started"""
+    global _metrics_server
+    if _metrics_server:
+        server, thread = _metrics_server
+        server.shutdown()
+        thread.join()
+        _metrics_server = None

prefect/utilities/templating.py

@@ -1,7 +1,17 @@
 import enum
 import os
 import re
-from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Set, Type, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    NamedTuple,
+    Optional,
+    Set,
+    Type,
+    TypeVar,
+    Union,
+)
 
 from prefect.client.utilities import inject_client
 from prefect.utilities.annotations import NotSet
@@ -301,7 +311,7 @@ async def resolve_block_document_references(
 
 
 @inject_client
-async def resolve_variables(template: T, client: "PrefectClient" = None):
+async def resolve_variables(template: T, client: Optional["PrefectClient"] = None):
     """
     Resolve variables in a template by replacing each variable placeholder with the
     value of the variable.

prefect/utilities/timeout.py

@@ -1,6 +1,6 @@
 from asyncio import CancelledError
 from contextlib import contextmanager
-from typing import Optional
+from typing import Optional, Type
 
 from prefect._internal.concurrency.cancellation import (
     cancel_async_after,
@@ -8,8 +8,19 @@ from prefect._internal.concurrency.cancellation import (
 )
 
 
+def fail_if_not_timeout_error(timeout_exc_type: Type[Exception]) -> None:
+    if not issubclass(timeout_exc_type, TimeoutError):
+        raise ValueError(
+            "The `timeout_exc_type` argument must be a subclass of `TimeoutError`."
+        )
+
+
 @contextmanager
-def timeout_async(seconds: Optional[float] = None):
+def timeout_async(
+    seconds: Optional[float] = None, timeout_exc_type: Type[TimeoutError] = TimeoutError
+):
+    fail_if_not_timeout_error(timeout_exc_type)
+
     if seconds is None:
         yield
         return
@@ -18,11 +29,15 @@ def timeout_async(seconds: Optional[float] = None):
         with cancel_async_after(timeout=seconds):
             yield
     except CancelledError:
-        raise TimeoutError(f"Scope timed out after {seconds} second(s).")
+        raise timeout_exc_type(f"Scope timed out after {seconds} second(s).")
 
 
 @contextmanager
-def timeout(seconds: Optional[float] = None):
+def timeout(
+    seconds: Optional[float] = None, timeout_exc_type: Type[TimeoutError] = TimeoutError
+):
+    fail_if_not_timeout_error(timeout_exc_type)
+
     if seconds is None:
         yield
         return
@@ -31,4 +46,4 @@ def timeout(seconds: Optional[float] = None):
         with cancel_sync_after(timeout=seconds):
             yield
     except CancelledError:
-        raise TimeoutError(f"Scope timed out after {seconds} second(s).")
+        raise timeout_exc_type(f"Scope timed out after {seconds} second(s).")