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.}

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

burr/integrations/persisters/b_redis.py

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?

burr/integrations/persisters/b_redis.py

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

I went through both classes and there is not much shared stuff.. The annoying thing is that most code is of the pattern: talk to db -- extract some useful data (e.g. sequence_id) -- talk to db again. So all the logic that appears duplicate is sandwich between db calls that are sync/async depending on the class.

Got it, yeah, I think that's OK given sync versus async is generally tricky

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 · 2025-01-22T14:12:44Z

I have increased the scope to include all existing persister implementations to make some consistency additions. There are 3 points to highlight:

Consistency for methods across persisters:

.cleanup() to be able to close the connection
__enter__ / __exit__ or __aenter__ / __aexit__ to be used as context managers
cls.from_config()
set_serde_kwargs

Naming convention for persister modules:

I opted to name the modules after the underlying dependency library. I am not a fan of that naming, but this gives us more flexibility down the line in case we implement the same persister with multiple different libraries since we can then keep the class name -- For example:

psycopg2.PostgreSQLPersister
PyGreSQL.PostgreSQLPersister

New table in persister docs to highlight which implementations we have and which dependencies there are.

elijahbenizzy

I think this looks good -- it's consistent, which is nice. Will think through the naming but naming it based on the libraries/plugins makes sense to me! @skrawcz thoughts?

@jernejfrank --maybe add a note in the references on naming scheme/rules. Right next to the table.

elijahbenizzy · 2025-01-15T18:26:26Z

burr/integrations/persisters/b_redis.py

+    def __getstate__(self) -> dict:
+        state = self.__dict__.copy()
+        if not hasattr(self.connection, "connection_pool"):
+            logger.warning("Redis connection is not serializable.")


So I don't think we want a warning here -- this specifically is what we need to recreate in setstate. E.G. the reason we have this is so we can serialize a redis connection and recreate it when needed.

elijahbenizzy · 2025-01-24T17:12:02Z

burr/core/persistence.py

+    @classmethod
+    def from_config(cls, config: dict) -> "SQLitePersister":
+        """Creates a new instance of the SQLitePersister from a configuration dictionary."""
+        return cls.from_values(


What's the use of this? Also should be able to just use ** -- the defaults work.

Some DB persister had from_config and from_values, others had from_values, and some (like SQLite) did not have either. I implemented both from_config and from_values so that all persisters have the same ability to initialize. This maybe useful to put on the abstract class level to enforce it when implementing new persisters?

elijahbenizzy · 2025-01-24T17:12:18Z

burr/core/persistence.py

-            db_path, **connect_kwargs if connect_kwargs is not None else {}
-        )
+
+        # Here for backwards compatibility


explain a little more

elijahbenizzy · 2025-01-24T17:13:14Z

burr/core/persistence.py

    def __del__(self):
+        # This should be deprecated


explain why

burr/integrations/persisters/b_aiosqlite.py

elijahbenizzy · 2025-01-24T17:14:32Z

burr/integrations/persisters/b_mongodb.py

@@ -45,114 +25,22 @@ def from_values(
        mongo_client_kwargs: dict = None,
    ) -> "MongoDBBasePersister":
        """Initializes the MongoDBBasePersister class."""
+        logger.warning(


We should be able to put this logger at the import level I think. Just at the top of the file?

elijahbenizzy

@skrawcz any thoughts on naming for integraitons?

elijahbenizzy · 2025-01-31T05:33:34Z

burr/core/persistence.py

@@ -304,6 +338,7 @@ def __init__(
        table_name: str = "burr_state",
        serde_kwargs: dict = None,
        connect_kwargs: dict = None,
+        connection=None,


Async implementation to persist the state to a postgresql db. In addition adds methods to postgresql sync persister to make it consistent. Migrates postgresql to psycopg2 module for consistent renaming.

Implements the async function from redis-py.asyncio for state persistance. Add methods to sync redis persister for consistency.

The methods added make it consistent with the implementations of other db persisters.

Migrates the MongoDBPersister to b_pymongo.py for consistent naming convention.

Some libraries are only supported for python <= 3.12 this sets the python version to 3.12 for the docs and validate_examples hooks so that Ray gets installed.

ellipsis-dev bot reviewed Jan 13, 2025

View reviewed changes

elijahbenizzy reviewed Jan 15, 2025

View reviewed changes

jernejfrank force-pushed the async_db_support branch from 9bce9c2 to 63c306c Compare January 22, 2025 13:56

elijahbenizzy reviewed Jan 24, 2025

View reviewed changes

jernejfrank force-pushed the async_db_support branch from 63c306c to 8c834af Compare January 25, 2025 13:56

elijahbenizzy approved these changes Jan 31, 2025

View reviewed changes

jernejfrank added 5 commits February 1, 2025 09:02

Add async PostgreSQL persister support

eae9427

Async implementation to persist the state to a postgresql db. In addition adds methods to postgresql sync persister to make it consistent. Migrates postgresql to psycopg2 module for consistent renaming.

Add async support for Redis state persister

3c0bbf0

Implements the async function from redis-py.asyncio for state persistance. Add methods to sync redis persister for consistency.

Add general methods to sync and async SQLite persisters.

6c48452

The methods added make it consistent with the implementations of other db persisters.

Add methods to MongoDBPersister and migrate

a89c1c0

Migrates the MongoDBPersister to b_pymongo.py for consistent naming convention.

Update persister docs to reflect db dependencies

150b844

jernejfrank force-pushed the async_db_support branch 2 times, most recently from b2660e1 to 024625f Compare February 1, 2025 01:13

Update package imports

bac8e28

jernejfrank force-pushed the async_db_support branch 2 times, most recently from 61821b7 to c666cba Compare February 1, 2025 13:44

Pin python<3.12 in hooks to install dependencies

32d3eb4

Some libraries are only supported for python <= 3.12 this sets the python version to 3.12 for the docs and validate_examples hooks so that Ray gets installed.

jernejfrank force-pushed the async_db_support branch from c666cba to 32d3eb4 Compare February 2, 2025 00:13

elijahbenizzy merged commit dca4ab3 into DAGWorks-Inc:main Feb 3, 2025
10 of 11 checks passed

jernejfrank deleted the async_db_support branch February 4, 2025 02:33

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

jernejfrank Jan 22, 2025

elijahbenizzy Jan 31, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

elijahbenizzy Jan 14, 2025

jernejfrank commented Jan 22, 2025

elijahbenizzy left a comment

elijahbenizzy Jan 15, 2025

elijahbenizzy Jan 24, 2025

jernejfrank Jan 25, 2025

elijahbenizzy Jan 24, 2025

elijahbenizzy Jan 24, 2025

elijahbenizzy Jan 24, 2025

elijahbenizzy left a comment

elijahbenizzy Jan 31, 2025

Async db support #497

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 22, 2025

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

elijahbenizzy left a 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