prefect-client 3.4.6.dev1__py3-none-any.whl → 3.4.7__py3-none-any.whl

prefect/tasks.py

@@ -1,5 +1,5 @@
 """
-Module containing the base workflow task class and decorator - for most use cases, using the [`@task` decorator][prefect.tasks.task] is preferred.
+Module containing the base workflow task class and decorator - for most use cases, using the `@task` decorator is preferred.
 """
 
 # This file requires type-checking with pyright because mypy does not yet support PEP612
@@ -47,8 +47,8 @@ from prefect.cache_policies import DEFAULT, NO_CACHE, CachePolicy
 from prefect.client.orchestration import get_client
 from prefect.client.schemas import TaskRun
 from prefect.client.schemas.objects import (
+    RunInput,
     StateDetails,
-    TaskRunInput,
     TaskRunPolicy,
     TaskRunResult,
 )
@@ -244,12 +244,17 @@ def _infer_parent_task_runs(
     # tracked within the same flow run.
     if flow_run_context:
         for v in parameters.values():
+            upstream_state = None
+
             if isinstance(v, State):
                 upstream_state = v
             elif isinstance(v, PrefectFuture):
                 upstream_state = v.state
             else:
-                upstream_state = flow_run_context.task_run_results.get(id(v))
+                res = flow_run_context.run_results.get(id(v))
+                if res:
+                    upstream_state, _ = res
+
             if upstream_state and upstream_state.is_running():
                 parents.append(
                     TaskRunResult(id=upstream_state.state_details.task_run_id)
@@ -296,9 +301,6 @@ class Task(Generic[P, R]):
     """
     A Prefect task definition.
 
-    !!! note
-        We recommend using [the `@task` decorator][prefect.tasks.task] for most use-cases.
-
     Wraps a function with an entrypoint to the Prefect engine. Calling this class within a flow function
     creates a new task run.
 
@@ -840,7 +842,7 @@ class Task(Generic[P, R]):
         flow_run_context: Optional[FlowRunContext] = None,
         parent_task_run_context: Optional[TaskRunContext] = None,
         wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
-        extra_task_inputs: Optional[dict[str, set[TaskRunInput]]] = None,
+        extra_task_inputs: Optional[dict[str, set[RunInput]]] = None,
         deferred: bool = False,
     ) -> TaskRun:
         from prefect.utilities._engine import dynamic_key_for_task_run
@@ -943,7 +945,7 @@ class Task(Generic[P, R]):
         flow_run_context: Optional[FlowRunContext] = None,
         parent_task_run_context: Optional[TaskRunContext] = None,
         wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
-        extra_task_inputs: Optional[dict[str, set[TaskRunInput]]] = None,
+        extra_task_inputs: Optional[dict[str, set[RunInput]]] = None,
         deferred: bool = False,
     ) -> TaskRun:
         from prefect.utilities._engine import dynamic_key_for_task_run
@@ -1530,7 +1532,7 @@ class Task(Generic[P, R]):
         args: Optional[tuple[Any, ...]] = None,
         kwargs: Optional[dict[str, Any]] = None,
         wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
-        dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+        dependencies: Optional[dict[str, set[RunInput]]] = None,
     ) -> PrefectDistributedFuture[R]:
         """
         Create a pending task run for a task worker to execute.
@@ -2033,3 +2035,44 @@ class MaterializingTask(Task[P, R]):
             Asset(key=a) if isinstance(a, str) else a for a in assets
         ]
         self.materialized_by = materialized_by
+
+    def with_options(
+        self,
+        assets: Optional[Sequence[Union[str, Asset]]] = None,
+        **task_kwargs: Unpack[TaskOptions],
+    ) -> "MaterializingTask[P, R]":
+        import inspect
+
+        sig = inspect.signature(Task.__init__)
+
+        # Map parameter names to attribute names where they differ
+        # from parameter to attribute.
+        param_to_attr = {
+            "on_completion": "on_completion_hooks",
+            "on_failure": "on_failure_hooks",
+            "on_rollback": "on_rollback_hooks",
+            "on_commit": "on_commit_hooks",
+        }
+
+        # Build kwargs for Task constructor
+        init_kwargs = {}
+        for param_name in sig.parameters:
+            if param_name in ("self", "fn", "assets", "materialized_by"):
+                continue
+
+            attr_name = param_to_attr.get(param_name, param_name)
+            init_kwargs[param_name] = task_kwargs.get(
+                param_name, getattr(self, attr_name)
+            )
+
+        return MaterializingTask(
+            fn=self.fn,
+            assets=(
+                [Asset(key=a) if isinstance(a, str) else a for a in assets]
+                if assets is not None
+                else self.assets
+            ),
+            materialized_by=self.materialized_by,
+            # Now, the rest
+            **init_kwargs,
+        )

prefect/types/__init__.py

@@ -15,6 +15,7 @@ from .names import (
     WITHOUT_BANNED_CHARACTERS,
     MAX_VARIABLE_NAME_LENGTH,
     URILike,
+    ValidAssetKey,
 )
 from pydantic import (
     BeforeValidator,
@@ -216,6 +217,7 @@ __all__ = [
     "Name",
     "NameOrEmpty",
     "NonEmptyishName",
+    "ValidAssetKey",
     "SecretDict",
     "StatusCode",
     "StrictVariableValue",

prefect/types/names.py

@@ -160,3 +160,53 @@ URILike = Annotated[
         examples=["s3://bucket/folder/data.csv", "postgres://dbtable"],
     ),
 ]
+
+
+MAX_ASSET_KEY_LENGTH = 512
+
+RESTRICTED_ASSET_CHARACTERS = [
+    "\n",
+    "\r",
+    "\t",
+    "\0",
+    " ",
+    "#",
+    "?",
+    "&",
+    "%",
+    '"',
+    "'",
+    "<",
+    ">",
+    "[",
+    "]",
+    "{",
+    "}",
+    "|",
+    "\\",
+    "^",
+    "`",
+]
+
+
+def validate_valid_asset_key(value: str) -> str:
+    """Validate asset key with character restrictions and length limit."""
+    for char in RESTRICTED_ASSET_CHARACTERS:
+        if char in value:
+            raise ValueError(f"Asset key cannot contain '{char}'")
+
+    if len(value) > MAX_ASSET_KEY_LENGTH:
+        raise ValueError(f"Asset key cannot exceed {MAX_ASSET_KEY_LENGTH} characters")
+
+    return validate_uri(value)
+
+
+ValidAssetKey = Annotated[
+    str,
+    AfterValidator(validate_valid_asset_key),
+    Field(
+        max_length=MAX_ASSET_KEY_LENGTH,
+        description=f"A URI-like string with a lowercase protocol, restricted characters, and max {MAX_ASSET_KEY_LENGTH} characters",
+        examples=["s3://bucket/folder/data.csv", "postgres://dbtable"],
+    ),
+]

prefect/utilities/_ast.py

@@ -1,6 +1,6 @@
 import ast
 import math
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import anyio
 from typing_extensions import TypeAlias
@@ -17,7 +17,7 @@ OPEN_FILE_SEMAPHORE = LazySemaphore(lambda: math.floor(get_open_file_limit() * 0
 # this potentially could be a TypedDict, but you
 # need some way to convince the type checker that
 # Literal["flow_name", "task_name"] are being provided
-DecoratedFnMetadata: TypeAlias = dict[str, str]
+DecoratedFnMetadata: TypeAlias = dict[str, Any]
 
 
 async def find_prefect_decorated_functions_in_file(

prefect/utilities/callables.py

@@ -654,7 +654,7 @@ def _get_docstring_from_source(source_code: str, func_name: str) -> Optional[str
         and isinstance(func_def.body[0], ast.Expr)
         and isinstance(func_def.body[0].value, ast.Constant)
     ):
-        return func_def.body[0].value.value
+        return str(func_def.body[0].value.value)
     return None
 
 

prefect/utilities/collections.py

@@ -629,12 +629,12 @@ def get_from_dict(
         The fetched value if the key exists, or the default value if it does not.
 
     Examples:
-    >>> get_from_dict({'a': {'b': {'c': [1, 2, 3, 4]}}}, 'a.b.c[1]')
-    2
-    >>> get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, ['a', 'b', 1, 'c', 1])
-    2
-    >>> get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, 'a.b.1.c.2', 'default')
-    'default'
+
+    ```python
+    get_from_dict({'a': {'b': {'c': [1, 2, 3, 4]}}}, 'a.b.c[1]') # 2
+    get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, ['a', 'b', 1, 'c', 1]) # 2
+    get_from_dict({'a': {'b': [0, {'c': [1, 2]}]}}, 'a.b.1.c.2', 'default') # 'default'
+    ```
     """
     if isinstance(keys, str):
         keys = keys.replace("[", ".").replace("]", "").split(".")

prefect/utilities/engine.py

@@ -24,8 +24,11 @@ from typing_extensions import TypeIs
 import prefect
 import prefect.exceptions
 from prefect._internal.concurrency.cancellation import get_deadline
-from prefect.client.schemas import OrchestrationResult, TaskRun
-from prefect.client.schemas.objects import TaskRunInput, TaskRunResult
+from prefect.client.schemas import FlowRunResult, OrchestrationResult, TaskRun
+from prefect.client.schemas.objects import (
+    RunType,
+    TaskRunResult,
+)
 from prefect.client.schemas.responses import (
     SetStateStatus,
     StateAbortDetails,
@@ -60,20 +63,25 @@ engine_logger: Logger = get_logger("engine")
 T = TypeVar("T")
 
 
-async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRunResult]:
+async def collect_task_run_inputs(
+    expr: Any, max_depth: int = -1
+) -> set[Union[TaskRunResult, FlowRunResult]]:
     """
     This function recurses through an expression to generate a set of any discernible
     task run inputs it finds in the data structure. It produces a set of all inputs
     found.
 
     Examples:
-        >>> task_inputs = {
-        >>>    k: await collect_task_run_inputs(v) for k, v in parameters.items()
-        >>> }
+
+        ```python
+        task_inputs = {
+            k: await collect_task_run_inputs(v) for k, v in parameters.items()
+         }
+        ```
     """
     # TODO: This function needs to be updated to detect parameters and constants
 
-    inputs: set[TaskRunResult] = set()
+    inputs: set[Union[TaskRunResult, FlowRunResult]] = set()
 
     def add_futures_and_states_to_inputs(obj: Any) -> None:
         if isinstance(obj, PrefectFuture):
@@ -89,9 +97,12 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRun
         elif isinstance(obj, quote):
             raise StopVisiting
         else:
-            state = get_state_for_result(obj)
-            if state and state.state_details.task_run_id:
-                inputs.add(TaskRunResult(id=state.state_details.task_run_id))
+            res = get_state_for_result(obj)
+            if res:
+                state, run_type = res
+                run_result = state.state_details.to_run_result(run_type)
+                if run_result:
+                    inputs.add(run_result)
 
     visit_collection(
         expr,
@@ -105,20 +116,22 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> set[TaskRun
 
 def collect_task_run_inputs_sync(
     expr: Any, future_cls: Any = PrefectFuture, max_depth: int = -1
-) -> set[TaskRunInput]:
+) -> set[Union[TaskRunResult, FlowRunResult]]:
     """
     This function recurses through an expression to generate a set of any discernible
     task run inputs it finds in the data structure. It produces a set of all inputs
     found.
 
     Examples:
-        >>> task_inputs = {
-        >>>    k: collect_task_run_inputs_sync(v) for k, v in parameters.items()
-        >>> }
+        ```python
+        task_inputs = {
+            k: collect_task_run_inputs_sync(v) for k, v in parameters.items()
+         }
+        ```
     """
     # TODO: This function needs to be updated to detect parameters and constants
 
-    inputs: set[TaskRunInput] = set()
+    inputs: set[Union[TaskRunResult, FlowRunResult]] = set()
 
     def add_futures_and_states_to_inputs(obj: Any) -> None:
         if isinstance(obj, future_cls) and hasattr(obj, "task_run_id"):
@@ -138,9 +151,12 @@ def collect_task_run_inputs_sync(
         elif isinstance(obj, quote):
             raise StopVisiting
         else:
-            state = get_state_for_result(obj)
-            if state and state.state_details.task_run_id:
-                inputs.add(TaskRunResult(id=state.state_details.task_run_id))
+            res = get_state_for_result(obj)
+            if res:
+                state, run_type = res
+                run_result = state.state_details.to_run_result(run_type)
+                if run_result:
+                    inputs.add(run_result)
 
     visit_collection(
         expr,
@@ -299,12 +315,11 @@ def _is_result_record(data: Any) -> TypeIs[ResultRecord[Any]]:
 async def propose_state(
     client: "PrefectClient",
     state: State[Any],
+    flow_run_id: UUID,
     force: bool = False,
-    task_run_id: Optional[UUID] = None,
-    flow_run_id: Optional[UUID] = None,
 ) -> State[Any]:
     """
-    Propose a new state for a flow run or task run, invoking Prefect orchestration logic.
+    Propose a new state for a flow run, invoking Prefect orchestration logic.
 
     If the proposed state is accepted, the provided `state` will be augmented with
      details and returned.
@@ -319,25 +334,21 @@ async def propose_state(
     error will be raised.
 
     Args:
-        state: a new state for the task or flow run
-        task_run_id: an optional task run id, used when proposing task run states
+        state: a new state for a flow run
         flow_run_id: an optional flow run id, used when proposing flow run states
 
     Returns:
-        a [State model][prefect.client.schemas.objects.State] representation of the
-            flow or task run state
+        a State model representation of the flow run state
 
     Raises:
-        ValueError: if neither task_run_id or flow_run_id is provided
         prefect.exceptions.Abort: if an ABORT instruction is received from
             the Prefect API
     """
 
-    # Determine if working with a task run or flow run
-    if not task_run_id and not flow_run_id:
-        raise ValueError("You must provide either a `task_run_id` or `flow_run_id`")
+    if not flow_run_id:
+        raise ValueError("You must provide a `flow_run_id`")
 
-    # Handle task and sub-flow tracing
+    # Handle sub-flow tracing
     if state.is_final():
         result: Any
         if _is_result_record(state.data):
@@ -345,7 +356,7 @@ async def propose_state(
         else:
             result = state.data
 
-        link_state_to_result(state, result)
+        link_state_to_flow_run_result(state, result)
 
     # Handle repeated WAITs in a loop instead of recursively, to avoid
     # reaching max recursion depth in extreme cases.
@@ -364,18 +375,8 @@ async def propose_state(
             response = await set_state_func()
         return response
 
-    # Attempt to set the state
-    if task_run_id:
-        set_state = partial(client.set_task_run_state, task_run_id, state, force=force)
-        response = await set_state_and_handle_waits(set_state)
-    elif flow_run_id:
-        set_state = partial(client.set_flow_run_state, flow_run_id, state, force=force)
-        response = await set_state_and_handle_waits(set_state)
-    else:
-        raise ValueError(
-            "Neither flow run id or task run id were provided. At least one must "
-            "be given."
-        )
+    set_state = partial(client.set_flow_run_state, flow_run_id, state, force=force)
+    response = await set_state_and_handle_waits(set_state)
 
     # Parse the response to return the new state
     if response.status == SetStateStatus.ACCEPT:
@@ -412,12 +413,11 @@ async def propose_state(
 def propose_state_sync(
     client: "SyncPrefectClient",
     state: State[Any],
+    flow_run_id: UUID,
     force: bool = False,
-    task_run_id: Optional[UUID] = None,
-    flow_run_id: Optional[UUID] = None,
 ) -> State[Any]:
     """
-    Propose a new state for a flow run or task run, invoking Prefect orchestration logic.
+    Propose a new state for a flow run, invoking Prefect orchestration logic.
 
     If the proposed state is accepted, the provided `state` will be augmented with
      details and returned.
@@ -432,32 +432,26 @@ def propose_state_sync(
     error will be raised.
 
     Args:
-        state: a new state for the task or flow run
-        task_run_id: an optional task run id, used when proposing task run states
+        state: a new state for the flow run
         flow_run_id: an optional flow run id, used when proposing flow run states
 
     Returns:
-        a [State model][prefect.client.schemas.objects.State] representation of the
-            flow or task run state
+        a State model representation of the flow run state
 
     Raises:
-        ValueError: if neither task_run_id or flow_run_id is provided
+        ValueError: if flow_run_id is not provided
         prefect.exceptions.Abort: if an ABORT instruction is received from
             the Prefect API
     """
 
-    # Determine if working with a task run or flow run
-    if not task_run_id and not flow_run_id:
-        raise ValueError("You must provide either a `task_run_id` or `flow_run_id`")
-
-    # Handle task and sub-flow tracing
+    # Handle sub-flow tracing
     if state.is_final():
         if _is_result_record(state.data):
             result = state.data.result
         else:
             result = state.data
 
-        link_state_to_result(state, result)
+        link_state_to_flow_run_result(state, result)
 
     # Handle repeated WAITs in a loop instead of recursively, to avoid
     # reaching max recursion depth in extreme cases.
@@ -477,17 +471,8 @@ def propose_state_sync(
         return response
 
     # Attempt to set the state
-    if task_run_id:
-        set_state = partial(client.set_task_run_state, task_run_id, state, force=force)
-        response = set_state_and_handle_waits(set_state)
-    elif flow_run_id:
-        set_state = partial(client.set_flow_run_state, flow_run_id, state, force=force)
-        response = set_state_and_handle_waits(set_state)
-    else:
-        raise ValueError(
-            "Neither flow run id or task run id were provided. At least one must "
-            "be given."
-        )
+    set_state = partial(client.set_flow_run_state, flow_run_id, state, force=force)
+    response = set_state_and_handle_waits(set_state)
 
     # Parse the response to return the new state
     if response.status == SetStateStatus.ACCEPT:
@@ -519,7 +504,7 @@ def propose_state_sync(
         )
 
 
-def get_state_for_result(obj: Any) -> Optional[State]:
+def get_state_for_result(obj: Any) -> Optional[tuple[State, RunType]]:
     """
     Get the state related to a result object.
 
@@ -527,10 +512,20 @@ def get_state_for_result(obj: Any) -> Optional[State]:
     """
     flow_run_context = FlowRunContext.get()
     if flow_run_context:
-        return flow_run_context.task_run_results.get(id(obj))
+        return flow_run_context.run_results.get(id(obj))
+
+
+def link_state_to_flow_run_result(state: State, result: Any) -> None:
+    """Creates a link between a state and flow run result"""
+    link_state_to_result(state, result, RunType.FLOW_RUN)
+
+
+def link_state_to_task_run_result(state: State, result: Any) -> None:
+    """Creates a link between a state and task run result"""
+    link_state_to_result(state, result, RunType.TASK_RUN)
 
 
-def link_state_to_result(state: State, result: Any) -> None:
+def link_state_to_result(state: State, result: Any, run_type: RunType) -> None:
     """
     Caches a link between a state and a result and its components using
     the `id` of the components to map to the state. The cache is persisted to the
@@ -586,7 +581,7 @@ def link_state_to_result(state: State, result: Any) -> None:
             ):
                 state.state_details.untrackable_result = True
                 return
-            flow_run_context.task_run_results[id(obj)] = linked_state
+            flow_run_context.run_results[id(obj)] = (linked_state, run_type)
 
         visit_collection(expr=result, visit_fn=link_if_trackable, max_depth=1)
 

prefect/utilities/pydantic.py

@@ -1,4 +1,5 @@
 import warnings
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -20,6 +21,7 @@ from pydantic import (
 from pydantic_core import to_jsonable_python
 from typing_extensions import Literal
 
+from prefect.utilities.collections import visit_collection
 from prefect.utilities.dispatch import get_dispatch_key, lookup_type, register_base_type
 from prefect.utilities.importtools import from_qualified_name, to_qualified_name
 from prefect.utilities.names import obfuscate
@@ -344,7 +346,23 @@ def handle_secret_render(value: object, context: dict[str, Any]) -> object:
             else obfuscate(value)
         )
     elif isinstance(value, BaseModel):
-        return value.model_dump(context=context)
+        # Pass the serialization mode if available in context
+        mode = context.get("serialization_mode", "python")
+        if mode == "json":
+            # For JSON mode with nested models, we need to recursively process fields
+            # because regular Pydantic models don't understand include_secrets
+
+            json_data = value.model_dump(mode="json")
+            for field_name in type(value).model_fields:
+                field_value = getattr(value, field_name)
+                json_data[field_name] = visit_collection(
+                    expr=field_value,
+                    visit_fn=partial(handle_secret_render, context=context),
+                    return_data=True,
+                )
+            return json_data
+        else:
+            return value.model_dump(context=context)
     return value
 
 

prefect/workers/base.py

@@ -208,10 +208,12 @@ class BaseJobConfiguration(BaseModel):
         Defaults to using the job configuration parameter name as the template variable name.
 
         e.g.
+        ```python
         {
             key1: '{{ key1 }}',     # default variable template
             key2: '{{ template2 }}', # `template2` specifically provide as template
         }
+        ```
         """
         configuration: dict[str, Any] = {}
         properties = cls.model_json_schema()["properties"]

{prefect_client-3.4.6.dev1.dist-info → prefect_client-3.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.4.6.dev1
+Version: 3.4.7
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io