ci: Enable mypy checks for checkpoint-postgres lib

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -5,7 +5,7 @@ from typing import Any, Iterator, List, Optional, Union
 from langchain_core.runnables import RunnableConfig
 from psycopg import Connection, Cursor, Pipeline
 from psycopg.errors import UndefinedTable
-from psycopg.rows import dict_row
+from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
 
@@ -22,9 +22,11 @@ from langgraph.checkpoint.postgres.base import (
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
+Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
+
 
 @contextmanager
-def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]:
+def _get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
     if isinstance(conn, Connection):
         yield conn
     elif isinstance(conn, ConnectionPool):
@@ -39,7 +41,7 @@ class PostgresSaver(BasePostgresSaver):
 
     def __init__(
         self,
-        conn: Union[Connection, ConnectionPool],
+        conn: Conn,
         pipe: Optional[Pipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -85,9 +87,13 @@ class PostgresSaver(BasePostgresSaver):
         """
         with self._cursor() as cur:
             try:
-                version = cur.execute(
+                row = cur.execute(
                     "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                ).fetchone()["v"]
+                ).fetchone()
+                if row is None:
+                    version = -1
+                else:
+                    version = row["v"]
             except UndefinedTable:
                 version = -1
             for v, migration in zip(
@@ -212,7 +218,7 @@ class PostgresSaver(BasePostgresSaver):
         checkpoint_id = get_checkpoint_id(config)
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         if checkpoint_id:
-            args = (thread_id, checkpoint_ns, checkpoint_id)
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
             where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
         else:
             args = (thread_id, checkpoint_ns)
@@ -306,7 +312,7 @@ class PostgresSaver(BasePostgresSaver):
                 self._dump_blobs(
                     thread_id,
                     checkpoint_ns,
-                    copy.pop("channel_values"),
+                    copy.pop("channel_values"),  # type: ignore[misc]
                     new_versions,
                 ),
             )
@@ -356,7 +362,7 @@ class PostgresSaver(BasePostgresSaver):
             )
 
     @contextmanager
-    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
+    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
         with _get_connection(self.conn) as conn:
             if self.pipe:
                 # a connection in pipeline mode can be used concurrently