Skip to content

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

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

fix: support async with on async persister factory methods#681
andreahlert wants to merge 1 commit 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
Self = None


class _AsyncPersisterContextManager:
Copy link
Contributor

Choose a reason for hiding this comment

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

move to common.py?

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.

Self = None


class _AsyncPersisterContextManager:
Copy link
Contributor

@skrawcz skrawcz Mar 21, 2026

Choose a reason for hiding this comment

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

+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__``.
"""

Copy link
Contributor

Choose a reason for hiding this comment

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

Nit: consider adding a type annotation to coro:

def __init__(self, coro: Any):

Or for stricter typing: Coroutine[Any, Any, T] with a TypeVar.

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