prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__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)

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/variables.py

@@ -1,9 +1,11 @@
-from typing import List, Optional
+from typing import List, Optional, Union
 
-from prefect._internal.compatibility.deprecated import deprecated_callable
 from prefect.client.schemas.actions import VariableCreate as VariableRequest
 from prefect.client.schemas.actions import VariableUpdate as VariableUpdateRequest
+from prefect.client.schemas.objects import Variable as VariableResponse
 from prefect.client.utilities import get_or_create_client
+from prefect.exceptions import ObjectNotFound
+from prefect.types import StrictVariableValue
 from prefect.utilities.asyncutils import sync_compatible
 
 
@@ -23,92 +25,112 @@ class Variable(VariableRequest):
     async def set(
         cls,
         name: str,
-        value: str,
+        value: StrictVariableValue,
         tags: Optional[List[str]] = None,
         overwrite: bool = False,
-    ) -> Optional[str]:
+        as_object: bool = False,
+    ):
         """
-        Sets a new variable. If one exists with the same name, user must pass `overwrite=True`
-        ```
-            from prefect.variables import Variable
+        Sets a new variable. If one exists with the same name, must pass `overwrite=True`
 
-            @flow
-            def my_flow():
-                var = Variable.set(name="my_var",value="test_value", tags=["hi", "there"], overwrite=True)
-        ```
-        or
-        ```
+        Returns the newly set value. If `as_object=True`, return the full Variable object
+
+        Args:
+            - name: The name of the variable to set.
+            - value: The value of the variable to set.
+            - tags: An optional list of strings to associate with the variable.
+            - overwrite: Whether to overwrite the variable if it already exists.
+            - as_object: Whether to return the full Variable object.
+
+        Example:
+            Set a new variable and overwrite it if it already exists.
+            ```
             from prefect.variables import Variable
 
             @flow
-            async def my_flow():
-                var = await Variable.set(name="my_var",value="test_value", tags=["hi", "there"], overwrite=True)
-        ```
+            def my_flow():
+                Variable.set(name="my_var",value="test_value", tags=["hi", "there"], overwrite=True)
+            ```
         """
         client, _ = get_or_create_client()
-        variable = await client.read_variable_by_name(name)
-        var_dict = {"name": name, "value": value}
-        var_dict["tags"] = tags or []
-        if variable:
+        variable_exists = await client.read_variable_by_name(name)
+        var_dict = {"name": name, "value": value, "tags": tags or []}
+
+        if variable_exists:
             if not overwrite:
                 raise ValueError(
-                    "You are attempting to save a variable with a name that is already in use. If you would like to overwrite the values that are saved, then call .set with `overwrite=True`."
+                    f"Variable {name!r} already exists. Use `overwrite=True` to update it."
                 )
-            var = VariableUpdateRequest(**var_dict)
-            await client.update_variable(variable=var)
+            await client.update_variable(variable=VariableUpdateRequest(**var_dict))
             variable = await client.read_variable_by_name(name)
         else:
-            var = VariableRequest(**var_dict)
-            variable = await client.create_variable(variable=var)
+            variable = await client.create_variable(
+                variable=VariableRequest(**var_dict)
+            )
 
-        return variable if variable else None
+        return variable if as_object else variable.value
 
     @classmethod
     @sync_compatible
-    async def get(cls, name: str, default: Optional[str] = None) -> Optional[str]:
+    async def get(
+        cls,
+        name: str,
+        default: StrictVariableValue = None,
+        as_object: bool = False,
+    ) -> Union[StrictVariableValue, VariableResponse]:
         """
-        Get a variable by name. If doesn't exist return the default.
-        ```
+        Get a variable's value by name.
+
+        If the variable does not exist, return the default value.
+
+        If `as_object=True`, return the full variable object. `default` is ignored in this case.
+
+        Args:
+            - name: The name of the variable to get.
+            - default: The default value to return if the variable does not exist.
+            - as_object: Whether to return the full variable object.
+
+        Example:
+            Get a variable's value by name.
+            ```python
+            from prefect import flow
             from prefect.variables import Variable
 
             @flow
             def my_flow():
                 var = Variable.get("my_var")
-        ```
-        or
-        ```
-            from prefect.variables import Variable
-
-            @flow
-            async def my_flow():
-                var = await Variable.get("my_var")
-        ```
+            ```
         """
         client, _ = get_or_create_client()
         variable = await client.read_variable_by_name(name)
-        return variable if variable else default
 
+        return variable if as_object else (variable.value if variable else default)
 
-@deprecated_callable(start_date="Apr 2024")
-@sync_compatible
-async def get(name: str, default: Optional[str] = None) -> Optional[str]:
-    """
-    Get a variable by name. If doesn't exist return the default.
-    ```
-        from prefect import variables
-
-        @flow
-        def my_flow():
-            var = variables.get("my_var")
-    ```
-    or
-    ```
-        from prefect import variables
-
-        @flow
-        async def my_flow():
-            var = await variables.get("my_var")
-    ```
-    """
-    variable = await Variable.get(name)
-    return variable.value if variable else default
+    @classmethod
+    @sync_compatible
+    async def unset(cls, name: str) -> bool:
+        """
+        Unset a variable by name.
+
+        Args:
+            - name: The name of the variable to unset.
+
+        Returns `True` if the variable was deleted, `False` if the variable did not exist.
+
+        Example:
+            Unset a variable by name.
+            ```python
+            from prefect import flow
+            from prefect.variables import Variable
+
+            @flow
+            def my_flow():
+                Variable.unset("my_var")
+            ```
+        """
+        client, _ = get_or_create_client()
+        try:
+            await client.delete_variable_by_name(name=name)
+            return True
+        except ObjectNotFound:
+            return False

prefect/workers/__init__.py

@@ -1,2 +1 @@
 from .process import ProcessWorker
-from .block import BlockWorker

prefect/workers/base.py

@@ -7,14 +7,7 @@ from uuid import uuid4
 import anyio
 import anyio.abc
 import pendulum
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-from prefect._internal.schemas.validators import return_v_or_none
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel, Field, PrivateAttr, validator
-else:
-    from pydantic import BaseModel, Field, PrivateAttr, validator
+from pydantic import BaseModel, Field, PrivateAttr, field_validator
 
 import prefect
 from prefect._internal.compatibility.experimental import (
@@ -22,6 +15,7 @@ from prefect._internal.compatibility.experimental import (
     ExperimentalFeature,
     experiment_enabled,
 )
+from prefect._internal.schemas.validators import return_v_or_none
 from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas.actions import WorkPoolCreate, WorkPoolUpdate
 from prefect.client.schemas.filters import (
@@ -37,7 +31,6 @@ from prefect.client.schemas.filters import (
 )
 from prefect.client.schemas.objects import StateType, WorkPool
 from prefect.client.utilities import inject_client
-from prefect.engine import propose_state
 from prefect.events import Event, RelatedResource, emit_event
 from prefect.events.related import object_as_related_resource, tags_as_related_resources
 from prefect.exceptions import (
@@ -57,6 +50,7 @@ from prefect.settings import (
 )
 from prefect.states import Crashed, Pending, exception_to_failed_state
 from prefect.utilities.dispatch import get_registry_for_type, register_base_type
+from prefect.utilities.engine import propose_state
 from prefect.utilities.slugify import slugify
 from prefect.utilities.templating import (
     apply_values,
@@ -107,16 +101,27 @@ class BaseJobConfiguration(BaseModel):
     def is_using_a_runner(self):
         return self.command is not None and "prefect flow-run execute" in self.command
 
-    @validator("command")
+    @field_validator("command")
+    @classmethod
     def _coerce_command(cls, v):
         return return_v_or_none(v)
 
+    @field_validator("env", mode="before")
+    @classmethod
+    def _coerce_env(cls, v):
+        return {k: str(v) if v is not None else None for k, v in v.items()}
+
     @staticmethod
     def _get_base_config_defaults(variables: dict) -> dict:
         """Get default values from base config for all variables that have them."""
         defaults = dict()
         for variable_name, attrs in variables.items():
-            if "default" in attrs:
+            # We remote `None` values because we don't want to use them in templating.
+            # The currently logic depends on keys not existing to populate the correct value
+            # in some cases.
+            # Pydantic will provide default values if the keys are missing when creating
+            # a configuration class.
+            if "default" in attrs and attrs.get("default") is not None:
                 defaults[variable_name] = attrs["default"]
 
         return defaults
@@ -124,7 +129,10 @@ class BaseJobConfiguration(BaseModel):
     @classmethod
     @inject_client
     async def from_template_and_values(
-        cls, base_job_template: dict, values: dict, client: "PrefectClient" = None
+        cls,
+        base_job_template: dict,
+        values: dict,
+        client: Optional["PrefectClient"] = None,
     ):
         """Creates a valid worker configuration object from the provided base
         configuration and overrides.
@@ -161,7 +169,7 @@ class BaseJobConfiguration(BaseModel):
         }
         """
         configuration = {}
-        properties = cls.schema()["properties"]
+        properties = cls.model_json_schema()["properties"]
         for k, v in properties.items():
             if v.get("template"):
                 template = v["template"]
@@ -420,7 +428,7 @@ class BaseWorker(abc.ABC):
     @classmethod
     def get_default_base_job_template(cls) -> Dict:
         if cls.job_configuration_variables is None:
-            schema = cls.job_configuration.schema()
+            schema = cls.job_configuration.model_json_schema()
             # remove "template" key from all dicts in schema['properties'] because it is not a
             # relevant field
             for key, value in schema["properties"].items():
@@ -428,7 +436,7 @@ class BaseWorker(abc.ABC):
                     schema["properties"][key].pop("template", None)
             variables_schema = schema
         else:
-            variables_schema = cls.job_configuration_variables.schema()
+            variables_schema = cls.job_configuration_variables.model_json_schema()
         variables_schema.pop("title", None)
         return {
             "job_configuration": cls.job_configuration.json_template(),
@@ -963,7 +971,7 @@ class BaseWorker(abc.ABC):
         return {
             "name": self.name,
             "work_pool": (
-                self._work_pool.dict(json_compatible=True)
+                self._work_pool.model_dump(mode="json")
                 if self._work_pool is not None
                 else None
             ),
@@ -1068,7 +1076,7 @@ class BaseWorker(abc.ABC):
         state_updates = state_updates or {}
         state_updates.setdefault("name", "Cancelled")
         state_updates.setdefault("type", StateType.CANCELLED)
-        state = flow_run.state.copy(update=state_updates)
+        state = flow_run.state.model_copy(update=state_updates)
 
         await self._client.set_flow_run_state(flow_run.id, state, force=True)
 
@@ -1155,7 +1163,7 @@ class BaseWorker(abc.ABC):
         if include_self:
             worker_resource = self._event_resource()
             worker_resource["prefect.resource.role"] = "worker"
-            related.append(RelatedResource.parse_obj(worker_resource))
+            related.append(RelatedResource.model_validate(worker_resource))
 
         return related
 

prefect/workers/process.py

@@ -26,13 +26,7 @@ from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
 import anyio
 import anyio.abc
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, validator
-else:
-    from pydantic import Field, validator
+from pydantic import Field, field_validator
 
 from prefect._internal.schemas.validators import validate_command
 from prefect.client.schemas import FlowRun
@@ -68,7 +62,8 @@ class ProcessJobConfiguration(BaseJobConfiguration):
     stream_output: bool = Field(default=True)
     working_dir: Optional[Path] = Field(default=None)
 
-    @validator("working_dir")
+    @field_validator("working_dir")
+    @classmethod
     def validate_command(cls, v):
         return validate_command(v)
 

prefect/workers/server.py

@@ -1,8 +1,8 @@
 from typing import Union
 
 import uvicorn
-from prefect._vendor.fastapi import APIRouter, FastAPI, status
-from prefect._vendor.fastapi.responses import JSONResponse
+from fastapi import APIRouter, FastAPI, status
+from fastapi.responses import JSONResponse
 
 from prefect.settings import (
     PREFECT_WORKER_WEBSERVER_HOST,