prefect-client 2.16.4__tar.gz → 2.16.6__tar.gz

{prefect-client-2.16.4 → prefect-client-2.16.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 2.16.4
+Version: 2.16.6
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect-client-2.16.4 → prefect-client-2.16.6}/requirements-dev.txt

@@ -31,6 +31,7 @@ pytkdocs >= 0.14.2
 pyyaml
 requests
 setuptools != 60.9.0; python_version < '3.8'
+vermin
 virtualenv
 watchfiles
 respx

{prefect-client-2.16.4 → prefect-client-2.16.6}/setup.cfg

@@ -81,6 +81,16 @@ omit =
 	src/prefect/server/database/migrations/versions/*
 ignore_errors = True
 
+[vermin]
+make_paths_absolute = no
+backports = 
+	typing_extensions
+format = parsable
+eval_annotations = yes
+only_show_violations = yes
+targets = 3.8-
+exclusion_regex = ^src/prefect/_vendor/.*$|^src/prefect/utilities/compat\.py$
+
 [egg_info]
 tag_build = 
 tag_date = 0

{prefect-client-2.16.4 → prefect-client-2.16.6}/src/prefect/_internal/concurrency/calls.py

@@ -2,6 +2,7 @@
 Implementation of the `Call` data structure for transport of deferred function calls
 and low-level management of call execution.
 """
+
 import abc
 import asyncio
 import concurrent.futures
@@ -40,8 +41,8 @@ P = ParamSpec("P")
 # we already have strong references to the `Call` objects in other places
 # and b) this is used for performance optimizations where we have fallback
 # behavior if this weakref is garbage collected. A fix for issue #10952.
-current_call: contextvars.ContextVar["weakref.ref[Call]"] = contextvars.ContextVar(
-    "current_call"
+current_call: contextvars.ContextVar["weakref.ref[Call]"] = (  # novm
+    contextvars.ContextVar("current_call")
 )
 
 # Create a strong reference to tasks to prevent destruction during execution errors

{prefect-client-2.16.4 → prefect-client-2.16.6}/src/prefect/_internal/pydantic/__init__.py

@@ -9,3 +9,5 @@
 from pydantic.version import VERSION as PYDANTIC_VERSION
 
 HAS_PYDANTIC_V2 = PYDANTIC_VERSION.startswith("2.")
+
+from ._compat import model_dump, model_json_schema, model_validate, IncEx

prefect-client-2.16.6/src/prefect/_internal/pydantic/_compat.py

@@ -0,0 +1,213 @@
+from typing import Any, Dict, Literal, Optional, Set, Type, Union
+
+import typing_extensions
+from pydantic import BaseModel
+
+from prefect._internal.pydantic import HAS_PYDANTIC_V2
+from prefect.logging.loggers import get_logger
+from prefect.settings import PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS
+
+IncEx: typing_extensions.TypeAlias = (
+    "Union[Set[int], Set[str], Dict[int, Any], Dict[str, Any], None]"
+)
+
+logger = get_logger("prefect._internal.pydantic")
+
+if HAS_PYDANTIC_V2:
+    from pydantic.json_schema import GenerateJsonSchema
+
+
+def is_pydantic_v2_compatible(
+    model_instance: Optional[BaseModel] = None, fn_name: Optional[str] = None
+) -> bool:
+    """
+    Determines if the current environment is compatible with Pydantic V2 features,
+    based on the presence of Pydantic V2 and a global setting that enables V2 functionalities.
+
+    This function primarily serves to facilitate conditional logic in code that needs to
+    operate differently depending on the availability of Pydantic V2 features. It checks
+    two conditions: whether Pydantic V2 is installed, and whether the use of V2 features
+    is explicitly enabled through a global setting (`PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS`).
+
+    Parameters:
+    -----------
+    model_instance : Optional[BaseModel], optional
+        An instance of a Pydantic model. This parameter is used to perform a type check
+        to ensure the passed object is a Pydantic model instance. If not provided or if
+        the object is not a Pydantic model, a TypeError is raised. Defaults to None.
+
+    fn_name : Optional[str], optional
+        The name of the function or feature for which V2 compatibility is being checked.
+        This is used for logging purposes to provide more context in debug messages.
+        Defaults to None.
+
+    Returns:
+    --------
+    bool
+        True if the current environment supports Pydantic V2 features and if the global
+        setting for enabling V2 features is set to True. False otherwise.
+
+    Raises:
+    -------
+    TypeError
+        If `model_instance` is provided but is not an instance of a Pydantic BaseModel.
+    """
+    if model_instance and not isinstance(model_instance, BaseModel):
+        raise TypeError(
+            f"Expected a Pydantic model, but got {type(model_instance).__name__}"
+        )
+
+    should_dump_as_v2_model = (
+        HAS_PYDANTIC_V2 and PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS
+    )
+
+    if should_dump_as_v2_model:
+        logger.debug(
+            f"Using Pydantic v2 compatibility layer for `{fn_name}`. This will be removed in a future release."
+        )
+
+        return True
+
+    elif HAS_PYDANTIC_V2:
+        logger.debug(
+            "Pydantic v2 compatibility layer is disabled. To enable, set `PREFECT_EXPERIMENTAL_ENABLE_PYDANTIC_V2_INTERNALS` to `True`."
+        )
+
+    else:
+        logger.debug("Pydantic v2 is not installed.")
+
+    return False
+
+
+def model_dump(
+    model_instance: BaseModel,
+    *,
+    mode: Union[Literal["json", "python"], str] = "python",
+    include: IncEx = None,
+    exclude: IncEx = None,
+    by_alias: bool = False,
+    exclude_unset: bool = False,
+    exclude_defaults: bool = False,
+    exclude_none: bool = False,
+    round_trip: bool = False,
+    warnings: bool = True,
+) -> Dict[str, Any]:
+    """
+    Generate a dictionary representation of the model, optionally specifying which fields to include or exclude.
+
+    Args:
+        mode: The mode in which `to_python` should run.
+            If mode is 'json', the output will only contain JSON serializable types.
+            If mode is 'python', the output may contain non-JSON-serializable Python objects.
+        include: A list of fields to include in the output.
+        exclude: A list of fields to exclude from the output.
+        by_alias: Whether to use the field's alias in the dictionary key if defined.
+        exclude_unset: Whether to exclude fields that have not been explicitly set.
+        exclude_defaults: Whether to exclude fields that are set to their default value.
+        exclude_none: Whether to exclude fields that have a value of `None`.
+        round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T].
+        warnings: Whether to log warnings when invalid fields are encountered.
+
+    Returns:
+        A dictionary representation of the model.
+    """
+    if is_pydantic_v2_compatible(model_instance=model_instance, fn_name="model_dump"):
+        return model_instance.model_dump(
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+        )
+
+    return model_instance.dict(
+        include=include,
+        exclude=exclude,
+        by_alias=by_alias,
+        exclude_unset=exclude_unset,
+        exclude_defaults=exclude_defaults,
+        exclude_none=exclude_none,
+    )
+
+
+DEFAULT_REF_TEMPLATE = "#/$defs/{model}"
+JsonSchemaMode = Literal["validation", "serialization"]
+
+
+def model_json_schema(
+    model: Type[BaseModel],
+    *,
+    by_alias: bool = True,
+    ref_template: str = DEFAULT_REF_TEMPLATE,
+    schema_generator=None,
+    mode: JsonSchemaMode = "validation",
+) -> Dict[str, Any]:
+    """
+    Generates a JSON schema for a model class.
+
+    Parameters
+    ----------
+    by_alias : bool, optional
+        Whether to use attribute aliases or not, by default True
+    ref_template : str, optional
+        The reference template, by default DEFAULT_REF_TEMPLATE
+    schema_generator : type[GenerateEmptySchemaForUserClasses], optional
+        To override the logic used to generate the JSON schema, as a subclass of GenerateEmptySchemaForUserClasses with your desired modifications, by default GenerateEmptySchemaForUserClasses
+    mode : JsonSchemaMode, optional
+        The mode in which to generate the schema, by default 'validation'
+
+    Returns
+    -------
+    dict[str, Any]
+        The JSON schema for the given model class.
+    """
+    if is_pydantic_v2_compatible(fn_name="model_json_schema"):
+        schema_generator = GenerateJsonSchema
+        return model.model_json_schema(
+            by_alias=by_alias,
+            ref_template=ref_template,
+            schema_generator=schema_generator,
+            mode=mode,
+        )
+
+    return model.schema(
+        by_alias=by_alias,
+        ref_template=ref_template,
+    )
+
+
+def model_validate(
+    model: Type[BaseModel],
+    obj: Any,
+    *,
+    strict: bool = False,
+    from_attributes: bool = False,
+    context: Optional[Dict[str, Any]] = None,
+) -> Union[BaseModel, Dict[str, Any]]:
+    """Validate a pydantic model instance.
+
+    Args:
+        obj: The object to validate.
+        strict: Whether to enforce types strictly.
+        from_attributes: Whether to extract data from object attributes.
+        context: Additional context to pass to the validator.
+
+    Raises:
+        ValidationError: If the object could not be validated.
+
+    Returns:
+        The validated model instance.
+    """
+    if is_pydantic_v2_compatible(fn_name="model_validate"):
+        return model.model_validate(
+            obj=obj,
+            strict=strict,
+            from_attributes=from_attributes,
+            context=context,
+        )
+
+    return model.parse_obj(obj)

prefect-client-2.16.6/src/prefect/_internal/pydantic/v1_schema.py

@@ -0,0 +1,48 @@
+import inspect
+import typing
+
+from prefect._internal.pydantic import HAS_PYDANTIC_V2
+
+if HAS_PYDANTIC_V2:
+    from pydantic.v1 import BaseModel as V1BaseModel
+else:
+    from pydantic import BaseModel as V1BaseModel
+
+
+def is_v1_model(v) -> bool:
+    if isinstance(v, V1BaseModel):
+        return True
+    try:
+        if inspect.isclass(v) and issubclass(v, V1BaseModel):
+            return True
+    except TypeError:
+        pass
+
+    return False
+
+
+def is_v1_type(v) -> bool:
+    if HAS_PYDANTIC_V2:
+        if is_v1_model(v):
+            return True
+
+        try:
+            return v.__module__.startswith("pydantic.v1.types")
+        except AttributeError:
+            return False
+
+    return True
+
+
+def has_v1_type_as_param(signature: inspect.Signature) -> bool:
+    parameters = signature.parameters.values()
+    for p in parameters:
+        # check if this parameter is a v1 model
+        if is_v1_type(p.annotation):
+            return True
+
+        # check if this parameter is a collection of types
+        for v in typing.get_args(p.annotation):
+            if is_v1_type(v):
+                return True
+    return False

{prefect-client-2.16.4 → prefect-client-2.16.6}/src/prefect/_internal/pydantic/v2_schema.py

@@ -7,7 +7,7 @@ import typing as t
 import pendulum
 import pydantic
 from pydantic import BaseModel as V2BaseModel
-from pydantic import ConfigDict, create_model
+from pydantic import ConfigDict, PydanticUndefinedAnnotation, create_model
 from pydantic.type_adapter import TypeAdapter
 
 from prefect._internal.pydantic.annotations.pendulum import (
@@ -106,7 +106,11 @@ def create_v2_schema(
     model = create_model(
         name_, __config__=model_cfg, __base__=model_base, **model_fields
     )
-    adapter = TypeAdapter(model)
+    try:
+        adapter = TypeAdapter(model)
+    except PydanticUndefinedAnnotation as exc:
+        # in v1 this raises a TypeError, which is handled by parameter_schema
+        raise TypeError(exc.message)
 
     # root model references under #definitions
     schema = adapter.json_schema(

prefect-client-2.16.6/src/prefect/artifacts.py

@@ -0,0 +1,185 @@
+"""
+Interface for creating and reading artifacts.
+"""
+
+import json
+import math
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas.actions import ArtifactCreate
+from prefect.client.utilities import inject_client
+from prefect.context import FlowRunContext, TaskRunContext
+from prefect.utilities.asyncutils import sync_compatible
+
+INVALID_TABLE_TYPE_ERROR = (
+    "`create_table_artifact` requires a `table` argument of type `dict[list]` or"
+    " `list[dict]`."
+)
+
+
+@inject_client
+async def _create_artifact(
+    type: str,
+    key: Optional[str] = None,
+    description: Optional[str] = None,
+    data: Optional[Union[Dict[str, Any], Any]] = None,
+    client: Optional[PrefectClient] = None,
+) -> UUID:
+    """
+    Helper function to create an artifact.
+
+    Arguments:
+        type: A string identifying the type of artifact.
+        key: A user-provided string identifier.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+        data: A JSON payload that allows for a result to be retrieved.
+        client: The PrefectClient
+
+    Returns:
+        - The table artifact ID.
+    """
+    artifact_args = {}
+    task_run_ctx = TaskRunContext.get()
+    flow_run_ctx = FlowRunContext.get()
+
+    if task_run_ctx:
+        artifact_args["task_run_id"] = task_run_ctx.task_run.id
+        artifact_args["flow_run_id"] = task_run_ctx.task_run.flow_run_id
+    elif flow_run_ctx:
+        artifact_args["flow_run_id"] = flow_run_ctx.flow_run.id
+
+    if key is not None:
+        artifact_args["key"] = key
+    if type is not None:
+        artifact_args["type"] = type
+    if description is not None:
+        artifact_args["description"] = description
+    if data is not None:
+        artifact_args["data"] = data
+
+    artifact = ArtifactCreate(**artifact_args)
+
+    return await client.create_artifact(artifact=artifact)
+
+
+@sync_compatible
+async def create_link_artifact(
+    link: str,
+    link_text: Optional[str] = None,
+    key: Optional[str] = None,
+    description: Optional[str] = None,
+) -> UUID:
+    """
+    Create a link artifact.
+
+    Arguments:
+        link: The link to create.
+        link_text: The link text.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+
+    Returns:
+        The table artifact ID.
+    """
+    formatted_link = f"[{link_text}]({link})" if link_text else f"[{link}]({link})"
+    artifact = await _create_artifact(
+        key=key,
+        type="markdown",
+        description=description,
+        data=formatted_link,
+    )
+
+    return artifact.id
+
+
+@sync_compatible
+async def create_markdown_artifact(
+    markdown: str,
+    key: Optional[str] = None,
+    description: Optional[str] = None,
+) -> UUID:
+    """
+    Create a markdown artifact.
+
+    Arguments:
+        markdown: The markdown to create.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The table artifact ID.
+    """
+    artifact = await _create_artifact(
+        key=key,
+        type="markdown",
+        description=description,
+        data=markdown,
+    )
+
+    return artifact.id
+
+
+@sync_compatible
+async def create_table_artifact(
+    table: Union[Dict[str, List[Any]], List[Dict[str, Any]], List[List[Any]]],
+    key: Optional[str] = None,
+    description: Optional[str] = None,
+) -> UUID:
+    """
+    Create a table artifact.
+
+    Arguments:
+        table: The table to create.
+        key: A user-provided string identifier.
+          Required for the artifact to show in the Artifacts page in the UI.
+          The key must only contain lowercase letters, numbers, and dashes.
+        description: A user-specified description of the artifact.
+
+    Returns:
+        The table artifact ID.
+    """
+
+    def _sanitize_nan_values(item):
+        """
+        Sanitize NaN values in a given item. The item can be a dict, list or float.
+        """
+
+        if isinstance(item, list):
+            return [_sanitize_nan_values(sub_item) for sub_item in item]
+
+        elif isinstance(item, dict):
+            return {k: _sanitize_nan_values(v) for k, v in item.items()}
+
+        elif isinstance(item, float) and math.isnan(item):
+            return None
+
+        else:
+            return item
+
+    sanitized_table = _sanitize_nan_values(table)
+
+    if isinstance(table, dict) and all(isinstance(v, list) for v in table.values()):
+        pass
+    elif isinstance(table, list) and all(isinstance(v, (list, dict)) for v in table):
+        pass
+    else:
+        raise TypeError(INVALID_TABLE_TYPE_ERROR)
+
+    formatted_table = json.dumps(sanitized_table)
+
+    artifact = await _create_artifact(
+        key=key,
+        type="table",
+        description=description,
+        data=formatted_table,
+    )
+
+    return artifact.id

{prefect-client-2.16.4 → prefect-client-2.16.6}/src/prefect/client/base.py

@@ -1,9 +1,10 @@
 import copy
 import sys
 import threading
+import uuid
 from collections import defaultdict
 from contextlib import asynccontextmanager
-from functools import partial
+from datetime import datetime, timezone
 from typing import (
     Any,
     AsyncGenerator,
@@ -11,6 +12,7 @@ from typing import (
     Callable,
     Dict,
     MutableMapping,
+    Optional,
     Protocol,
     Set,
     Tuple,
@@ -21,10 +23,11 @@ from typing import (
 import anyio
 import httpx
 from asgi_lifespan import LifespanManager
-from httpx import HTTPStatusError, Response
+from httpx import HTTPStatusError, Request, Response
 from prefect._vendor.starlette import status
 from typing_extensions import Self
 
+from prefect.client.schemas.objects import CsrfToken
 from prefect.exceptions import PrefectHTTPStatusError
 from prefect.logging import get_logger
 from prefect.settings import (
@@ -188,9 +191,20 @@ class PrefectHttpxClient(httpx.AsyncClient):
     [Configuring Cloudflare Rate Limiting](https://support.cloudflare.com/hc/en-us/articles/115001635128-Configuring-Rate-Limiting-from-UI)
     """
 
+    def __init__(self, *args, enable_csrf_support: bool = False, **kwargs):
+        self.enable_csrf_support: bool = enable_csrf_support
+        self.csrf_token: Optional[str] = None
+        self.csrf_token_expiration: Optional[datetime] = None
+        self.csrf_client_id: uuid.UUID = uuid.uuid4()
+
+        super().__init__(*args, **kwargs)
+
     async def _send_with_retry(
         self,
-        request: Callable,
+        request: Request,
+        send: Callable[[Request], Awaitable[Response]],
+        send_args: Tuple,
+        send_kwargs: Dict,
         retry_codes: Set[int] = set(),
         retry_exceptions: Tuple[Exception, ...] = tuple(),
     ):
@@ -207,21 +221,34 @@ class PrefectHttpxClient(httpx.AsyncClient):
         try_count = 0
         response = None
 
+        is_change_request = request.method.lower() in {"post", "put", "patch", "delete"}
+
+        if self.enable_csrf_support and is_change_request:
+            await self._add_csrf_headers(request=request)
+
         while try_count <= PREFECT_CLIENT_MAX_RETRIES.value():
             try_count += 1
             retry_seconds = None
             exc_info = None
 
             try:
-                response = await request()
+                response = await send(request, *send_args, **send_kwargs)
             except retry_exceptions:  # type: ignore
                 if try_count > PREFECT_CLIENT_MAX_RETRIES.value():
                     raise
                 # Otherwise, we will ignore this error but capture the info for logging
                 exc_info = sys.exc_info()
             else:
-                # We got a response; return immediately if it is not retryable
-                if response.status_code not in retry_codes:
+                # We got a response; check if it's a CSRF error, otherwise
+                # return immediately if it is not retryable
+                if (
+                    response.status_code == status.HTTP_403_FORBIDDEN
+                    and "Invalid CSRF token" in response.text
+                ):
+                    # We got a CSRF error, clear the token and try again
+                    self.csrf_token = None
+                    await self._add_csrf_headers(request)
+                elif response.status_code not in retry_codes:
                     return response
 
                 if "Retry-After" in response.headers:
@@ -268,19 +295,24 @@ class PrefectHttpxClient(httpx.AsyncClient):
         # We ran out of retries, return the failed response
         return response
 
-    async def send(self, *args, **kwargs) -> Response:
+    async def send(self, request: Request, *args, **kwargs) -> Response:
         """
         Send a request with automatic retry behavior for the following status codes:
 
+        - 403 Forbidden, if the request failed due to CSRF protection
+        - 408 Request Timeout
         - 429 CloudFlare-style rate limiting
         - 502 Bad Gateway
         - 503 Service unavailable
+        - Any additional status codes provided in `PREFECT_CLIENT_RETRY_EXTRA_CODES`
         """
 
-        api_request = partial(super().send, *args, **kwargs)
-
+        super_send = super().send
         response = await self._send_with_retry(
-            request=api_request,
+            request=request,
+            send=super_send,
+            send_args=args,
+            send_kwargs=kwargs,
             retry_codes={
                 status.HTTP_429_TOO_MANY_REQUESTS,
                 status.HTTP_503_SERVICE_UNAVAILABLE,
@@ -312,3 +344,41 @@ class PrefectHttpxClient(httpx.AsyncClient):
         response.raise_for_status()
 
         return response
+
+    async def _add_csrf_headers(self, request: Request):
+        now = datetime.now(timezone.utc)
+
+        if not self.enable_csrf_support:
+            return
+
+        if not self.csrf_token or (
+            self.csrf_token_expiration and now > self.csrf_token_expiration
+        ):
+            token_request = self.build_request(
+                "GET", f"/csrf-token?client={self.csrf_client_id}"
+            )
+
+            try:
+                token_response = await self.send(token_request)
+            except PrefectHTTPStatusError as exc:
+                old_server = exc.response.status_code == status.HTTP_404_NOT_FOUND
+                unconfigured_server = (
+                    exc.response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+                    and "CSRF protection is disabled." in exc.response.text
+                )
+
+                if old_server or unconfigured_server:
+                    # The token endpoint is either unavailable, suggesting an
+                    # older server, or CSRF protection is disabled. In either
+                    # case we should disable CSRF support.
+                    self.enable_csrf_support = False
+                    return
+
+                raise
+
+            token: CsrfToken = CsrfToken.parse_obj(token_response.json())
+            self.csrf_token = token.token
+            self.csrf_token_expiration = token.expiration
+
+        request.headers["Prefect-Csrf-Token"] = self.csrf_token
+        request.headers["Prefect-Csrf-Client"] = str(self.csrf_client_id)