Skip to content

fix: support async with on async persister factory methods#681

Open
andreahlert wants to merge 4 commits intoapache:mainfrom
andreahlert:fix/issue-546-async-context-manager
Open

fix: support async with on async persister factory methods#681
andreahlert wants to merge 4 commits intoapache:mainfrom
andreahlert:fix/issue-546-async-context-manager

Conversation

@andreahlert
Copy link
Contributor

Summary

Fixes #546

AsyncSQLitePersister.from_values() and AsyncPostgreSQLPersister.from_values() are async classmethods that return coroutines. Using them directly with async with fails with:

TypeError: 'coroutine' object does not support the asynchronous context manager protocol

This PR introduces _AsyncPersisterContextManager, a thin wrapper that implements both __await__ and __aenter__/__aexit__, so factory methods now support both usage patterns:

# await (backwards compatible)
persister = await AsyncSQLitePersister.from_values(db_path="test.db")

# async with (what the issue requested)
async with AsyncSQLitePersister.from_values(db_path="test.db") as persister:
    await persister.initialize()
    ...

The same fix is applied to AsyncPostgreSQLPersister in b_asyncpg.py.

Test plan

  • Reproduced the original TypeError in a Docker container (Python 3.11)
  • Verified async with ... from_values() works after the fix
  • Verified await ... from_values() still works (backwards compatibility)
  • @skrawcz review
  • @kajocina validate against original use case

`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
Copy link
Contributor

@skrawcz skrawcz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. Move _AsyncPersisterContextManager to burr/common/. Having b_asyncpg.py import from b_aiosqlite.py is an unwanted cross-dependency between unrelated integrations. Something like burr/common/types.py or a new burr/common/async_utils.py would be a better home.

  2. Add tests for the async with pattern. All existing tests use await, 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
@andreahlert
Copy link
Contributor Author

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:

  1. Move _AsyncPersisterContextManager to burr/common/. Having b_asyncpg.py import from b_aiosqlite.py is an unwanted cross-dependency between unrelated integrations. Something like burr/common/types.py or a new burr/common/async_utils.py would be a better home.
  2. Add tests for the async with pattern. All existing tests use await, 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.

Addressed all three points. While testing the fix in an isolated Docker container (Python 3.11), I caught two additional bugs:

  • aexit would crash with AttributeError if aenter failed (e.g. bad db path), because self._persister was still None. Now guarded.
  • The wrapper could be silently consumed twice (coroutine exhaustion). Added a _consumed flag that raises RuntimeError with a clear message.

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?

@andreahlert andreahlert requested a review from skrawcz March 24, 2026 12:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Async state persister doesn't work as a context manager

2 participants