prefect-client 3.0.0rc11__tar.gz → 3.0.0rc12__tar.gz

{prefect-client-3.0.0rc11/src/prefect_client.egg-info → prefect-client-3.0.0rc12}/PKG-INFO

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

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/requirements-client.txt

@@ -3,7 +3,7 @@ asgi-lifespan >= 1.0, < 3.0
 cachetools >= 5.3, < 6.0
 cloudpickle >= 2.0, < 4.0
 coolname >= 1.0.4, < 3.0.0
-croniter >= 1.0.12, < 3.0.0
+croniter >= 1.0.12, < 4.0.0
 exceptiongroup >= 1.0.0
 fastapi >= 0.111.0, < 1.0.0
 fsspec >= 2022.5.0

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/requirements-dev.txt

@@ -1,7 +1,6 @@
 ruff
 cairosvg
 codespell>=2.2.6
-ddtrace
 ipython
 jinja2
 moto >= 5

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/_internal/concurrency/services.py

@@ -151,6 +151,7 @@ class QueueService(abc.ABC, Generic[T]):
 
             if item is None:
                 logger.debug("Exiting service %r", self)
+                self._queue.task_done()
                 break
 
             try:
@@ -164,6 +165,8 @@ class QueueService(abc.ABC, Generic[T]):
                     item,
                     exc_info=log_traceback,
                 )
+            finally:
+                self._queue.task_done()
 
     @abc.abstractmethod
     async def _handle(self, item: T):
@@ -235,6 +238,12 @@ class QueueService(abc.ABC, Generic[T]):
         else:
             return concurrent.futures.wait(futures, timeout=timeout)
 
+    def wait_until_empty(self):
+        """
+        Wait until the queue is empty and all items have been processed.
+        """
+        self._queue.join()
+
     @classmethod
     def instance(cls: Type[Self], *args) -> Self:
         """

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/artifacts.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import json  # noqa: I001
 import math
+import warnings
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID
 
@@ -54,8 +55,19 @@ class Artifact(ArtifactRequest):
         Returns:
             - The created artifact.
         """
+        from prefect.context import MissingContextError, get_run_context
+
         client, _ = get_or_create_client(client)
         task_run_id, flow_run_id = get_task_and_flow_run_ids()
+
+        try:
+            get_run_context()
+        except MissingContextError:
+            warnings.warn(
+                "Artifact creation outside of a flow or task run is deprecated and will be removed in a later version.",
+                FutureWarning,
+            )
+
         return await client.create_artifact(
             artifact=ArtifactRequest(
                 type=self.type,

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/client/schemas/actions.py

@@ -377,6 +377,10 @@ class DeploymentFlowRunCreate(ActionBaseModel):
     parameters: Dict[str, Any] = Field(
         default_factory=dict, description="The parameters for the flow run."
     )
+    enforce_parameter_schema: Optional[bool] = Field(
+        default=None,
+        description="Whether or not to enforce the parameter schema on this run.",
+    )
     context: Dict[str, Any] = Field(
         default_factory=dict, description="The context for the flow run."
     )

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/events/schemas/events.py

@@ -60,6 +60,16 @@ class Resource(Labelled):
     def name(self) -> Optional[str]:
         return self.get("prefect.resource.name")
 
+    def prefect_object_id(self, kind: str) -> UUID:
+        """Extracts the UUID from an event's resource ID if it's the expected kind
+        of prefect resource"""
+        prefix = f"{kind}." if not kind.endswith(".") else kind
+
+        if not self.id.startswith(prefix):
+            raise ValueError(f"Resource ID {self.id} does not start with {prefix}")
+
+        return UUID(self.id[len(prefix) :])
+
 
 class RelatedResource(Resource):
     """A Resource with a specific role in an Event"""

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/flow_engine.py

@@ -91,9 +91,12 @@ def load_flow_and_flow_run(flow_run_id: UUID) -> Tuple[FlowRun, Flow]:
 
     flow_run = client.read_flow_run(flow_run_id)
     if entrypoint:
-        flow = load_flow_from_entrypoint(entrypoint)
+        # we should not accept a placeholder flow at runtime
+        flow = load_flow_from_entrypoint(entrypoint, use_placeholder_flow=False)
     else:
-        flow = run_coro_as_sync(load_flow_from_flow_run(flow_run))
+        flow = run_coro_as_sync(
+            load_flow_from_flow_run(flow_run, use_placeholder_flow=False)
+        )
 
     return flow_run, flow
 

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/flows.py

@@ -798,7 +798,7 @@ class Flow(Generic[P, R]):
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[List["FlexibleScheduleList"]] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
         schedule: Optional[SCHEDULE_TYPES] = None,
         is_schedule_active: Optional[bool] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
@@ -1704,6 +1704,7 @@ def select_flow(
 
 def load_flow_from_entrypoint(
     entrypoint: str,
+    use_placeholder_flow: bool = True,
 ) -> Flow:
     """
     Extract a flow object from a script at an entrypoint by running all of the code in the file.
@@ -1711,6 +1712,8 @@ def load_flow_from_entrypoint(
     Args:
         entrypoint: a string in the format `<path_to_script>:<flow_func_name>` or a module path
             to a flow function
+        use_placeholder_flow: if True, use a placeholder Flow object if the actual flow object
+            cannot be loaded from the entrypoint (e.g. dependencies are missing)
 
     Returns:
         The flow object from the script
@@ -1737,8 +1740,10 @@ def load_flow_from_entrypoint(
         # drawback of this approach is that we're unable to actually load the
         # function, so we create a placeholder flow that will re-raise this
         # exception when called.
-
-        flow = load_placeholder_flow(entrypoint=entrypoint, raises=exc)
+        if use_placeholder_flow:
+            flow = load_placeholder_flow(entrypoint=entrypoint, raises=exc)
+        else:
+            raise
 
     if not isinstance(flow, Flow):
         raise MissingFlowError(
@@ -1856,6 +1861,7 @@ async def load_flow_from_flow_run(
     flow_run: "FlowRun",
     ignore_storage: bool = False,
     storage_base_path: Optional[str] = None,
+    use_placeholder_flow: bool = True,
 ) -> Flow:
     """
     Load a flow from the location/script provided in a deployment's storage document.
@@ -1882,7 +1888,9 @@ async def load_flow_from_flow_run(
             f"Importing flow code from module path {deployment.entrypoint}"
         )
         flow = await run_sync_in_worker_thread(
-            load_flow_from_entrypoint, deployment.entrypoint
+            load_flow_from_entrypoint,
+            deployment.entrypoint,
+            use_placeholder_flow=use_placeholder_flow,
         )
         return flow
 
@@ -1924,7 +1932,11 @@ async def load_flow_from_flow_run(
     import_path = relative_path_to_current_platform(deployment.entrypoint)
     run_logger.debug(f"Importing flow code from '{import_path}'")
 
-    flow = await run_sync_in_worker_thread(load_flow_from_entrypoint, str(import_path))
+    flow = await run_sync_in_worker_thread(
+        load_flow_from_entrypoint,
+        str(import_path),
+        use_placeholder_flow=use_placeholder_flow,
+    )
 
     return flow
 

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/results.py

@@ -25,7 +25,7 @@ from typing_extensions import ParamSpec, Self
 import prefect
 from prefect.blocks.core import Block
 from prefect.client.utilities import inject_client
-from prefect.exceptions import MissingResult, ObjectAlreadyExists
+from prefect.exceptions import MissingResult
 from prefect.filesystems import (
     LocalFileSystem,
     WritableFileSystem,
@@ -64,51 +64,6 @@ R = TypeVar("R")
 _default_storages: Dict[Tuple[str, str], WritableFileSystem] = {}
 
 
-async def _get_or_create_default_storage(block_document_slug: str) -> ResultStorage:
-    """
-    Generate a default file system for storage.
-    """
-    default_storage_name, storage_path = cache_key = (
-        block_document_slug,
-        PREFECT_LOCAL_STORAGE_PATH.value(),
-    )
-
-    async def get_storage() -> WritableFileSystem:
-        try:
-            return await Block.load(default_storage_name)
-        except ValueError as e:
-            if "Unable to find" not in str(e):
-                raise e
-
-        block_type_slug, name = default_storage_name.split("/")
-        if block_type_slug == "local-file-system":
-            block = LocalFileSystem(basepath=storage_path)
-        else:
-            raise ValueError(
-                "The default storage block does not exist, but it is of type "
-                f"'{block_type_slug}' which cannot be created implicitly.  Please create "
-                "the block manually."
-            )
-
-        try:
-            await block.save(name, overwrite=False)
-        except ValueError as e:
-            if "already in use" not in str(e):
-                raise e
-        except ObjectAlreadyExists:
-            # Another client created the block before we reached this line
-            block = await Block.load(default_storage_name)
-
-        return block
-
-    try:
-        return _default_storages[cache_key]
-    except KeyError:
-        storage = await get_storage()
-        _default_storages[cache_key] = storage
-        return storage
-
-
 @sync_compatible
 async def get_default_result_storage() -> ResultStorage:
     """

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/runner/runner.py

@@ -167,9 +167,7 @@ class Runner:
         self.query_seconds = query_seconds or PREFECT_RUNNER_POLL_FREQUENCY.value()
         self._prefetch_seconds = prefetch_seconds
 
-        self._limiter: Optional[anyio.CapacityLimiter] = anyio.CapacityLimiter(
-            self.limit
-        )
+        self._limiter: Optional[anyio.CapacityLimiter] = None
         self._client = get_client()
         self._submitting_flow_run_ids = set()
         self._cancelling_flow_run_ids = set()
@@ -1227,6 +1225,8 @@ class Runner:
         self._client = get_client()
         self._tmp_dir.mkdir(parents=True)
 
+        self._limiter = anyio.CapacityLimiter(self.limit)
+
         if not hasattr(self, "_loop") or not self._loop:
             self._loop = asyncio.get_event_loop()
 

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/settings.py

@@ -481,18 +481,6 @@ PREFECT_HOME = Setting(
 directory may be created automatically when required.
 """
 
-PREFECT_EXTRA_ENTRYPOINTS = Setting(
-    str,
-    default="",
-)
-"""
-Modules for Prefect to import when Prefect is imported.
-
-Values should be separated by commas, e.g. `my_module,my_other_module`.
-Objects within modules may be specified by a ':' partition, e.g. `my_module:my_object`.
-If a callable object is provided, it will be called with no arguments on import.
-"""
-
 PREFECT_DEBUG_MODE = Setting(
     bool,
     default=False,
@@ -1438,11 +1426,6 @@ a task worker should move a task from PENDING to RUNNING very quickly, so runs s
 PENDING for a while is a sign that the task worker may have crashed.
 """
 
-PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT = Setting(bool, default=False)
-"""
-Whether or not to disable the sync_compatible decorator utility.
-"""
-
 PREFECT_EXPERIMENTAL_ENABLE_SCHEDULE_CONCURRENCY = Setting(bool, default=False)
 
 # Defaults -----------------------------------------------------------------------------

{prefect-client-3.0.0rc11 → prefect-client-3.0.0rc12}/src/prefect/task_engine.py

@@ -5,6 +5,7 @@ import time
 from asyncio import CancelledError
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
+from functools import wraps
 from textwrap import dedent
 from typing import (
     Any,
@@ -53,6 +54,7 @@ from prefect.records.result_store import ResultFactoryStore
 from prefect.results import BaseResult, ResultFactory, _format_user_supplied_storage_key
 from prefect.settings import (
     PREFECT_DEBUG_MODE,
+    PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION,
     PREFECT_TASKS_REFRESH_CACHE,
 )
 from prefect.states import (
@@ -124,8 +126,7 @@ class TaskRunEngine(Generic[P, R]):
             raise ValueError("Task run is not set")
         return self.task_run.state
 
-    @property
-    def can_retry(self) -> bool:
+    def can_retry(self, exc: Exception) -> bool:
         retry_condition: Optional[
             Callable[[Task[P, Coroutine[Any, Any, R]], TaskRun, State], bool]
         ] = self.task.retry_condition_fn
@@ -136,9 +137,19 @@ class TaskRunEngine(Generic[P, R]):
                 f"Running `retry_condition_fn` check {retry_condition!r} for task"
                 f" {self.task.name!r}"
             )
-            return not retry_condition or retry_condition(
-                self.task, self.task_run, self.state
+            state = Failed(
+                data=exc,
+                message=f"Task run encountered unexpected exception: {repr(exc)}",
             )
+            if inspect.iscoroutinefunction(retry_condition):
+                should_retry = run_coro_as_sync(
+                    retry_condition(self.task, self.task_run, state)
+                )
+            elif inspect.isfunction(retry_condition):
+                should_retry = retry_condition(self.task, self.task_run, state)
+            else:
+                should_retry = not retry_condition
+            return should_retry
         except Exception:
             self.logger.error(
                 (
@@ -269,6 +280,17 @@ class TaskRunEngine(Generic[P, R]):
             return
 
         new_state = Running()
+
+        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            self.task_run.start_time = new_state.timestamp
+            self.task_run.run_count += 1
+
+            flow_run_context = FlowRunContext.get()
+            if flow_run_context:
+                # Carry forward any task run information from the flow run
+                flow_run = flow_run_context.flow_run
+                self.task_run.flow_run_run_count = flow_run.run_count
+
         state = self.set_state(new_state)
 
         # TODO: this is temporary until the API stops rejecting state transitions
@@ -298,24 +320,37 @@ class TaskRunEngine(Generic[P, R]):
         last_state = self.state
         if not self.task_run:
             raise ValueError("Task run is not set")
-        try:
-            new_state = propose_state_sync(
-                self.client, state, task_run_id=self.task_run.id, force=force
-            )
-        except Pause as exc:
-            # We shouldn't get a pause signal without a state, but if this happens,
-            # just use a Paused state to assume an in-process pause.
-            new_state = exc.state if exc.state else Paused()
-            if new_state.state_details.pause_reschedule:
-                # If we're being asked to pause and reschedule, we should exit the
-                # task and expect to be resumed later.
-                raise
 
-        # currently this is a hack to keep a reference to the state object
-        # that has an in-memory result attached to it; using the API state
+        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            self.task_run.state = new_state = state
+
+            # Ensure that the state_details are populated with the current run IDs
+            new_state.state_details.task_run_id = self.task_run.id
+            new_state.state_details.flow_run_id = self.task_run.flow_run_id
+
+            # Predictively update the de-normalized task_run.state_* attributes
+            self.task_run.state_id = new_state.id
+            self.task_run.state_type = new_state.type
+            self.task_run.state_name = new_state.name
+        else:
+            try:
+                new_state = propose_state_sync(
+                    self.client, state, task_run_id=self.task_run.id, force=force
+                )
+            except Pause as exc:
+                # We shouldn't get a pause signal without a state, but if this happens,
+                # just use a Paused state to assume an in-process pause.
+                new_state = exc.state if exc.state else Paused()
+                if new_state.state_details.pause_reschedule:
+                    # If we're being asked to pause and reschedule, we should exit the
+                    # task and expect to be resumed later.
+                    raise
+
+            # currently this is a hack to keep a reference to the state object
+            # that has an in-memory result attached to it; using the API state
+            # could result in losing that reference
+            self.task_run.state = new_state
 
-        # could result in losing that reference
-        self.task_run.state = new_state
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
             task_run=self.task_run,
@@ -323,6 +358,7 @@ class TaskRunEngine(Generic[P, R]):
             validated_state=self.task_run.state,
             follows=self._last_event,
         )
+
         return new_state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -367,11 +403,19 @@ class TaskRunEngine(Generic[P, R]):
         )
         transaction.stage(
             terminal_state.data,
-            on_rollback_hooks=self.task.on_rollback_hooks,
-            on_commit_hooks=self.task.on_commit_hooks,
+            on_rollback_hooks=[
+                _with_transaction_hook_logging(hook, "rollback", self.logger)
+                for hook in self.task.on_rollback_hooks
+            ],
+            on_commit_hooks=[
+                _with_transaction_hook_logging(hook, "commit", self.logger)
+                for hook in self.task.on_commit_hooks
+            ],
         )
         if transaction.is_committed():
             terminal_state.name = "Cached"
+
+        self.record_terminal_state_timing(terminal_state)
         self.set_state(terminal_state)
         self._return_value = result
         return result
@@ -383,7 +427,7 @@ class TaskRunEngine(Generic[P, R]):
           - If the task has a retry delay, place in AwaitingRetry state with a delayed scheduled time.
         - If the task has no retries left, or the retry condition is not met, return False.
         """
-        if self.retries < self.task.retries and self.can_retry:
+        if self.retries < self.task.retries and self.can_retry(exc):
             if self.task.retry_delay_seconds:
                 delay = (
                     self.task.retry_delay_seconds[
@@ -398,6 +442,8 @@ class TaskRunEngine(Generic[P, R]):
             else:
                 delay = None
                 new_state = Retrying()
+                if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+                    self.task_run.run_count += 1
 
             self.logger.info(
                 "Task run failed with exception: %r - " "Retry %s/%s will start %s",
@@ -432,6 +478,7 @@ class TaskRunEngine(Generic[P, R]):
                     result_factory=getattr(context, "result_factory", None),
                 )
             )
+            self.record_terminal_state_timing(state)
             self.set_state(state)
             self._raised = exc
 
@@ -454,9 +501,20 @@ class TaskRunEngine(Generic[P, R]):
         state = run_coro_as_sync(exception_to_crashed_state(exc))
         self.logger.error(f"Crash detected! {state.message}")
         self.logger.debug("Crash details:", exc_info=exc)
+        self.record_terminal_state_timing(state)
         self.set_state(state, force=True)
         self._raised = exc
 
+    def record_terminal_state_timing(self, state: State) -> None:
+        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            if self.task_run.start_time and not self.task_run.end_time:
+                self.task_run.end_time = state.timestamp
+
+                if self.task_run.state.is_running():
+                    self.task_run.total_run_time += (
+                        state.timestamp - self.task_run.state.timestamp
+                    )
+
     @contextmanager
     def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
         from prefect.utilities.engine import (
@@ -469,7 +527,8 @@ class TaskRunEngine(Generic[P, R]):
         if not self.task_run:
             raise ValueError("Task run is not set")
 
-        self.task_run = client.read_task_run(self.task_run.id)
+        if not PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            self.task_run = client.read_task_run(self.task_run.id)
         with ExitStack() as stack:
             if log_prints := should_log_prints(self.task):
                 stack.enter_context(patch_print())
@@ -483,23 +542,24 @@ class TaskRunEngine(Generic[P, R]):
                     client=client,
                 )
             )
-            # set the logger to the task run logger
+
             self.logger = task_run_logger(task_run=self.task_run, task=self.task)  # type: ignore
 
-            # update the task run name if necessary
-            if not self._task_name_set and self.task.task_run_name:
-                task_run_name = _resolve_custom_task_run_name(
-                    task=self.task, parameters=self.parameters
-                )
-                self.client.set_task_run_name(
-                    task_run_id=self.task_run.id, name=task_run_name
-                )
-                self.logger.extra["task_run_name"] = task_run_name
-                self.logger.debug(
-                    f"Renamed task run {self.task_run.name!r} to {task_run_name!r}"
-                )
-                self.task_run.name = task_run_name
-                self._task_name_set = True
+            if not PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+                # update the task run name if necessary
+                if not self._task_name_set and self.task.task_run_name:
+                    task_run_name = _resolve_custom_task_run_name(
+                        task=self.task, parameters=self.parameters
+                    )
+                    self.client.set_task_run_name(
+                        task_run_id=self.task_run.id, name=task_run_name
+                    )
+                    self.logger.extra["task_run_name"] = task_run_name
+                    self.logger.debug(
+                        f"Renamed task run {self.task_run.name!r} to {task_run_name!r}"
+                    )
+                    self.task_run.name = task_run_name
+                    self._task_name_set = True
             yield
 
     @contextmanager
@@ -511,22 +571,47 @@ class TaskRunEngine(Generic[P, R]):
         """
         Enters a client context and creates a task run if needed.
         """
+
         with hydrated_context(self.context):
             with ClientContext.get_or_create() as client_ctx:
                 self._client = client_ctx.sync_client
                 self._is_started = True
                 try:
                     if not self.task_run:
-                        self.task_run = run_coro_as_sync(
-                            self.task.create_run(
-                                id=task_run_id,
-                                parameters=self.parameters,
-                                flow_run_context=FlowRunContext.get(),
-                                parent_task_run_context=TaskRunContext.get(),
-                                wait_for=self.wait_for,
-                                extra_task_inputs=dependencies,
+                        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+                            # TODO - this maybe should be a method on Task?
+                            from prefect.utilities.engine import (
+                                _resolve_custom_task_run_name,
+                            )
+
+                            task_run_name = None
+                            if not self._task_name_set and self.task.task_run_name:
+                                task_run_name = _resolve_custom_task_run_name(
+                                    task=self.task, parameters=self.parameters
+                                )
+
+                            self.task_run = run_coro_as_sync(
+                                self.task.create_local_run(
+                                    id=task_run_id,
+                                    parameters=self.parameters,
+                                    flow_run_context=FlowRunContext.get(),
+                                    parent_task_run_context=TaskRunContext.get(),
+                                    wait_for=self.wait_for,
+                                    extra_task_inputs=dependencies,
+                                    task_run_name=task_run_name,
+                                )
+                            )
+                        else:
+                            self.task_run = run_coro_as_sync(
+                                self.task.create_run(
+                                    id=task_run_id,
+                                    parameters=self.parameters,
+                                    flow_run_context=FlowRunContext.get(),
+                                    parent_task_run_context=TaskRunContext.get(),
+                                    wait_for=self.wait_for,
+                                    extra_task_inputs=dependencies,
+                                )
                             )
-                        )
                     # Emit an event to capture that the task run was in the `PENDING` state.
                     self._last_event = emit_task_run_state_change_event(
                         task_run=self.task_run,
@@ -916,3 +1001,28 @@ def run_task(
         return run_task_async(**kwargs)
     else:
         return run_task_sync(**kwargs)
+
+
+def _with_transaction_hook_logging(
+    hook: Callable[[Transaction], None],
+    hook_type: Literal["rollback", "commit"],
+    logger: logging.Logger,
+) -> Callable[[Transaction], None]:
+    @wraps(hook)
+    def _hook(txn: Transaction) -> None:
+        hook_name = _get_hook_name(hook)
+        logger.info(f"Running {hook_type} hook {hook_name!r}")
+
+        try:
+            hook(txn)
+        except Exception as exc:
+            logger.error(
+                f"An error was encountered while running {hook_type} hook {hook_name!r}",
+            )
+            raise exc
+        else:
+            logger.info(
+                f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
+            )
+
+    return _hook