prefect-client 3.1.13__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
@@ -647,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
 
     # --------------------------
     #
@@ -1208,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
 
     # --------------------------
     #
@@ -1330,12 +1338,14 @@ def run_flow_sync(
     parameters: Optional[Dict[str, Any]] = 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():
@@ -1352,9 +1362,14 @@ async def run_flow_async(
     parameters: Optional[Dict[str, Any]] = 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():
@@ -1371,12 +1386,17 @@ def run_generator_flow_sync(
     parameters: Optional[Dict[str, Any]] = 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():
@@ -1407,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():
@@ -1446,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(
@@ -1458,6 +1498,7 @@ def run_flow(
             ),
             wait_for=wait_for,
             return_type=return_type,
+            context=context,
         )
 
         if flow.isasync and flow.isgenerator:

prefect/flows.py

@@ -127,9 +127,9 @@ if TYPE_CHECKING:
     import logging
 
     from prefect.client.orchestration import PrefectClient
+    from prefect.client.schemas.objects import FlowRun
     from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
     from prefect.deployments.runner import RunnerDeployment
-    from prefect.flows import FlowRun
     from prefect.runner.storage import RunnerStorage
 
 logger: "logging.Logger" = get_logger("flows")

prefect/logging/configuration.py

@@ -94,12 +94,9 @@ def setup_logging(incremental: Optional[bool] = None) -> dict[str, Any]:
 
     for logger_name in PREFECT_LOGGING_EXTRA_LOGGERS.value():
         logger = logging.getLogger(logger_name)
-        for handler in extra_config.handlers:
-            if not config["incremental"]:
+        if not config["incremental"]:
+            for handler in extra_config.handlers:
                 logger.addHandler(handler)
-            if logger.level == logging.NOTSET:
-                logger.setLevel(extra_config.level)
-            logger.propagate = extra_config.propagate
 
     PROCESS_LOGGING_CONFIG = config
 

prefect/logging/loggers.py

@@ -24,7 +24,6 @@ else:
         LoggingAdapter = logging.LoggerAdapter
 
 if TYPE_CHECKING:
-    from prefect.client.schemas import FlowRun as ClientFlowRun
     from prefect.client.schemas.objects import FlowRun, TaskRun
     from prefect.context import RunContext
     from prefect.flows import Flow
@@ -164,7 +163,7 @@ def get_run_logger(
 
 
 def flow_run_logger(
-    flow_run: Union["FlowRun", "ClientFlowRun"],
+    flow_run: "FlowRun",
     flow: Optional["Flow[Any, Any]"] = None,
     **kwargs: str,
 ) -> PrefectLogAdapter: