prefect-client 3.1.12__py3-none-any.whl → 3.1.14__py3-none-any.whl

prefect/flow_engine.py

@@ -39,6 +39,7 @@ from prefect.context import (
     FlowRunContext,
     SyncClientContext,
     TagsContext,
+    hydrated_context,
 )
 from prefect.exceptions import (
     Abort,
@@ -137,6 +138,7 @@ class BaseFlowRunEngine(Generic[P, R]):
     flow_run_id: Optional[UUID] = None
     logger: logging.Logger = field(default_factory=lambda: get_logger("engine"))
     wait_for: Optional[Iterable[PrefectFuture[Any]]] = None
+    context: Optional[dict[str, Any]] = None
     # holds the return value from the user code
     _return_value: Union[R, Type[NotSet]] = NotSet
     # holds the exception raised by the user code, if any
@@ -146,7 +148,7 @@ class BaseFlowRunEngine(Generic[P, R]):
     _flow_run_name_set: bool = False
     _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.flow is None and self.flow_run_id is None:
             raise ValueError("Either a flow or a flow_run_id must be provided.")
 
@@ -167,7 +169,7 @@ class BaseFlowRunEngine(Generic[P, R]):
             return False  # TODO: handle this differently?
         return getattr(self, "flow_run").state.is_pending()
 
-    def cancel_all_tasks(self):
+    def cancel_all_tasks(self) -> None:
         if hasattr(self.flow.task_runner, "cancel_all"):
             self.flow.task_runner.cancel_all()  # type: ignore
 
@@ -208,6 +210,8 @@ class BaseFlowRunEngine(Generic[P, R]):
 @dataclass
 class FlowRunEngine(BaseFlowRunEngine[P, R]):
     _client: Optional[SyncPrefectClient] = None
+    flow_run: FlowRun | None = None
+    parameters: dict[str, Any] | None = None
 
     @property
     def client(self) -> SyncPrefectClient:
@@ -502,7 +506,7 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
             tags=TagsContext.get().current_tags,
         )
 
-    def call_hooks(self, state: Optional[State] = None):
+    def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         flow = self.flow
@@ -600,7 +604,9 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
 
             # set the logger to the flow run logger
 
-            self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
+            self.logger: "logging.Logger" = flow_run_logger(
+                flow_run=self.flow_run, flow=self.flow
+            )  # type: ignore
 
             # update the flow run name if necessary
             if not self._flow_run_name_set and self.flow.flow_run_name:
@@ -643,65 +649,68 @@ class FlowRunEngine(BaseFlowRunEngine[P, R]):
         """
         Enters a client context and creates a flow run if needed.
         """
-        with SyncClientContext.get_or_create() as client_ctx:
-            self._client = client_ctx.client
-            self._is_started = True
+        with hydrated_context(self.context):
+            with SyncClientContext.get_or_create() as client_ctx:
+                self._client = client_ctx.client
+                self._is_started = True
 
-            if not self.flow_run:
-                self.flow_run = self.create_flow_run(self.client)
-            else:
-                # Update the empirical policy to match the flow if it is not set
-                if self.flow_run.empirical_policy.retry_delay is None:
-                    self.flow_run.empirical_policy.retry_delay = (
-                        self.flow.retry_delay_seconds
-                    )
+                if not self.flow_run:
+                    self.flow_run = self.create_flow_run(self.client)
+                else:
+                    # Update the empirical policy to match the flow if it is not set
+                    if self.flow_run.empirical_policy.retry_delay is None:
+                        self.flow_run.empirical_policy.retry_delay = (
+                            self.flow.retry_delay_seconds
+                        )
 
-                if self.flow_run.empirical_policy.retries is None:
-                    self.flow_run.empirical_policy.retries = self.flow.retries
+                    if self.flow_run.empirical_policy.retries is None:
+                        self.flow_run.empirical_policy.retries = self.flow.retries
 
-                self.client.update_flow_run(
-                    flow_run_id=self.flow_run.id,
-                    flow_version=self.flow.version,
-                    empirical_policy=self.flow_run.empirical_policy,
-                )
+                    self.client.update_flow_run(
+                        flow_run_id=self.flow_run.id,
+                        flow_version=self.flow.version,
+                        empirical_policy=self.flow_run.empirical_policy,
+                    )
 
-            self._telemetry.start_span(
-                run=self.flow_run,
-                client=self.client,
-                parameters=self.parameters,
-            )
+                self._telemetry.start_span(
+                    run=self.flow_run,
+                    client=self.client,
+                    parameters=self.parameters,
+                )
 
-            try:
-                yield self
+                try:
+                    yield self
 
-            except TerminationSignal as exc:
-                self.cancel_all_tasks()
-                self.handle_crash(exc)
-                raise
-            except Exception:
-                # regular exceptions are caught and re-raised to the user
-                raise
-            except (Abort, Pause):
-                raise
-            except GeneratorExit:
-                # Do not capture generator exits as crashes
-                raise
-            except BaseException as exc:
-                # BaseExceptions are caught and handled as crashes
-                self.handle_crash(exc)
-                raise
-            finally:
-                # If debugging, use the more complete `repr` than the usual `str` description
-                display_state = (
-                    repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
-                )
-                self.logger.log(
-                    level=logging.INFO if self.state.is_completed() else logging.ERROR,
-                    msg=f"Finished in state {display_state}",
-                )
+                except TerminationSignal as exc:
+                    self.cancel_all_tasks()
+                    self.handle_crash(exc)
+                    raise
+                except Exception:
+                    # regular exceptions are caught and re-raised to the user
+                    raise
+                except (Abort, Pause):
+                    raise
+                except GeneratorExit:
+                    # Do not capture generator exits as crashes
+                    raise
+                except BaseException as exc:
+                    # BaseExceptions are caught and handled as crashes
+                    self.handle_crash(exc)
+                    raise
+                finally:
+                    # If debugging, use the more complete `repr` than the usual `str` description
+                    display_state = (
+                        repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
+                    )
+                    self.logger.log(
+                        level=logging.INFO
+                        if self.state.is_completed()
+                        else logging.ERROR,
+                        msg=f"Finished in state {display_state}",
+                    )
 
-                self._is_started = False
-                self._client = None
+                    self._is_started = False
+                    self._client = None
 
     # --------------------------
     #
@@ -768,6 +777,8 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
     """
 
     _client: Optional[PrefectClient] = None
+    parameters: dict[str, Any] | None = None
+    flow_run: FlowRun | None = None
 
     @property
     def client(self) -> PrefectClient:
@@ -1061,7 +1072,7 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             tags=TagsContext.get().current_tags,
         )
 
-    async def call_hooks(self, state: Optional[State] = None):
+    async def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         flow = self.flow
@@ -1158,7 +1169,9 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
             stack.enter_context(ConcurrencyContext())
 
             # set the logger to the flow run logger
-            self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
+            self.logger: "logging.Logger" = flow_run_logger(
+                flow_run=self.flow_run, flow=self.flow
+            )
 
             # update the flow run name if necessary
 
@@ -1200,71 +1213,74 @@ class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
         """
         Enters a client context and creates a flow run if needed.
         """
-        async with AsyncClientContext.get_or_create() as client_ctx:
-            self._client = client_ctx.client
-            self._is_started = True
+        with hydrated_context(self.context):
+            async with AsyncClientContext.get_or_create() as client_ctx:
+                self._client = client_ctx.client
+                self._is_started = True
+
+                if not self.flow_run:
+                    self.flow_run = await self.create_flow_run(self.client)
+                    flow_run_url = url_for(self.flow_run)
+
+                    if flow_run_url:
+                        self.logger.info(
+                            f"View at {flow_run_url}", extra={"send_to_api": False}
+                        )
+                else:
+                    # Update the empirical policy to match the flow if it is not set
+                    if self.flow_run.empirical_policy.retry_delay is None:
+                        self.flow_run.empirical_policy.retry_delay = (
+                            self.flow.retry_delay_seconds
+                        )
 
-            if not self.flow_run:
-                self.flow_run = await self.create_flow_run(self.client)
-                flow_run_url = url_for(self.flow_run)
+                    if self.flow_run.empirical_policy.retries is None:
+                        self.flow_run.empirical_policy.retries = self.flow.retries
 
-                if flow_run_url:
-                    self.logger.info(
-                        f"View at {flow_run_url}", extra={"send_to_api": False}
-                    )
-            else:
-                # Update the empirical policy to match the flow if it is not set
-                if self.flow_run.empirical_policy.retry_delay is None:
-                    self.flow_run.empirical_policy.retry_delay = (
-                        self.flow.retry_delay_seconds
+                    await self.client.update_flow_run(
+                        flow_run_id=self.flow_run.id,
+                        flow_version=self.flow.version,
+                        empirical_policy=self.flow_run.empirical_policy,
                     )
 
-                if self.flow_run.empirical_policy.retries is None:
-                    self.flow_run.empirical_policy.retries = self.flow.retries
-
-                await self.client.update_flow_run(
-                    flow_run_id=self.flow_run.id,
-                    flow_version=self.flow.version,
-                    empirical_policy=self.flow_run.empirical_policy,
+                await self._telemetry.async_start_span(
+                    run=self.flow_run,
+                    client=self.client,
+                    parameters=self.parameters,
                 )
 
-            await self._telemetry.async_start_span(
-                run=self.flow_run,
-                client=self.client,
-                parameters=self.parameters,
-            )
-
-            try:
-                yield self
+                try:
+                    yield self
 
-            except TerminationSignal as exc:
-                self.cancel_all_tasks()
-                await self.handle_crash(exc)
-                raise
-            except Exception:
-                # regular exceptions are caught and re-raised to the user
-                raise
-            except (Abort, Pause):
-                raise
-            except GeneratorExit:
-                # Do not capture generator exits as crashes
-                raise
-            except BaseException as exc:
-                # BaseExceptions are caught and handled as crashes
-                await self.handle_crash(exc)
-                raise
-            finally:
-                # If debugging, use the more complete `repr` than the usual `str` description
-                display_state = (
-                    repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
-                )
-                self.logger.log(
-                    level=logging.INFO if self.state.is_completed() else logging.ERROR,
-                    msg=f"Finished in state {display_state}",
-                )
+                except TerminationSignal as exc:
+                    self.cancel_all_tasks()
+                    await self.handle_crash(exc)
+                    raise
+                except Exception:
+                    # regular exceptions are caught and re-raised to the user
+                    raise
+                except (Abort, Pause):
+                    raise
+                except GeneratorExit:
+                    # Do not capture generator exits as crashes
+                    raise
+                except BaseException as exc:
+                    # BaseExceptions are caught and handled as crashes
+                    await self.handle_crash(exc)
+                    raise
+                finally:
+                    # If debugging, use the more complete `repr` than the usual `str` description
+                    display_state = (
+                        repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
+                    )
+                    self.logger.log(
+                        level=logging.INFO
+                        if self.state.is_completed()
+                        else logging.ERROR,
+                        msg=f"Finished in state {display_state}",
+                    )
 
-                self._is_started = False
-                self._client = None
+                    self._is_started = False
+                    self._client = None
 
     # --------------------------
     #
@@ -1320,14 +1336,16 @@ def run_flow_sync(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
     return_type: Literal["state", "result"] = "result",
+    context: Optional[dict[str, Any]] = None,
 ) -> Union[R, State, None]:
     engine = FlowRunEngine[P, R](
         flow=flow,
         parameters=parameters,
         flow_run=flow_run,
         wait_for=wait_for,
+        context=context,
     )
 
     with engine.start():
@@ -1342,11 +1360,16 @@ async def run_flow_async(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
     return_type: Literal["state", "result"] = "result",
+    context: Optional[dict[str, Any]] = None,
 ) -> Union[R, State, None]:
     engine = AsyncFlowRunEngine[P, R](
-        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+        flow=flow,
+        parameters=parameters,
+        flow_run=flow_run,
+        wait_for=wait_for,
+        context=context,
     )
 
     async with engine.start():
@@ -1361,14 +1384,19 @@ def run_generator_flow_sync(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
     return_type: Literal["state", "result"] = "result",
+    context: Optional[dict[str, Any]] = None,
 ) -> Generator[R, None, None]:
     if return_type != "result":
         raise ValueError("The return_type for a generator flow must be 'result'")
 
     engine = FlowRunEngine[P, R](
-        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+        flow=flow,
+        parameters=parameters,
+        flow_run=flow_run,
+        wait_for=wait_for,
+        context=context,
     )
 
     with engine.start():
@@ -1399,12 +1427,17 @@ async def run_generator_flow_async(
     parameters: Optional[Dict[str, Any]] = None,
     wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
+    context: Optional[dict[str, Any]] = None,
 ) -> AsyncGenerator[R, None]:
     if return_type != "result":
         raise ValueError("The return_type for a generator flow must be 'result'")
 
     engine = AsyncFlowRunEngine[P, R](
-        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+        flow=flow,
+        parameters=parameters,
+        flow_run=flow_run,
+        wait_for=wait_for,
+        context=context,
     )
 
     async with engine.start():
@@ -1438,8 +1471,23 @@ def run_flow(
     wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
     error_logger: Optional[logging.Logger] = None,
-) -> Union[R, State, None]:
-    ret_val: Union[R, State, None] = None
+    context: Optional[dict[str, Any]] = None,
+) -> (
+    R
+    | State
+    | None
+    | Coroutine[Any, Any, R | State | None]
+    | Generator[R, None, None]
+    | AsyncGenerator[R, None]
+):
+    ret_val: Union[
+        R,
+        State,
+        None,
+        Coroutine[Any, Any, R | State | None],
+        Generator[R, None, None],
+        AsyncGenerator[R, None],
+    ] = None
 
     try:
         kwargs: dict[str, Any] = dict(
@@ -1450,6 +1498,7 @@ def run_flow(
             ),
             wait_for=wait_for,
             return_type=return_type,
+            context=context,
         )
 
         if flow.isasync and flow.isgenerator: