Async db support #497

jernejfrank · 2025-01-13T08:11:26Z

Adding async support for Postgres and Redis state persisters. This expands on #488 and addresses #484 when using Postgres or Redis.

Changes

Potgres state persisted integrates commands using asyncpg.
Redis supports async via redis.asyncio

How I tested this

unit test

Notes

For Postgres: asyncpg uses a coroutine to open/close the connection. In this case pickling / unpickling the state needs a workaround (not yet implemented) since __getstate__ and __setstate__ do not have async support.
Redis pickling/unpickling is straightforward (and implemented) since the connection to the db is sync.

Checklist

PR has an informative and human-readable title (this will be pulled into the release notes)
Changes are limited to a single goal (no scope creep)
Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
Any change in functionality is tested
New functions are documented (with a description, list of inputs, and expected output)
Placeholder code is flagged / future TODOs are captured in comments
Project documentation has been updated if adding/changing functionality.

Important

Add asynchronous support for Redis and PostgreSQL state persisters using redis.asyncio and asyncpg.

Async Support:
- Added AsyncRedisBasePersister using redis.asyncio for async Redis operations in b_redis.py.
- Added AsyncPostgreSQLPersister using asyncpg for async PostgreSQL operations in postgresql.py.
Testing:
- Added async tests in test_b_redis.py for AsyncRedisBasePersister.
- Added async tests in test_postgresql.py for AsyncPostgreSQLPersister.
Documentation:
- Updated persister.rst to include AsyncRedisBasePersister and AsyncPostgreSQLPersister.
Dependencies:
- Added asyncpg to pyproject.toml under postgresql dependencies.

^{This description was created by}^{for 304eb21. It will automatically update as commits are pushed.}

Async implementation to persist the state to a postgresql db.

Implements the async function from redis-py.asyncio for state persistance.

ellipsis-dev

👍 Looks good to me! Reviewed everything up to 304eb21 in 37 seconds

More details

Looked at 662 lines of code in 6 files
Skipped 0 files when reviewing.
Skipped posting 2 drafted comments based on config settings.

1. burr/integrations/persisters/b_redis.py:315

Draft comment:
The create_key method does not perform any asynchronous operations and can be converted to a regular method.
Reason this comment was not posted:
Confidence changes required: 50%
The create_key method in AsyncRedisBasePersister is unnecessarily defined as an async function. It does not perform any asynchronous operations and can be converted to a regular method.

2. burr/integrations/persisters/postgresql.py:362

Draft comment:
The create_key method does not perform any asynchronous operations and can be converted to a regular method.
Reason this comment was not posted:
Confidence changes required: 50%
The create_key method in AsyncPostgreSQLPersister is unnecessarily defined as an async function. It does not perform any asynchronous operations and can be converted to a regular method.

Workflow ID: wflow_m5wjdK3DlH4vtpsj

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

elijahbenizzy

Some comments, will take a look again

elijahbenizzy · 2025-01-14T19:13:21Z

burr/integrations/persisters/b_redis.py

@@ -2,6 +2,7 @@

 try:
    import redis  # can't name module redis because this import wouldn't work.
+    import redis.asyncio as aredis


Just checking -- redis.asyncio is included in redis, right? It's not a separate library with separate dependencies?

Choice:

Keep async + sync redis in the same place (makes sense cause of the library)

Move async redis to b_aredis.py or something like that (makes sense cause it's a separate functionality and you could imagine them separating it out or using an alternative client at some point)

elijahbenizzy · 2025-01-14T19:13:32Z

burr/integrations/persisters/b_redis.py

+class AsyncRedisBasePersister(persistence.BaseStatePersister):
+    """Main class for Async Redis persister.
+
+    Use this class if you want to directly control injecting the async Redis client.


What does this mean?

elijahbenizzy · 2025-01-14T20:39:18Z

burr/integrations/persisters/b_redis.py

+    .. warning::
+        The synchronous persister closes the connection on deletion of the class using the ``__del__`` method.
+        In an async context that is not reliable (the event loop may already be closed by the time ``__del__``
+        gets invoked). Therefore, you are responsible for closing the connection yourself (i.e. manual cleanup).


Worth thinking -- can this be done by a context manager? E.G. if you use from_values() how do you close it? Might be worth either exposing the client or something.

elijahbenizzy · 2025-01-14T20:41:15Z

burr/integrations/persisters/b_redis.py

+            "status": data[b"status"].decode(),
+        }
+
+    async def create_key(self, app_id, partition_key, sequence_id):


Should this be async?

Wonder if this should subclass the synchronous redis persister to share code -- there's a lot of duplication...

Alternatively helper functions to help with common operations

elijahbenizzy · 2025-01-14T20:42:35Z

burr/integrations/persisters/postgresql.py

@@ -1,6 +1,7 @@
 from burr.integrations import base

 try:
+    import asyncpg


Ditto with the above -- I think this should probably be a separate one so we don't force people who just need sync to also have async libraries

elijahbenizzy · 2025-01-14T20:45:20Z

burr/integrations/persisters/postgresql.py

+        if self._initialized:
+            return True
+
+        async with self.connection.transaction():


Does this need to be a transaction? It's a single query, and it's purely a read (so it doesn't/can't be rolled back)

Hmm, also feels like we should be able to set self._initialized to be true here if _initialized is true...

elijahbenizzy · 2025-01-14T20:46:35Z

burr/integrations/persisters/postgresql.py

+            partition_key = self.PARTITION_KEY_DEFAULT
+        logger.debug("Loading %s, %s, %s", partition_key, app_id, sequence_id)
+
+        async with self.connection.transaction():


ditto with transaction -- not needed I think

elijahbenizzy · 2025-01-14T20:56:50Z

burr/integrations/persisters/postgresql.py

+            status,
+        )
+
+        async with self.connection.transaction():


Also should not need to be a transaction

jernejfrank added 2 commits January 13, 2025 15:54

Add async PostgreSQL persister support

d3ba49d

Async implementation to persist the state to a postgresql db.

Add async support for Redis state persister

304eb21

Implements the async function from redis-py.asyncio for state persistance.

ellipsis-dev bot reviewed Jan 13, 2025

View reviewed changes

jernejfrank added 2 commits January 13, 2025 16:36

Fix asyncpg dependency import fro tests

bdd613e

Fix redis async tests have their own table db=1

9bce9c2

elijahbenizzy reviewed Jan 15, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Async db support #497

Async db support #497

jernejfrank commented Jan 13, 2025 •

edited by ellipsis-dev bot

Loading

ellipsis-dev bot left a comment

elijahbenizzy left a comment

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

Async db support #497

Are you sure you want to change the base?

Async db support #497

Conversation

jernejfrank commented Jan 13, 2025 • edited by ellipsis-dev bot Loading

Changes

How I tested this

Notes

Checklist

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

elijahbenizzy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jernejfrank commented Jan 13, 2025 •

edited by ellipsis-dev bot

Loading