From aca67107c14516aa5b0cfd750adcb4c15aae21f0 Mon Sep 17 00:00:00 2001
From: Phoenix Logan <plogan@chanzuckerberg.com>
Date: Tue, 3 Dec 2024 20:26:06 -0800
Subject: [PATCH] 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.
---
 libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/__init__.py | 2 +-
 libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/aio.py      | 2 +-
 libs/checkpoint-duckdb/langgraph/store/duckdb/aio.py           | 2 +-
 libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py  | 2 +-
 libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py | 2 +-
 libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py      | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/__init__.py b/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/__init__.py
index 1002eebe8..7a873ab4b 100644
--- a/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/__init__.py
+++ b/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/__init__.py
@@ -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.
diff --git a/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/aio.py b/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/aio.py
index aa52feb18..54c1924cc 100644
--- a/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/aio.py
+++ b/libs/checkpoint-duckdb/langgraph/checkpoint/duckdb/aio.py
@@ -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.
diff --git a/libs/checkpoint-duckdb/langgraph/store/duckdb/aio.py b/libs/checkpoint-duckdb/langgraph/store/duckdb/aio.py
index d6fd7dd89..f050f449b 100644
--- a/libs/checkpoint-duckdb/langgraph/store/duckdb/aio.py
+++ b/libs/checkpoint-duckdb/langgraph/store/duckdb/aio.py
@@ -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.
diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
index 4c0f5295c..3a1e13db1 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
@@ -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.
diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
index b552a75f4..ea749473c 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
@@ -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.
diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
index 21cde06e0..72fca5bea 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
@@ -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.