fix: support async with on async persister factory methods#681
fix: support async with on async persister factory methods#681andreahlert wants to merge 1 commit intoapache:mainfrom
async with on async persister factory methods#681Conversation
`AsyncSQLitePersister.from_values()` and `AsyncPostgreSQLPersister.from_values()` were async classmethods returning coroutines, which cannot be used directly with `async with`. This wraps them in `_AsyncPersisterContextManager` that supports both `await` (backwards compatible) and `async with` protocols. Closes apache#546
| Self = None | ||
|
|
||
|
|
||
| class _AsyncPersisterContextManager: |
There was a problem hiding this comment.
Good solution to #546! The _AsyncPersisterContextManager approach is clean — supporting both await and async with from the same factory method is a nice UX improvement.
Two items to address:
-
Move
_AsyncPersisterContextManagertoburr/common/. Havingb_asyncpg.pyimport fromb_aiosqlite.pyis an unwanted cross-dependency between unrelated integrations. Something likeburr/common/types.pyor a newburr/common/async_utils.pywould be a better home. -
Add tests for the
async withpattern. All existing tests useawait, so the new context manager functionality is untested. At minimum, a test like:async def test_async_sqlite_from_values_as_context_manager(tmp_path): db_path = str(tmp_path / "test.db") async with AsyncSQLitePersister.from_values(db_path=db_path) as persister: await persister.initialize() await persister.save("pk", "app1", 1, "pos", State({"k": "v"}), "completed") loaded = await persister.load("pk", "app1") assert loaded is not None
Also checked: AsyncRedisBasePersister.from_values is already a regular (non-async) method, so it doesn't have this issue. The fix correctly covers the only two affected persisters.
| Self = None | ||
|
|
||
|
|
||
| class _AsyncPersisterContextManager: |
There was a problem hiding this comment.
+1 to my earlier suggestion to move this to burr/common/. Having b_asyncpg import from b_aiosqlite creates an unwanted coupling between unrelated integrations. burr/common/types.py or a new burr/common/async_utils.py would be a better location.
| The wrapper awaits the coroutine on ``__aenter__`` and delegates | ||
| ``__aexit__`` to the persister's own ``__aexit__``. | ||
| """ | ||
|
|
There was a problem hiding this comment.
Nit: consider adding a type annotation to coro:
def __init__(self, coro: Any):Or for stricter typing: Coroutine[Any, Any, T] with a TypeVar.
Summary
Fixes #546
AsyncSQLitePersister.from_values()andAsyncPostgreSQLPersister.from_values()are async classmethods that return coroutines. Using them directly withasync withfails with:This PR introduces
_AsyncPersisterContextManager, a thin wrapper that implements both__await__and__aenter__/__aexit__, so factory methods now support both usage patterns:The same fix is applied to
AsyncPostgreSQLPersisterinb_asyncpg.py.Test plan
TypeErrorin a Docker container (Python 3.11)async with ... from_values()works after the fixawait ... from_values()still works (backwards compatibility)