prefect-client 3.3.5.dev1__py3-none-any.whl → 3.3.5.dev3__py3-none-any.whl

prefect/_build_info.py

@@ -1,5 +1,5 @@
 # Generated by versioningit
-__version__ = "3.3.5.dev1"
-__build_date__ = "2025-04-12 08:07:21.126697+00:00"
-__git_commit__ = "291bb48b175219a7d9d2a42bd9d09194c31350b6"
+__version__ = "3.3.5.dev3"
+__build_date__ = "2025-04-16 08:08:24.630955+00:00"
+__git_commit__ = "d01b30e8ec3cffd839c23ab42edc62983c3cb1a4"
 __dirty__ = False

prefect/events/clients.py

@@ -39,6 +39,7 @@ import prefect.types._datetime
 from prefect.events import Event
 from prefect.logging import get_logger
 from prefect.settings import (
+    PREFECT_API_AUTH_STRING,
     PREFECT_API_KEY,
     PREFECT_API_SSL_CERT_FILE,
     PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
@@ -367,7 +368,7 @@ class PrefectEventsClient(EventsClient):
         await self._connect.__aexit__(exc_type, exc_val, exc_tb)
         return await super().__aexit__(exc_type, exc_val, exc_tb)
 
-    def _log_debug(self, message: str, *args, **kwargs) -> None:
+    def _log_debug(self, message: str, *args: Any, **kwargs: Any) -> None:
         message = f"EventsClient(id={id(self)}): " + message
         logger.debug(message, *args, **kwargs)
 
@@ -578,6 +579,7 @@ class PrefectEventSubscriber:
     _seen_events: MutableMapping[UUID, bool]
 
     _api_key: Optional[str]
+    _auth_token: Optional[str]
 
     def __init__(
         self,
@@ -593,6 +595,8 @@ class PrefectEventSubscriber:
                 the client should attempt to reconnect
         """
         self._api_key = None
+        self._auth_token = PREFECT_API_AUTH_STRING.value()
+
         if not api_url:
             api_url = cast(str, PREFECT_API_URL.value())
 
@@ -641,8 +645,10 @@ class PrefectEventSubscriber:
         await pong
 
         logger.debug("  authenticating...")
+        # Use the API key (for Cloud) OR the auth token (for self-hosted with auth string)
+        token = self._api_key or self._auth_token
         await self._websocket.send(
-            orjson.dumps({"type": "auth", "token": self._api_key}).decode()
+            orjson.dumps({"type": "auth", "token": token}).decode()
         )
 
         try:
@@ -652,13 +658,13 @@ class PrefectEventSubscriber:
         except AssertionError as e:
             raise Exception(
                 "Unable to authenticate to the event stream. Please ensure the "
-                "provided api_key you are using is valid for this environment. "
+                "provided api_key or auth_token you are using is valid for this environment. "
                 f"Reason: {e.args[0]}"
             )
         except ConnectionClosedError as e:
             reason = getattr(e.rcvd, "reason", None)
             msg = "Unable to authenticate to the event stream. Please ensure the "
-            msg += "provided api_key you are using is valid for this environment. "
+            msg += "provided api_key or auth_token you are using is valid for this environment. "
             msg += f"Reason: {reason}" if reason else ""
             raise Exception(msg) from e
 

prefect/flows.py

@@ -1130,7 +1130,7 @@ class Flow(Generic[P, R]):
     @classmethod
     async def afrom_source(
         cls,
-        source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
+        source: Union[str, Path, "RunnerStorage", ReadableDeploymentStorage],
         entrypoint: str,
     ) -> "Flow[..., Any]":
         """
@@ -1244,7 +1244,7 @@ class Flow(Generic[P, R]):
     @async_dispatch(afrom_source)
     def from_source(
         cls,
-        source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
+        source: Union[str, Path, "RunnerStorage", ReadableDeploymentStorage],
         entrypoint: str,
     ) -> "Flow[..., Any]":
         """
@@ -1976,7 +1976,7 @@ class FlowDecorator:
         # manually here.
         @staticmethod
         def from_source(
-            source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
+            source: Union[str, Path, "RunnerStorage", ReadableDeploymentStorage],
             entrypoint: str,
         ) -> Union["Flow[..., Any]", Coroutine[Any, Any, "Flow[..., Any]"]]: ...
 

prefect/locking/memory.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import asyncio
 import threading
-from typing import Any, Optional, TypedDict
+from typing import Any, TypedDict
 
 from typing_extensions import Self
 
@@ -19,7 +21,7 @@ class _LockInfo(TypedDict):
 
     holder: str
     lock: threading.Lock
-    expiration_timer: Optional[threading.Timer]
+    expiration_timer: threading.Timer | None
 
 
 class MemoryLockManager(LockManager):
@@ -63,8 +65,8 @@ class MemoryLockManager(LockManager):
         self,
         key: str,
         holder: str,
-        acquire_timeout: Optional[float] = None,
-        hold_timeout: Optional[float] = None,
+        acquire_timeout: float | None = None,
+        hold_timeout: float | None = None,
     ) -> bool:
         with self._locks_dict_lock:
             if key not in self._locks:
@@ -116,13 +118,13 @@ class MemoryLockManager(LockManager):
         self,
         key: str,
         holder: str,
-        acquire_timeout: Optional[float] = None,
-        hold_timeout: Optional[float] = None,
+        acquire_timeout: float | None = None,
+        hold_timeout: float | None = None,
     ) -> bool:
         with self._locks_dict_lock:
             if key not in self._locks:
                 lock = threading.Lock()
-                await asyncio.to_thread(lock.acquire)
+                lock.acquire()
                 expiration_timer = None
                 if hold_timeout is not None:
                     expiration_timer = threading.Timer(
@@ -192,7 +194,7 @@ class MemoryLockManager(LockManager):
             and lock_info["holder"] == holder
         )
 
-    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+    def wait_for_lock(self, key: str, timeout: float | None = None) -> bool:
         lock_info: _LockInfo | None = self._locks.get(key)
         if lock_info is None:
             return True
@@ -206,7 +208,7 @@ class MemoryLockManager(LockManager):
             return lock_acquired
         return True
 
-    async def await_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+    async def await_for_lock(self, key: str, timeout: float | None = None) -> bool:
         lock_info: _LockInfo | None = self._locks.get(key, None)
         if lock_info is None:
             return True

prefect/results.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import inspect
 import os
 import socket
@@ -33,6 +34,7 @@ from typing_extensions import ParamSpec, Self
 import prefect
 import prefect.types._datetime
 from prefect._internal.compatibility.async_dispatch import async_dispatch
+from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect._result_records import R, ResultRecord, ResultRecordMetadata
 from prefect.blocks.core import Block
 from prefect.exceptions import (
@@ -475,10 +477,17 @@ class ResultStore(BaseModel):
         Returns:
             str: A unique identifier string.
         """
+        current_loop = get_running_loop()
         hostname = socket.gethostname()
         pid = os.getpid()
         thread_name = threading.current_thread().name
         thread_id = threading.get_ident()
+        if current_loop:
+            current_task = asyncio.current_task()
+            if current_task:
+                # include the task id to ensure uniqueness because there might be
+                # multiple tasks running in the same thread
+                return f"{hostname}:{pid}:{thread_id}:{thread_name}:{id(current_task)}"
         return f"{hostname}:{pid}:{thread_id}:{thread_name}"
 
     @sync_compatible

prefect/task_engine.py

@@ -81,7 +81,13 @@ from prefect.states import (
     return_value_to_state,
 )
 from prefect.telemetry.run_telemetry import RunTelemetry
-from prefect.transactions import IsolationLevel, Transaction, transaction
+from prefect.transactions import (
+    AsyncTransaction,
+    IsolationLevel,
+    Transaction,
+    atransaction,
+    transaction,
+)
 from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
@@ -1020,7 +1026,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # otherwise, return the exception
             return self._raised
 
-    async def handle_success(self, result: R, transaction: Transaction) -> R:
+    async def handle_success(self, result: R, transaction: AsyncTransaction) -> R:
         if self.task.cache_expiration is not None:
             expiration = prefect.types._datetime.now("UTC") + self.task.cache_expiration
         else:
@@ -1302,7 +1308,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                         await self.call_hooks()
 
     @asynccontextmanager
-    async def transaction_context(self) -> AsyncGenerator[Transaction, None]:
+    async def transaction_context(self) -> AsyncGenerator[AsyncTransaction, None]:
         # refresh cache setting is now repurposes as overwrite transaction record
         overwrite = (
             self.task.refresh_cache
@@ -1317,7 +1323,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else None
         )
 
-        with transaction(
+        async with atransaction(
             key=self.compute_transaction_key(),
             store=get_result_store(),
             overwrite=overwrite,
@@ -1349,7 +1355,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 await self.handle_exception(exc)
 
     async def call_task_fn(
-        self, transaction: Transaction
+        self, transaction: AsyncTransaction
     ) -> Union[R, Coroutine[Any, Any, R]]:
         """
         Convenience method to call the task function. Returns a coroutine if the
@@ -1357,7 +1363,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         """
         parameters = self.parameters or {}
         if transaction.is_committed():
-            result = transaction.read()
+            result = await transaction.read()
         else:
             result = await call_with_parameters(self.task.fn, parameters)
         await self.handle_success(result, transaction=transaction)

prefect/tasks.py

@@ -764,11 +764,6 @@ class Task(Generic[P, R]):
     def on_rollback(
         self, fn: Callable[["Transaction"], None]
     ) -> Callable[["Transaction"], None]:
-        if asyncio.iscoroutinefunction(fn):
-            raise ValueError(
-                "Asynchronous rollback hooks are not yet supported. Rollback hooks must be synchronous functions."
-            )
-
         self.on_rollback_hooks.append(fn)
         return fn
 

prefect/telemetry/instrumentation.py

@@ -1,6 +1,6 @@
 import logging
-import os
 import re
+import socket
 from typing import TYPE_CHECKING
 from urllib.parse import urljoin
 from uuid import UUID
@@ -74,7 +74,7 @@ def setup_exporters(
     resource = Resource.create(
         {
             "service.name": "prefect",
-            "service.instance.id": os.uname().nodename,
+            "service.instance.id": socket.gethostname(),
             "prefect.account": str(account_id),
             "prefect.workspace": str(workspace_id),
         }

prefect/transactions.py

@@ -1,19 +1,26 @@
+from __future__ import annotations
+
+import abc
+import asyncio
 import copy
+import inspect
 import logging
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
 from functools import partial
 from typing import (
     Any,
+    AsyncGenerator,
     Callable,
-    Dict,
+    ClassVar,
     Generator,
-    List,
+    NoReturn,
     Optional,
     Type,
     Union,
 )
 
+import anyio.to_thread
 from pydantic import Field, PrivateAttr
 from typing_extensions import Self
 
@@ -32,8 +39,11 @@ from prefect.results import (
 )
 from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
+from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
 
+logger: logging.Logger = get_logger("transactions")
+
 
 class IsolationLevel(AutoEnum):
     READ_COMMITTED = AutoEnum.auto()
@@ -54,29 +64,27 @@ class TransactionState(AutoEnum):
     ROLLED_BACK = AutoEnum.auto()
 
 
-class Transaction(ContextModel):
+class BaseTransaction(ContextModel, abc.ABC):
     """
     A base model for transaction state.
     """
 
     store: Optional[ResultStore] = None
     key: Optional[str] = None
-    children: List["Transaction"] = Field(default_factory=list)
+    children: list[Self] = Field(default_factory=list)
     commit_mode: Optional[CommitMode] = None
     isolation_level: Optional[IsolationLevel] = IsolationLevel.READ_COMMITTED
     state: TransactionState = TransactionState.PENDING
-    on_commit_hooks: List[Callable[["Transaction"], None]] = Field(default_factory=list)
-    on_rollback_hooks: List[Callable[["Transaction"], None]] = Field(
-        default_factory=list
-    )
+    on_commit_hooks: list[Callable[[Self], None]] = Field(default_factory=list)
+    on_rollback_hooks: list[Callable[[Self], None]] = Field(default_factory=list)
     overwrite: bool = False
     logger: Union[logging.Logger, LoggingAdapter] = Field(
         default_factory=partial(get_logger, "transactions")
     )
     write_on_commit: bool = True
-    _stored_values: Dict[str, Any] = PrivateAttr(default_factory=dict)
-    _staged_value: Any = None
-    __var__: ContextVar[Self] = ContextVar("transaction")
+    _stored_values: dict[str, Any] = PrivateAttr(default_factory=dict)
+    _staged_value: ResultRecord[Any] | Any = None
+    __var__: ClassVar[ContextVar[Self]] = ContextVar("transaction")
 
     def set(self, name: str, value: Any) -> None:
         """
@@ -174,7 +182,8 @@ class Transaction(ContextModel):
     def is_active(self) -> bool:
         return self.state == TransactionState.ACTIVE
 
-    def __enter__(self) -> Self:
+    def prepare_transaction(self) -> None:
+        """Helper method to prepare transaction state and validate configuration."""
         if self._token is not None:
             raise RuntimeError(
                 "Context already entered. Context enter calls cannot be nested."
@@ -203,6 +212,51 @@ class Transaction(ContextModel):
 
         # this needs to go before begin, which could set the state to committed
         self.state = TransactionState.ACTIVE
+
+    def add_child(self, transaction: Self) -> None:
+        self.children.append(transaction)
+
+    def get_parent(self) -> Self | None:
+        parent = None
+        if self._token:
+            prev_var = self._token.old_value
+            if prev_var != Token.MISSING:
+                parent = prev_var
+        else:
+            # `_token` has been reset so we need to get the active transaction from the context var
+            parent = self.get_active()
+        return parent
+
+    def stage(
+        self,
+        value: Any,
+        on_rollback_hooks: Optional[list[Callable[..., Any]]] = None,
+        on_commit_hooks: Optional[list[Callable[..., Any]]] = None,
+    ) -> None:
+        """
+        Stage a value to be committed later.
+        """
+        on_commit_hooks = on_commit_hooks or []
+        on_rollback_hooks = on_rollback_hooks or []
+
+        if self.state != TransactionState.COMMITTED:
+            self._staged_value = value
+            self.on_rollback_hooks += on_rollback_hooks
+            self.on_commit_hooks += on_commit_hooks
+            self.state = TransactionState.STAGED
+
+    @classmethod
+    def get_active(cls: Type[Self]) -> Optional[Self]:
+        return cls.__var__.get(None)
+
+
+class Transaction(BaseTransaction):
+    """
+    A model representing the state of a transaction.
+    """
+
+    def __enter__(self) -> Self:
+        self.prepare_transaction()
         self.begin()
         self._token = self.__var__.set(self)
         return self
@@ -252,11 +306,9 @@ class Transaction(ContextModel):
         ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> Optional[ResultRecord[Any]]:
+    def read(self) -> ResultRecord[Any] | None:
         if self.store and self.key:
-            record = self.store.read(key=self.key)
-            if isinstance(record, ResultRecord):
-                return record
+            return self.store.read(key=self.key)
         return None
 
     def reset(self) -> None:
@@ -274,20 +326,6 @@ class Transaction(ContextModel):
         if parent and self.state == TransactionState.ROLLED_BACK:
             parent.rollback()
 
-    def add_child(self, transaction: "Transaction") -> None:
-        self.children.append(transaction)
-
-    def get_parent(self) -> Optional["Transaction"]:
-        parent = None
-        if self._token:
-            prev_var = getattr(self._token, "old_value")
-            if prev_var != Token.MISSING:
-                parent = prev_var
-        else:
-            # `_token` has been reset so we need to get the active transaction from the context var
-            parent = self.get_active()
-        return parent
-
     def commit(self) -> bool:
         if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
             if (
@@ -302,21 +340,19 @@ class Transaction(ContextModel):
 
         try:
             for child in self.children:
-                child.commit()
+                if inspect.iscoroutinefunction(child.commit):
+                    run_coro_as_sync(child.commit())
+                else:
+                    child.commit()
 
             for hook in self.on_commit_hooks:
                 self.run_hook(hook, "commit")
 
             if self.store and self.key and self.write_on_commit:
-                if isinstance(self.store, ResultStore):
-                    if isinstance(self._staged_value, ResultRecord):
-                        self.store.persist_result_record(
-                            result_record=self._staged_value
-                        )
-                    else:
-                        self.store.write(key=self.key, obj=self._staged_value)
+                if isinstance(self._staged_value, ResultRecord):
+                    self.store.persist_result_record(result_record=self._staged_value)
                 else:
-                    self.store.write(key=self.key, result=self._staged_value)
+                    self.store.write(key=self.key, obj=self._staged_value)
 
             self.state = TransactionState.COMMITTED
             if (
@@ -353,7 +389,10 @@ class Transaction(ContextModel):
             self.logger.info(f"Running {hook_type} hook {hook_name!r}")
 
         try:
-            hook(self)
+            if asyncio.iscoroutinefunction(hook):
+                run_coro_as_sync(hook(self))
+            else:
+                hook(self)
         except Exception as exc:
             if should_log:
                 self.logger.error(
@@ -366,24 +405,6 @@ class Transaction(ContextModel):
                     f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
                 )
 
-    def stage(
-        self,
-        value: Any,
-        on_rollback_hooks: Optional[list[Callable[..., Any]]] = None,
-        on_commit_hooks: Optional[list[Callable[..., Any]]] = None,
-    ) -> None:
-        """
-        Stage a value to be committed later.
-        """
-        on_commit_hooks = on_commit_hooks or []
-        on_rollback_hooks = on_rollback_hooks or []
-
-        if self.state != TransactionState.COMMITTED:
-            self._staged_value = value
-            self.on_rollback_hooks += on_rollback_hooks
-            self.on_commit_hooks += on_commit_hooks
-            self.state = TransactionState.STAGED
-
     def rollback(self) -> bool:
         if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
             return False
@@ -392,10 +413,13 @@ class Transaction(ContextModel):
             for hook in reversed(self.on_rollback_hooks):
                 self.run_hook(hook, "rollback")
 
-            self.state = TransactionState.ROLLED_BACK
+            self.state: TransactionState = TransactionState.ROLLED_BACK
 
             for child in reversed(self.children):
-                child.rollback()
+                if inspect.iscoroutinefunction(child.rollback):
+                    run_coro_as_sync(child.rollback())
+                else:
+                    child.rollback()
 
             return True
         except Exception:
@@ -414,24 +438,221 @@ class Transaction(ContextModel):
                 self.logger.debug(f"Releasing lock for transaction {self.key!r}")
                 self.store.release_lock(self.key)
 
-    @classmethod
-    def get_active(cls: Type[Self]) -> Optional[Self]:
-        return cls.__var__.get(None)
 
+class AsyncTransaction(BaseTransaction):
+    """
+    A model representing the state of an asynchronous transaction.
+    """
+
+    async def begin(self) -> None:
+        if (
+            self.store
+            and self.key
+            and self.isolation_level == IsolationLevel.SERIALIZABLE
+        ):
+            self.logger.debug(f"Acquiring lock for transaction {self.key!r}")
+            await self.store.aacquire_lock(self.key)
+        if (
+            not self.overwrite
+            and self.store
+            and self.key
+            and await self.store.aexists(key=self.key)
+        ):
+            self.state = TransactionState.COMMITTED
+
+    async def read(self) -> ResultRecord[Any] | None:
+        if self.store and self.key:
+            return await self.store.aread(key=self.key)
+        return None
 
-def get_transaction() -> Optional[Transaction]:
-    return Transaction.get_active()
+    async def reset(self) -> None:
+        parent = self.get_parent()
+
+        if parent:
+            # parent takes responsibility
+            parent.add_child(self)
+
+        if self._token:
+            self.__var__.reset(self._token)
+            self._token = None
+
+        # do this below reset so that get_transaction() returns the relevant txn
+        if parent and self.state == TransactionState.ROLLED_BACK:
+            await parent.rollback()
+
+    async def commit(self) -> bool:
+        if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+
+            return False
+
+        try:
+            for child in self.children:
+                if isinstance(child, AsyncTransaction):
+                    await child.commit()
+                else:
+                    child.commit()
+
+            for hook in self.on_commit_hooks:
+                await self.run_hook(hook, "commit")
+
+            if self.store and self.key and self.write_on_commit:
+                if isinstance(self._staged_value, ResultRecord):
+                    await self.store.apersist_result_record(
+                        result_record=self._staged_value
+                    )
+                else:
+                    await self.store.awrite(key=self.key, obj=self._staged_value)
+
+            self.state = TransactionState.COMMITTED
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+            return True
+        except SerializationError as exc:
+            if self.logger:
+                self.logger.warning(
+                    f"Encountered an error while serializing result for transaction {self.key!r}: {exc}"
+                    " Code execution will continue, but the transaction will not be committed.",
+                )
+            await self.rollback()
+            return False
+        except Exception:
+            if self.logger:
+                self.logger.exception(
+                    f"An error was encountered while committing transaction {self.key!r}",
+                    exc_info=True,
+                )
+            await self.rollback()
+            return False
+
+    async def run_hook(self, hook: Callable[..., Any], hook_type: str) -> None:
+        hook_name = get_hook_name(hook)
+        # Undocumented way to disable logging for a hook. Subject to change.
+        should_log = getattr(hook, "log_on_run", True)
+
+        if should_log:
+            self.logger.info(f"Running {hook_type} hook {hook_name!r}")
+
+        try:
+            if asyncio.iscoroutinefunction(hook):
+                await hook(self)
+            else:
+                await anyio.to_thread.run_sync(hook, self)
+        except Exception as exc:
+            if should_log:
+                self.logger.error(
+                    f"An error was encountered while running {hook_type} hook {hook_name!r}",
+                )
+            raise exc
+        else:
+            if should_log:
+                self.logger.info(
+                    f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
+                )
+
+    async def rollback(self) -> bool:
+        if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
+            return False
+
+        try:
+            for hook in reversed(self.on_rollback_hooks):
+                await self.run_hook(hook, "rollback")
+
+            self.state: TransactionState = TransactionState.ROLLED_BACK
+
+            for child in reversed(self.children):
+                if isinstance(child, AsyncTransaction):
+                    await child.rollback()
+                else:
+                    child.rollback()
+
+            return True
+        except Exception:
+            if self.logger:
+                self.logger.exception(
+                    f"An error was encountered while rolling back transaction {self.key!r}",
+                    exc_info=True,
+                )
+            return False
+        finally:
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+
+    async def __aenter__(self) -> Self:
+        self.prepare_transaction()
+        await self.begin()
+        self._token = self.__var__.set(self)
+        return self
+
+    async def __aexit__(self, *exc_info: Any) -> None:
+        exc_type, exc_val, _ = exc_info
+        if not self._token:
+            raise RuntimeError(
+                "Asymmetric use of context. Context exit called without an enter."
+            )
+        if exc_type:
+            await self.rollback()
+            await self.reset()
+            raise exc_val
+
+        if self.commit_mode == CommitMode.EAGER:
+            await self.commit()
+
+        # if parent, let them take responsibility
+        if self.get_parent():
+            await self.reset()
+            return
+
+        if self.commit_mode == CommitMode.OFF:
+            # if no one took responsibility to commit, rolling back
+            # note that rollback returns if already committed
+            await self.rollback()
+        elif self.commit_mode == CommitMode.LAZY:
+            # no one left to take responsibility for committing
+            await self.commit()
+
+        await self.reset()
+
+    def __enter__(self) -> NoReturn:
+        raise NotImplementedError(
+            "AsyncTransaction does not support the `with` statement. Use the `async with` statement instead."
+        )
+
+    def __exit__(self, *exc_info: Any) -> NoReturn:
+        raise NotImplementedError(
+            "AsyncTransaction does not support the `with` statement. Use the `async with` statement instead."
+        )
+
+
+def get_transaction() -> BaseTransaction | None:
+    return BaseTransaction.get_active()
 
 
 @contextmanager
 def transaction(
-    key: Optional[str] = None,
-    store: Optional[ResultStore] = None,
-    commit_mode: Optional[CommitMode] = None,
-    isolation_level: Optional[IsolationLevel] = None,
+    key: str | None = None,
+    store: ResultStore | None = None,
+    commit_mode: CommitMode | None = None,
+    isolation_level: IsolationLevel | None = None,
     overwrite: bool = False,
     write_on_commit: bool = True,
-    logger: Optional[Union[logging.Logger, LoggingAdapter]] = None,
+    logger: logging.Logger | LoggingAdapter | None = None,
 ) -> Generator[Transaction, None, None]:
     """
     A context manager for opening and managing a transaction.
@@ -473,3 +694,56 @@ def transaction(
         logger=_logger,
     ) as txn:
         yield txn
+
+
+@asynccontextmanager
+async def atransaction(
+    key: str | None = None,
+    store: ResultStore | None = None,
+    commit_mode: CommitMode | None = None,
+    isolation_level: IsolationLevel | None = None,
+    overwrite: bool = False,
+    write_on_commit: bool = True,
+    logger: logging.Logger | LoggingAdapter | None = None,
+) -> AsyncGenerator[AsyncTransaction, None]:
+    """
+    An asynchronous context manager for opening and managing an asynchronous transaction.
+
+    Args:
+        - key: An identifier to use for the transaction
+        - store: The store to use for persisting the transaction result. If not provided,
+            a default store will be used based on the current run context.
+        - commit_mode: The commit mode controlling when the transaction and
+            child transactions are committed
+        - overwrite: Whether to overwrite an existing transaction record in the store
+        - write_on_commit: Whether to write the result to the store on commit. If not provided,
+            the default will be determined by the current run context. If no run context is
+            available, the value of `PREFECT_RESULTS_PERSIST_BY_DEFAULT` will be used.
+
+    Yields:
+        - AsyncTransaction: An object representing the transaction state
+    """
+
+    # if there is no key, we won't persist a record
+    if key and not store:
+        store = get_result_store()
+
+    # Avoid inheriting a NullFileSystem for metadata_storage from a flow's result store
+    if store and isinstance(store.metadata_storage, NullFileSystem):
+        store = store.model_copy(update={"metadata_storage": None})
+
+    try:
+        _logger: Union[logging.Logger, LoggingAdapter] = logger or get_run_logger()
+    except MissingContextError:
+        _logger = get_logger("transactions")
+
+    async with AsyncTransaction(
+        key=key,
+        store=store,
+        commit_mode=commit_mode,
+        isolation_level=isolation_level,
+        overwrite=overwrite,
+        write_on_commit=write_on_commit,
+        logger=_logger,
+    ) as txn:
+        yield txn

{prefect_client-3.3.5.dev1.dist-info → prefect_client-3.3.5.dev3.dist-info}/METADATA

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

{prefect_client-3.3.5.dev1.dist-info → prefect_client-3.3.5.dev3.dist-info}/RECORD

@@ -1,7 +1,7 @@
 prefect/.prefectignore,sha256=awSprvKT0vI8a64mEOLrMxhxqcO-b0ERQeYpA2rNKVQ,390
 prefect/__init__.py,sha256=iCdcC5ZmeewikCdnPEP6YBAjPNV5dvfxpYCTpw30Hkw,3685
 prefect/__main__.py,sha256=WFjw3kaYJY6pOTA7WDOgqjsz8zUEUZHCcj3P5wyVa-g,66
-prefect/_build_info.py,sha256=WzWBU-7Qx6CRKSDzzF5Hy7t0Qb8X-5nU-NHHi1kXwak,185
+prefect/_build_info.py,sha256=8rZCXd-Nz2-o4Nm_4yyvy6osOkdGCLiuhickzD7KvFM,185
 prefect/_result_records.py,sha256=S6QmsODkehGVSzbMm6ig022PYbI6gNKz671p_8kBYx4,7789
 prefect/_waiters.py,sha256=Ia2ITaXdHzevtyWIgJoOg95lrEXQqNEOquHvw3T33UQ,9026
 prefect/agent.py,sha256=dPvG1jDGD5HSH7aM2utwtk6RaJ9qg13XjkA0lAIgQmY,287
@@ -14,21 +14,21 @@ prefect/exceptions.py,sha256=wZLQQMRB_DyiYkeEdIC5OKwbba5A94Dlnics-lrWI7A,11581
 prefect/filesystems.py,sha256=v5YqGB4uXf9Ew2VuB9VCSkawvYMMVvEtZf7w1VmAmr8,18036
 prefect/flow_engine.py,sha256=hZpTYEtwTPMtwVoTCrfD93igN7rlKeG_0kyCvdU4aYE,58876
 prefect/flow_runs.py,sha256=dbHcXsOq1UsNM7vyJV9gboCTylmdUwQ_-W4NQt4R4ds,17267
-prefect/flows.py,sha256=0Es8TYUUEItxAz6G50eUmlIAXDaUTh4U0dQvgUyW2rk,109529
+prefect/flows.py,sha256=kBLT6M903ZFD4TUvmvRN-zOdNcjSlIO9d2kHhPYq5Oo,109547
 prefect/futures.py,sha256=ZD5rdgUHA4sfxwHaPToumOUKlyn4d989JHR7eI97-Hs,23271
 prefect/main.py,sha256=8V-qLB4GjEVCkGRgGXeaIk-JIXY8Z9FozcNluj4Sm9E,2589
 prefect/plugins.py,sha256=FPRLR2mWVBMuOnlzeiTD9krlHONZH2rtYLD753JQDNQ,2516
 prefect/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-prefect/results.py,sha256=9mMOOZsj8ueqEch3oqxCaZPhF6D76Sn3P5TkqgVpbg8,36679
+prefect/results.py,sha256=3zuO8a4EGtti5AP4zLSnoES1YMKP1CugergxJdZa0-8,37140
 prefect/schedules.py,sha256=dhq4OhImRvcmtxF7UH1m8RbwYdHT5RQsp_FrxVXfODE,7289
 prefect/serializers.py,sha256=QI0oEal_BO4HQaWSjr6ReSwT55Hn4sbSOXxGgQI1-y0,9249
 prefect/states.py,sha256=rh7l1bnIYpTXdlXt5nnpz66y9KLjBWAJrN9Eo5RwgQs,26023
-prefect/task_engine.py,sha256=IIvDRl2bnnO3aKXTmtWsz0pnV8kL7xjOaUyJTlL6LaM,61491
+prefect/task_engine.py,sha256=wLOj7dPuIBBDZkvsqxWv7wLpUOhGnRb3N5s0ia2y5So,61576
 prefect/task_runners.py,sha256=Ce_ngocfq_X-NA5zhPj13IdVmzZ5h6gXlmfxYWs2AXA,15828
 prefect/task_runs.py,sha256=7LIzfo3fondCyEUpU05sYFN5IfpZigBDXrhG5yc-8t0,9039
 prefect/task_worker.py,sha256=gMj_rl4EjTrnJ5YSOXinC6y-7KSK7fRQt_UYbZbrrV8,17879
-prefect/tasks.py,sha256=X83KIVowu-c-EDFbmh4ZurEVtEKTyUSkQ8Elo0nWy4k,74946
-prefect/transactions.py,sha256=BYvxr4ZSFmYDCODPhH8DO1_51inH35oJ75ZZOd_GI_w,16341
+prefect/tasks.py,sha256=EpMw5O1B9pAFVraC0KzytMOKi8iy7ZYnKWRs7WtvogU,74742
+prefect/transactions.py,sha256=uIoPNudzJzH6NrMJhrgr5lyh6JxOJQqT1GvrXt69yNw,26068
 prefect/variables.py,sha256=dCK3vX7TbkqXZhnNT_v7rcGh3ISRqoR6pJVLpoll3Js,8342
 prefect/_experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 prefect/_experimental/bundles.py,sha256=EIZc6hsjQS8V2CH1RbsxOiFeR11l1IgvhiXTSO3uDyM,6386
@@ -146,7 +146,7 @@ prefect/docker/__init__.py,sha256=z6wdc6UFfiBG2jb9Jk64uCWVM04JKVWeVyDWwuuon8M,52
 prefect/docker/docker_image.py,sha256=bR_pEq5-FDxlwTj8CP_7nwZ_MiGK6KxIi8v7DRjy1Kg,3138
 prefect/events/__init__.py,sha256=GtKl2bE--pJduTxelH2xy7SadlLJmmis8WR1EYixhuA,2094
 prefect/events/actions.py,sha256=A7jS8bo4zWGnrt3QfSoQs0uYC1xfKXio3IfU0XtTb5s,9129
-prefect/events/clients.py,sha256=XA33NpeRdgVglt7J47uFdpASa1bCvcKWyxsxQt3vEPQ,27290
+prefect/events/clients.py,sha256=g_LNcPtYd1Erbz1q4UI33e5dmPjQv_OeejoYGkASs5w,27581
 prefect/events/filters.py,sha256=2hVfzc3Rdgy0mBHDutWxT__LJY0zpVM8greWX3y6kjM,8233
 prefect/events/related.py,sha256=CTeexYUmmA93V4gsR33GIFmw-SS-X_ouOpRg-oeq-BU,6672
 prefect/events/utilities.py,sha256=ww34bTMENCNwcp6RhhgzG0KgXOvKGe0MKmBdSJ8NpZY,3043
@@ -171,7 +171,7 @@ prefect/input/actions.py,sha256=BDx26b6ZYCTr0kbWBp73Or7UXnLIv1lnm0jow6Simxw,3871
 prefect/input/run_input.py,sha256=GoM4LR3oqAFLf2sPCR1yITY9tNSZT8kAd4gaC-v-a-c,22703
 prefect/locking/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 prefect/locking/filesystem.py,sha256=RVE4_5lgKi1iea0NZVQlyct5GU4fVAtCPPEdRMDaQHw,8128
-prefect/locking/memory.py,sha256=Q8NqSeksdQb-AZfql_SXeTd4SRFZ3rWMBwA5shTVEZM,7860
+prefect/locking/memory.py,sha256=EFQnhAO94jEy4TyS880DbsJ42CHT5WNuNc6Wj8dYrKc,7842
 prefect/locking/protocol.py,sha256=RsfvlaHTTEJ0YvYWSqFGoZuT2w4FPPxyQlHqjoyNGuE,4240
 prefect/logging/__init__.py,sha256=DpRZzZeWeiDHFlMDEQdknRzbxpL0ObFh5IqqS9iaZwQ,170
 prefect/logging/configuration.py,sha256=ZBAOgwE34VSZFSiP4gBEd0S9m645_mEs3dFhiwPj58o,3303
@@ -269,7 +269,7 @@ prefect/settings/models/server/tasks.py,sha256=_CaOUfh3WDXvUhmHXmR-MkTRaQqocZck4
 prefect/settings/models/server/ui.py,sha256=hShsi4rPBtdJA2WnT1Er0tWqu-e5wUum8NkNgucShkk,1867
 prefect/telemetry/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 prefect/telemetry/bootstrap.py,sha256=XaYlK4OTaloQX63AG7IfvGD7yH5LOsXfol4pPfaXSBI,1537
-prefect/telemetry/instrumentation.py,sha256=1wVdROX8Nz3P4Ja3we92UjRH1xF45Pm2sAA3rZ5rqtw,4727
+prefect/telemetry/instrumentation.py,sha256=5b8b5183IrmBgH722c2TIQHpjBrE0vlRqcKpfwRG0a0,4732
 prefect/telemetry/logging.py,sha256=ktIVTXbdZ46v6fUhoHNidFrpvpNJR-Pj-hQ4V9b40W4,789
 prefect/telemetry/processors.py,sha256=jw6j6LviOVxw3IBJe7cSjsxFk0zzY43jUmy6C9pcfCE,2272
 prefect/telemetry/run_telemetry.py,sha256=_FbjiPqPemu4xvZuI2YBPwXeRJ2BcKRJ6qgO4UMzKKE,8571
@@ -316,7 +316,7 @@ prefect/workers/cloud.py,sha256=dPvG1jDGD5HSH7aM2utwtk6RaJ9qg13XjkA0lAIgQmY,287
 prefect/workers/process.py,sha256=uxOwcqA2Ps-V-W6WeSdKCQMINrCxBEVx1K1Un8pb7vs,8973
 prefect/workers/server.py,sha256=2pmVeJZiVbEK02SO6BEZaBIvHMsn6G8LzjW8BXyiTtk,1952
 prefect/workers/utilities.py,sha256=VfPfAlGtTuDj0-Kb8WlMgAuOfgXCdrGAnKMapPSBrwc,2483
-prefect_client-3.3.5.dev1.dist-info/METADATA,sha256=I-vHcTjY3Yn3NggXlY03fy-YY9uhtnhDLeXIZx5v9T8,7456
-prefect_client-3.3.5.dev1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-prefect_client-3.3.5.dev1.dist-info/licenses/LICENSE,sha256=MCxsn8osAkzfxKC4CC_dLcUkU8DZLkyihZ8mGs3Ah3Q,11357
-prefect_client-3.3.5.dev1.dist-info/RECORD,,
+prefect_client-3.3.5.dev3.dist-info/METADATA,sha256=LcejbNc0zC2HoeBq6bkhgj5Pyne1GPZwmi9SLwEwCBc,7456
+prefect_client-3.3.5.dev3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+prefect_client-3.3.5.dev3.dist-info/licenses/LICENSE,sha256=MCxsn8osAkzfxKC4CC_dLcUkU8DZLkyihZ8mGs3Ah3Q,11357
+prefect_client-3.3.5.dev3.dist-info/RECORD,,