fix: make database saver classes inheritance-friendly (#2615)

Replace hardcoded database saver class names with `cls` in
`from_conn_string` factory methods to improve subclassing support

## Changes
* Replaced direct class instantiations with `cls(conn)` in
`from_conn_string` classmethods across all database implementations
* Updated both synchronous and asynchronous variants for DuckDB,
PostgreSQL, and SQLite savers

## Why
This refactor makes the database saver classes more extensible by
following Python's convention of using `cls` in class methods. This
enables proper inheritance patterns where subclasses can reuse the
factory methods without needing to override them. Previously, the
hardcoded class names would always instantiate the parent class, even
when called from a subclass.

## Testing
The change is backward compatible and doesn't alter existing
functionality. All existing tests should continue to pass as this is
purely a structural refactoring that preserves the current behavior
while improving extensibility.

## Notes
This PR addresses follow up on comments from #2518 - AsyncPostgresSaver
didn't need to be fixed but many of the other DB saver classes did.
This commit is contained in:
Phoenix Logan
2024-12-03 20:26:06 -08:00
committed by GitHub
parent 5fa196ab38
commit aca67107c1
6 changed files with 6 additions and 6 deletions
@@ -42,7 +42,7 @@ class DuckDBSaver(BaseDuckDBSaver):
DuckDBSaver: A new DuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield DuckDBSaver(conn)
yield cls(conn)
def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
@@ -45,7 +45,7 @@ class AsyncDuckDBSaver(BaseDuckDBSaver):
AsyncDuckDBSaver: A new AsyncDuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield AsyncDuckDBSaver(conn)
yield cls(conn)
async def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
@@ -156,7 +156,7 @@ class AsyncDuckDBStore(AsyncBatchedBaseStore, BaseDuckDBStore):
AsyncDuckDBStore: A new AsyncDuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield AsyncDuckDBStore(conn)
yield cls(conn)
async def setup(self) -> None:
"""Set up the store database asynchronously.
@@ -54,7 +54,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
pipeline: bool = False,
serde: Optional[SerializerProtocol] = None,
) -> AsyncIterator["AsyncPostgresSaver"]:
"""Create a new PostgresSaver instance from a connection string.
"""Create a new AsyncPostgresSaver instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
@@ -110,7 +110,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
check_same_thread=False,
)
) as conn:
yield SqliteSaver(conn)
yield cls(conn)
def setup(self) -> None:
"""Set up the checkpoint database.
@@ -137,7 +137,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
AsyncSqliteSaver: A new AsyncSqliteSaver instance.
"""
async with aiosqlite.connect(conn_string) as conn:
yield AsyncSqliteSaver(conn)
yield cls(conn)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.