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

prefect/server/api/logs.py

@@ -2,14 +2,21 @@
 Routes for interacting with log objects.
 """
 
-from typing import List
+from typing import Optional, Sequence
 
-from fastapi import Body, Depends, status
+from fastapi import Body, Depends, WebSocket, status
+from pydantic import TypeAdapter
+from starlette.status import WS_1002_PROTOCOL_ERROR
 
 import prefect.server.api.dependencies as dependencies
 import prefect.server.models as models
-import prefect.server.schemas as schemas
 from prefect.server.database import PrefectDBInterface, provide_database_interface
+from prefect.server.logs import stream
+from prefect.server.schemas.actions import LogCreate
+from prefect.server.schemas.core import Log
+from prefect.server.schemas.filters import LogFilter
+from prefect.server.schemas.sorting import LogSort
+from prefect.server.utilities import subscriptions
 from prefect.server.utilities.server import PrefectRouter
 
 router: PrefectRouter = PrefectRouter(prefix="/logs", tags=["Logs"])
@@ -17,7 +24,7 @@ router: PrefectRouter = PrefectRouter(prefix="/logs", tags=["Logs"])
 
 @router.post("/", status_code=status.HTTP_201_CREATED)
 async def create_logs(
-    logs: List[schemas.actions.LogCreate],
+    logs: Sequence[LogCreate],
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> None:
     """
@@ -30,18 +37,66 @@ async def create_logs(
             await models.logs.create_logs(session=session, logs=batch)
 
 
+logs_adapter: TypeAdapter[Sequence[Log]] = TypeAdapter(Sequence[Log])
+
+
 @router.post("/filter")
 async def read_logs(
     limit: int = dependencies.LimitBody(),
     offset: int = Body(0, ge=0),
-    logs: schemas.filters.LogFilter = None,
-    sort: schemas.sorting.LogSort = Body(schemas.sorting.LogSort.TIMESTAMP_ASC),
+    logs: Optional[LogFilter] = None,
+    sort: LogSort = Body(LogSort.TIMESTAMP_ASC),
     db: PrefectDBInterface = Depends(provide_database_interface),
-) -> List[schemas.core.Log]:
+) -> Sequence[Log]:
     """
     Query for logs.
     """
     async with db.session_context() as session:
-        return await models.logs.read_logs(
-            session=session, log_filter=logs, offset=offset, limit=limit, sort=sort
+        return logs_adapter.validate_python(
+            await models.logs.read_logs(
+                session=session, log_filter=logs, offset=offset, limit=limit, sort=sort
+            )
         )
+
+
+@router.websocket("/out")
+async def stream_logs_out(websocket: WebSocket) -> None:
+    """Serve a WebSocket to stream live logs"""
+    websocket = await subscriptions.accept_prefect_socket(websocket)
+    if not websocket:
+        return
+
+    try:
+        # After authentication, the next message is expected to be a filter message, any
+        # other type of message will close the connection.
+        message = await websocket.receive_json()
+
+        if message["type"] != "filter":
+            return await websocket.close(
+                WS_1002_PROTOCOL_ERROR, reason="Expected 'filter' message"
+            )
+
+        try:
+            filter = LogFilter.model_validate(message["filter"])
+        except Exception as e:
+            return await websocket.close(
+                WS_1002_PROTOCOL_ERROR, reason=f"Invalid filter: {e}"
+            )
+
+        # No backfill support for logs - only live streaming
+        # Subscribe to the ongoing log stream
+        async with stream.logs(filter) as log_stream:
+            async for log in log_stream:
+                if not log:
+                    if await subscriptions.still_connected(websocket):
+                        continue
+                    break
+
+                await websocket.send_json(
+                    {"type": "log", "log": log.model_dump(mode="json")}
+                )
+
+    except subscriptions.NORMAL_DISCONNECT_EXCEPTIONS:  # pragma: no cover
+        pass  # it's fine if a client disconnects either normally or abnormally
+
+    return None

prefect/server/api/server.py

@@ -922,6 +922,8 @@ class SubprocessASGIServer:
                 self.server_process.wait(timeout=5)
             except subprocess.TimeoutExpired:
                 self.server_process.kill()
+                # Ensure the process is reaped to avoid ResourceWarning
+                self.server_process.wait()
             finally:
                 self.server_process = None
         if self.port in self._instances:

prefect/server/api/templates.py

@@ -8,11 +8,17 @@ from prefect.server.utilities.user_templates import (
     validate_user_template,
 )
 
-router: PrefectRouter = PrefectRouter(prefix="/templates", tags=["Automations"])
+router: PrefectRouter = PrefectRouter(tags=["Automations"])
 
 
+# deprecated and can be removed after the ui removes its dependency on it
+# use /templates/validate instead
 @router.post(
-    "/validate",
+    "/automations/templates/validate",
+    response_class=Response,
+)
+@router.post(
+    "/templates/validate",
     response_class=Response,
 )
 def validate_template(template: str = Body(default="")) -> Response:

prefect/settings/context.py

@@ -33,20 +33,23 @@ def temporary_settings(
     See `Settings.copy_with_update` for details on different argument behavior.
 
     Examples:
-        >>> from prefect.settings import PREFECT_API_URL
-        >>>
-        >>> with temporary_settings(updates={PREFECT_API_URL: "foo"}):
-        >>>    assert PREFECT_API_URL.value() == "foo"
-        >>>
-        >>>    with temporary_settings(set_defaults={PREFECT_API_URL: "bar"}):
-        >>>         assert PREFECT_API_URL.value() == "foo"
-        >>>
-        >>>    with temporary_settings(restore_defaults={PREFECT_API_URL}):
-        >>>         assert PREFECT_API_URL.value() is None
-        >>>
-        >>>         with temporary_settings(set_defaults={PREFECT_API_URL: "bar"})
-        >>>             assert PREFECT_API_URL.value() == "bar"
-        >>> assert PREFECT_API_URL.value() is None
+
+        ```python
+        from prefect.settings import PREFECT_API_URL
+
+        with temporary_settings(updates={PREFECT_API_URL: "foo"}):
+           assert PREFECT_API_URL.value() == "foo"
+
+           with temporary_settings(set_defaults={PREFECT_API_URL: "bar"}):
+                assert PREFECT_API_URL.value() == "foo"
+
+           with temporary_settings(restore_defaults={PREFECT_API_URL}):
+                assert PREFECT_API_URL.value() is None
+
+                with temporary_settings(set_defaults={PREFECT_API_URL: "bar"})
+                    assert PREFECT_API_URL.value() == "bar"
+        assert PREFECT_API_URL.value() is None
+        ```
     """
     import prefect.context
 

prefect/settings/models/server/logs.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import ClassVar
+
+from pydantic import Field
+from pydantic_settings import SettingsConfigDict
+
+from prefect.settings.base import PrefectBaseSettings, build_settings_config
+
+
+class ServerLogsSettings(PrefectBaseSettings):
+    """
+    Settings for controlling behavior of the logs subsystem
+    """
+
+    model_config: ClassVar[SettingsConfigDict] = build_settings_config(
+        ("server", "logs")
+    )
+
+    stream_out_enabled: bool = Field(
+        default=False,
+        description="Whether or not to stream logs out to the API via websockets.",
+    )
+
+    stream_publishing_enabled: bool = Field(
+        default=False,
+        description="Whether or not to publish logs to the streaming system.",
+    )

prefect/settings/models/server/root.py

@@ -13,6 +13,7 @@ from .deployments import ServerDeploymentsSettings
 from .ephemeral import ServerEphemeralSettings
 from .events import ServerEventsSettings
 from .flow_run_graph import ServerFlowRunGraphSettings
+from .logs import ServerLogsSettings
 from .services import ServerServicesSettings
 from .tasks import ServerTasksSettings
 from .ui import ServerUISettings
@@ -127,6 +128,10 @@ class ServerSettings(PrefectBaseSettings):
         default_factory=ServerFlowRunGraphSettings,
         description="Settings for controlling flow run graph behavior",
     )
+    logs: ServerLogsSettings = Field(
+        default_factory=ServerLogsSettings,
+        description="Settings for controlling server logs behavior",
+    )
     services: ServerServicesSettings = Field(
         default_factory=ServerServicesSettings,
         description="Settings for controlling server services behavior",

prefect/settings/models/server/services.py

@@ -448,6 +448,32 @@ class ServerServicesTriggersSettings(ServicesBaseSetting):
         ),
     )
 
+    pg_notify_reconnect_interval_seconds: int = Field(
+        default=10,
+        description="""
+        The number of seconds to wait before reconnecting to the PostgreSQL NOTIFY/LISTEN 
+        connection after an error. Only used when using PostgreSQL as the database.
+        Defaults to `10`.
+        """,
+        validation_alias=AliasChoices(
+            AliasPath("pg_notify_reconnect_interval_seconds"),
+            "prefect_server_services_triggers_pg_notify_reconnect_interval_seconds",
+        ),
+    )
+
+    pg_notify_heartbeat_interval_seconds: int = Field(
+        default=5,
+        description="""
+        The number of seconds between heartbeat checks for the PostgreSQL NOTIFY/LISTEN 
+        connection to ensure it's still alive. Only used when using PostgreSQL as the database.
+        Defaults to `5`.
+        """,
+        validation_alias=AliasChoices(
+            AliasPath("pg_notify_heartbeat_interval_seconds"),
+            "prefect_server_services_triggers_pg_notify_heartbeat_interval_seconds",
+        ),
+    )
+
 
 class ServerServicesSettings(PrefectBaseSettings):
     """

prefect/task_engine.py

@@ -37,7 +37,7 @@ import prefect.types._datetime
 from prefect.cache_policies import CachePolicy
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import TaskRun
-from prefect.client.schemas.objects import State, TaskRunInput
+from prefect.client.schemas.objects import RunInput, State
 from prefect.concurrency.context import ConcurrencyContext
 from prefect.concurrency.v1.asyncio import concurrency as aconcurrency
 from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
@@ -96,7 +96,7 @@ from prefect.utilities.callables import call_with_parameters, parameters_to_args
 from prefect.utilities.collections import visit_collection
 from prefect.utilities.engine import (
     emit_task_run_state_change_event,
-    link_state_to_result,
+    link_state_to_task_run_result,
     resolve_to_final_result,
 )
 from prefect.utilities.math import clamped_poisson_interval
@@ -453,7 +453,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 result = state.data
 
-            link_state_to_result(new_state, result)
+            link_state_to_task_run_result(new_state, result)
 
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
@@ -683,7 +683,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
     def initialize_run(
         self,
         task_run_id: Optional[UUID] = None,
-        dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+        dependencies: Optional[dict[str, set[RunInput]]] = None,
     ) -> Generator[Self, Any, Any]:
         """
         Enters a client context and creates a task run if needed.
@@ -777,7 +777,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
     def start(
         self,
         task_run_id: Optional[UUID] = None,
-        dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+        dependencies: Optional[dict[str, set[RunInput]]] = None,
     ) -> Generator[None, None, None]:
         with self.initialize_run(task_run_id=task_run_id, dependencies=dependencies):
             with (
@@ -1038,7 +1038,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 result = new_state.data
 
-            link_state_to_result(new_state, result)
+            link_state_to_task_run_result(new_state, result)
 
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
@@ -1267,7 +1267,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
     async def initialize_run(
         self,
         task_run_id: Optional[UUID] = None,
-        dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+        dependencies: Optional[dict[str, set[RunInput]]] = None,
     ) -> AsyncGenerator[Self, Any]:
         """
         Enters a client context and creates a task run if needed.
@@ -1359,7 +1359,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
     async def start(
         self,
         task_run_id: Optional[UUID] = None,
-        dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+        dependencies: Optional[dict[str, set[RunInput]]] = None,
     ) -> AsyncGenerator[None, None]:
         async with self.initialize_run(
             task_run_id=task_run_id, dependencies=dependencies
@@ -1465,7 +1465,7 @@ def run_task_sync(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["state", "result"] = "result",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> Union[R, State, None]:
     engine = SyncTaskRunEngine[P, R](
@@ -1496,7 +1496,7 @@ async def run_task_async(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["state", "result"] = "result",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> Union[R, State, None]:
     engine = AsyncTaskRunEngine[P, R](
@@ -1527,7 +1527,7 @@ def run_generator_task_sync(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["state", "result"] = "result",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> Generator[R, None, None]:
     if return_type != "result":
@@ -1568,7 +1568,7 @@ def run_generator_task_sync(
                             # dictionary in an unbounded way, so finding a
                             # way to periodically clean it up (using
                             # weakrefs or similar) would be good
-                            link_state_to_result(engine.state, gen_result)
+                            link_state_to_task_run_result(engine.state, gen_result)
                             yield gen_result
                     except StopIteration as exc:
                         engine.handle_success(exc.value, transaction=txn)
@@ -1586,7 +1586,7 @@ async def run_generator_task_async(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["state", "result"] = "result",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> AsyncGenerator[R, None]:
     if return_type != "result":
@@ -1627,7 +1627,7 @@ async def run_generator_task_async(
                             # dictionary in an unbounded way, so finding a
                             # way to periodically clean it up (using
                             # weakrefs or similar) would be good
-                            link_state_to_result(engine.state, gen_result)
+                            link_state_to_task_run_result(engine.state, gen_result)
                             yield gen_result
                     except (StopAsyncIteration, GeneratorExit) as exc:
                         await engine.handle_success(None, transaction=txn)
@@ -1647,7 +1647,7 @@ def run_task(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["state"] = "state",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> State[R]: ...
 
@@ -1660,7 +1660,7 @@ def run_task(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["result"] = "result",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> R: ...
 
@@ -1672,7 +1672,7 @@ def run_task(
     parameters: Optional[dict[str, Any]] = None,
     wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
     return_type: Literal["state", "result"] = "result",
-    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    dependencies: Optional[dict[str, set[RunInput]]] = None,
     context: Optional[dict[str, Any]] = None,
 ) -> Union[R, State, None, Coroutine[Any, Any, Union[R, State, None]]]:
     """

prefect/task_runners.py

@@ -19,7 +19,7 @@ from typing import (
 from typing_extensions import ParamSpec, Self, TypeVar
 
 from prefect._internal.uuid7 import uuid7
-from prefect.client.schemas.objects import TaskRunInput
+from prefect.client.schemas.objects import RunInput
 from prefect.exceptions import MappingLengthMismatch, MappingMissingIterable
 from prefect.futures import (
     PrefectConcurrentFuture,
@@ -81,7 +81,7 @@ class TaskRunner(abc.ABC, Generic[F]):
         task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> F: ...
 
     @overload
@@ -91,7 +91,7 @@ class TaskRunner(abc.ABC, Generic[F]):
         task: "Task[Any, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> F: ...
 
     @abc.abstractmethod
@@ -100,7 +100,7 @@ class TaskRunner(abc.ABC, Generic[F]):
         task: "Task[P, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> F: ...
 
     def map(
@@ -262,7 +262,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[R]]):
         task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectConcurrentFuture[R]: ...
 
     @overload
@@ -271,7 +271,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[R]]):
         task: "Task[Any, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectConcurrentFuture[R]: ...
 
     def submit(
@@ -279,7 +279,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[R]]):
         task: "Task[P, R | Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectConcurrentFuture[R]:
         """
         Submit a task to the task run engine running in a separate thread.
@@ -415,7 +415,7 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture[R]]):
         task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectDistributedFuture[R]: ...
 
     @overload
@@ -424,7 +424,7 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture[R]]):
         task: "Task[Any, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectDistributedFuture[R]: ...
 
     def submit(
@@ -432,7 +432,7 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture[R]]):
         task: "Task[P, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectDistributedFuture[R]:
         """
         Submit a task to the task run engine running in a separate thread.

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",