fix: support async with on async persister factory methods#681
fix: support async with on async persister factory methods#681andreahlert wants to merge 4 commits 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
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.
…ls.py and add tests Address review feedback: - Move _AsyncPersisterContextManager from b_aiosqlite.py to burr/common/async_utils.py to avoid cross-dependency between unrelated integrations - Add type annotation to coro parameter - Add tests for async with pattern on from_values and from_config
…manager wrapper - __aexit__ now returns False when __aenter__ failed (persister is None), preventing AttributeError that would mask the original exception - Add _consumed flag to prevent silent coroutine reuse, raising RuntimeError with clear message on second await/async with - Add tests for both edge cases
- Add missing `await` to from_values() calls in parallelism.rst docs - Remove AsyncSQLiteContextManager helper class from both test files, now that from_values() natively supports async with - Replace deprecated .close() calls with .cleanup() in test fixtures
Addressed all three points. While testing the fix in an isolated Docker container (Python 3.11), I caught two additional bugs:
Also fixed the parallelism.rst docs that were missing await on from_values() calls, and removed the now-redundant AsyncSQLiteContextManager helper from both test files. LGTM. Could you please take an extra look? |
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)