Merge pull request #27 from tjni/wfh/provisioning_cursor_reuse

Fixup initial provisioning of aio postgres db

langgraph-tests/tests/conftest.py

@@ -59,7 +59,7 @@ def checkpointer_pymysql():
         # drop unique db
         with pymysql.connect(**PyMySQLSaver.parse_conn_string(DEFAULT_MYSQL_URI), autocommit=True) as conn:
             with conn.cursor() as cursor:
-                cursor.execute(f"drop databASE {database}")
+                cursor.execute(f"DROP DATABASE {database}")
 
 
 @asynccontextmanager

langgraph/checkpoint/mysql/__init__.py

@@ -11,7 +11,6 @@
     Protocol,
     TypeVar,
     Union,
-    cast,
 )
 
 from langchain_core.runnables import RunnableConfig
@@ -34,7 +33,7 @@
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
 
-class DictCursor(Protocol):
+class DictCursor(ContextManager, Protocol):
     """
     Protocol that a cursor should implement.
 
@@ -54,7 +53,7 @@ def fetchone(self) -> Optional[dict[str, Any]]: ...
     def fetchall(self) -> Sequence[dict[str, Any]]: ...
 
 
-R = TypeVar("R", bound=ContextManager, covariant=True)  # cursor type
+R = TypeVar("R", bound=DictCursor)  # cursor type
 
 
 Conn = _internal.Conn  # For backward compatibility
@@ -73,10 +72,6 @@ def __init__(
         self.conn = conn
         self.lock = threading.Lock()
 
-    @staticmethod
-    def _is_no_such_table_error(e: Exception) -> bool:
-        raise NotImplementedError
-
     @staticmethod
     def _get_cursor_from_connection(conn: _internal.C) -> R:
         raise NotImplementedError
@@ -110,21 +105,14 @@ def setup(self) -> None:
         already exist and runs database migrations. It MUST be called directly by the user
         the first time checkpointer is used.
         """
-        with self._cursor() as cur_:
-            cur = cast(DictCursor, cur_)
-            try:
-                cur.execute(
-                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                )
-                row = cur.fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except Exception as e:
-                if not self._is_no_such_table_error(e):
-                    raise
+        with self._cursor() as cur:
+            cur.execute(self.MIGRATIONS[0])
+            cur.execute("SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1")
+            row = cur.fetchone()
+            if row is None:
                 version = -1
+            else:
+                version = row["v"]
             for v, migration in zip(
                 range(version + 1, len(self.MIGRATIONS)),
                 self.MIGRATIONS[version + 1 :],
@@ -177,8 +165,7 @@ def list(
         if limit:
             query += f" LIMIT {limit}"
         # if we change this to use .stream() we need to make sure to close the cursor
-        with self._cursor() as cur_:
-            cur = cast(DictCursor, cur_)
+        with self._cursor() as cur:
             cur.execute(query, args)
             values = cur.fetchall()
             for value in values:
@@ -257,8 +244,7 @@ def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
             args = (thread_id, checkpoint_ns)
             where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
 
-        with self._cursor() as cur_:
-            cur = cast(DictCursor, cur_)
+        with self._cursor() as cur:
             cur.execute(
                 self.SELECT_SQL + where,
                 args,
@@ -342,8 +328,7 @@ def put(
             }
         }
 
-        with self._cursor(pipeline=True) as cur_:
-            cur = cast(DictCursor, cur_)
+        with self._cursor(pipeline=True) as cur:
             cur.executemany(
                 self.UPSERT_CHECKPOINT_BLOBS_SQL,
                 self._dump_blobs(
@@ -386,8 +371,7 @@ def put_writes(
             if all(w[0] in WRITES_IDX_MAP for w in writes)
             else self.INSERT_CHECKPOINT_WRITES_SQL
         )
-        with self._cursor(pipeline=True) as cur_:
-            cur = cast(DictCursor, cur_)
+        with self._cursor(pipeline=True) as cur:
             cur.executemany(
                 query,
                 self._dump_writes(

langgraph/checkpoint/mysql/aio.py

@@ -8,7 +8,6 @@
 import aiomysql  # type: ignore
 import pymysql
 import pymysql.connections
-import pymysql.constants.ER
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
@@ -100,19 +99,15 @@ async def setup(self) -> None:
         the first time checkpointer is used.
         """
         async with self._cursor() as cur:
-            try:
-                await cur.execute(
-                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                )
-                row = await cur.fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except pymysql.ProgrammingError as e:
-                if e.args[0] != pymysql.constants.ER.NO_SUCH_TABLE:
-                    raise
+            await cur.execute(self.MIGRATIONS[0])
+            await cur.execute(
+                "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+            )
+            row = await cur.fetchone()
+            if row is None:
                 version = -1
+            else:
+                version = row["v"]
             for v, migration in zip(
                 range(version + 1, len(self.MIGRATIONS)),
                 self.MIGRATIONS[version + 1 :],

langgraph/checkpoint/mysql/pymysql.py

@@ -56,14 +56,6 @@ def from_conn_string(
         ) as conn:
             yield cls(conn)
 
-    @override
-    @staticmethod
-    def _is_no_such_table_error(e: Exception) -> bool:
-        return (
-            isinstance(e, pymysql.ProgrammingError)
-            and e.args[0] == pymysql.constants.ER.NO_SUCH_TABLE
-        )
-
     @override
     @staticmethod
     def _get_cursor_from_connection(conn: pymysql.Connection) -> DictCursor: