WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-12 20:57:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1257 from langchain-ai/nc/7aug/postgres-migrations

Browse Source 
postgres: Add migration tracking
This commit is contained in:
Nuno Campos
2024-08-07 09:33:10 -07:00
committed by GitHub
parent ca4b65e1a6 54da68cfe7
commit 4c2542aac6
9 changed files with 91 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+13 -18
View File
@@ -4,6 +4,7 @@ from typing import Any, Iterator, List, Optional
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.types.json import Jsonb
@@ -23,8 +24,6 @@ from langgraph.checkpoint.serde.base import SerializerProtocol
class PostgresSaver(BasePostgresSaver):
lock: threading.Lock
is_setup: bool
def __init__(
self,
conn: Connection,
@@ -35,7 +34,6 @@ class PostgresSaver(BasePostgresSaver):
self.conn = conn
self.pipe = pipe
self.lock = threading.Lock()
self.is_setup = False
@classmethod
@contextmanager
@@ -67,23 +65,23 @@ class PostgresSaver(BasePostgresSaver):
already exist. It is called automatically when needed and should not be called
directly by the user.
"""
if self.is_setup:
return
with self.lock:
create_table_queries = [
self.CREATE_CHECKPOINTS_SQL,
self.CREATE_CHECKPOINT_BLOBS_SQL,
self.CREATE_CHECKPOINT_WRITES_SQL,
]
with self.conn.cursor(binary=True) as cur:
for query in create_table_queries:
cur.execute(query)
try:
version = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
).fetchone()["v"]
except UndefinedTable:
version = -1
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
):
cur.execute(migration)
cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})")
if self.pipe:
self.pipe.sync()
self.is_setup = True
def list(
self,
config: Optional[RunnableConfig],
@@ -92,7 +90,6 @@ class PostgresSaver(BasePostgresSaver):
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
self.setup()
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
@@ -124,7 +121,6 @@ class PostgresSaver(BasePostgresSaver):
)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
self.setup()
thread_id = config["configurable"]["thread_id"]
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -235,7 +231,6 @@ class PostgresSaver(BasePostgresSaver):
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
self.setup()
if self.pipe:
# a connection in pipeline mode can be used concurrently
# in multiple threads/coroutines, but only one cursor can be
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+18 -20
View File
@@ -4,6 +4,7 @@ from typing import Any, AsyncIterator, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb
@@ -21,8 +22,6 @@ from langgraph.checkpoint.serde.base import SerializerProtocol
class AsyncPostgresSaver(BasePostgresSaver):
lock: asyncio.Lock
is_setup: bool
def __init__(
self,
conn: AsyncConnection,
@@ -33,7 +32,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
self.conn = conn
self.pipe = pipe
self.lock = asyncio.Lock()
self.is_setup = False
@classmethod
@asynccontextmanager
@@ -65,23 +63,27 @@ class AsyncPostgresSaver(BasePostgresSaver):
already exist. It is called automatically when needed and should not be called
directly by the user.
"""
if self.is_setup:
return
async with self.lock:
create_table_queries = [
self.CREATE_CHECKPOINTS_SQL,
self.CREATE_CHECKPOINT_BLOBS_SQL,
self.CREATE_CHECKPOINT_WRITES_SQL,
]
async with self.conn.cursor() as cur:
for query in create_table_queries:
await cur.execute(query)
async with self.conn.cursor(binary=True) as cur:
try:
version = (
await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
).fetchone()["v"]
except UndefinedTable:
version = -1
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
):
await cur.execute(migration)
await cur.execute(
f"INSERT INTO checkpoint_migrations (v) VALUES ({v})"
)
if self.pipe:
await self.pipe.sync()
self.is_setup = True
async def alist(
self,
config: Optional[RunnableConfig],
@@ -90,7 +92,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[CheckpointTuple]:
await self.setup()
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
@@ -124,7 +125,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
await self.setup()
thread_id = config["configurable"]["thread_id"]
checkpoint_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -177,7 +177,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
await self.setup()
configurable = config["configurable"].copy()
thread_id = configurable.pop("thread_id")
checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -240,7 +239,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
@asynccontextmanager
async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
await self.setup()
if self.pipe:
# a connection in pipeline mode can be used concurrently
# in multiple threads/coroutines, but only one cursor can be
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+41 -40
View File
@@ -16,6 +16,46 @@ from langgraph.checkpoint.serde.types import ChannelProtocol
MetadataInput = Optional[dict[str, Any]]
"""
To add a new migration, add a new string to the MIGRATIONS list.
The position of the migration in the list is the version number.
"""
MIGRATIONS = [
"""CREATE TABLE IF NOT EXISTS checkpoint_migrations (
v INTEGER PRIMARY KEY
);""",
"""CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
parent_checkpoint_id TEXT,
type TEXT,
checkpoint JSONB NOT NULL,
metadata JSONB NOT NULL DEFAULT '{}',
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);""",
"""CREATE TABLE IF NOT EXISTS checkpoint_blobs (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
channel TEXT NOT NULL,
version TEXT NOT NULL,
type TEXT NOT NULL,
blob BYTEA NOT NULL,
PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
);""",
"""CREATE TABLE IF NOT EXISTS checkpoint_writes (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
type TEXT,
blob BYTEA NOT NULL,
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);""",
]
SELECT_SQL = """
select
thread_id,
@@ -43,43 +83,6 @@ select
) as pending_writes
from checkpoints """
CREATE_CHECKPOINTS_SQL = """
CREATE TABLE IF NOT EXISTS checkpoints (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
parent_checkpoint_id TEXT,
type TEXT,
checkpoint JSONB NOT NULL,
metadata JSONB NOT NULL DEFAULT '{}',
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);
"""
CREATE_CHECKPOINT_BLOBS_SQL = """
CREATE TABLE IF NOT EXISTS checkpoint_blobs (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
channel TEXT NOT NULL,
version TEXT NOT NULL,
type TEXT NOT NULL,
blob BYTEA NOT NULL,
PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
);"""
CREATE_CHECKPOINT_WRITES_SQL = """
CREATE TABLE IF NOT EXISTS checkpoint_writes (
thread_id TEXT NOT NULL,
checkpoint_ns TEXT NOT NULL DEFAULT '',
checkpoint_id TEXT NOT NULL,
task_id TEXT NOT NULL,
idx INTEGER NOT NULL,
channel TEXT NOT NULL,
type TEXT,
blob BYTEA NOT NULL,
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);
"""
UPSERT_CHECKPOINT_BLOBS_SQL = """
INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
VALUES (%s, %s, %s, %s, %s, %s)
@@ -104,9 +107,7 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
class BasePostgresSaver(BaseCheckpointSaver):
SELECT_SQL = SELECT_SQL
CREATE_CHECKPOINTS_SQL = CREATE_CHECKPOINTS_SQL
CREATE_CHECKPOINT_BLOBS_SQL = CREATE_CHECKPOINT_BLOBS_SQL
CREATE_CHECKPOINT_WRITES_SQL = CREATE_CHECKPOINT_WRITES_SQL
MIGRATIONS = MIGRATIONS
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
libs/checkpoint-postgres/tests/compose-postgres.yml
+2 -2
View File
@@ -2,7 +2,7 @@ services:
postgres-test:
image: postgres:16
ports:
- "5432:5432"
- "5441:5432"
environment:
POSTGRES_DB: postgres
POSTGRES_USER: postgres
@@ -13,4 +13,4 @@ services:
timeout: 1s
retries: 5
interval: 60s
start_interval: 1s
start_interval: 1s
libs/checkpoint-postgres/tests/conftest.py
+1 -1
View File
@@ -3,7 +3,7 @@ from psycopg import AsyncConnection
from psycopg.errors import UndefinedTable
from psycopg.rows import dict_row
DEFAULT_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
@pytest.fixture(scope="function")
libs/checkpoint-postgres/tests/test_async.py
+3 -1
View File
@@ -13,7 +13,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
class TestAsyncPostgresSaver:
@pytest.fixture(autouse=True)
def setup(self):
async def setup(self):
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
@@ -55,6 +55,8 @@ class TestAsyncPostgresSaver:
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
async def test_asearch(self):
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
libs/checkpoint-postgres/tests/test_sync.py
+2
View File
@@ -55,6 +55,8 @@ class TestPostgresSaver:
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
saver.setup()
def test_search(self):
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
libs/langgraph/tests/compose-postgres.yml
+2 -2
View File
@@ -2,7 +2,7 @@ services:
postgres-test:
image: postgres:16
ports:
- "5432:5432"
- "5442:5432"
environment:
POSTGRES_DB: postgres
POSTGRES_USER: postgres
@@ -13,4 +13,4 @@ services:
timeout: 1s
retries: 5
interval: 60s
start_interval: 1s
start_interval: 1s
libs/langgraph/tests/conftest.py
+9 -1
View File
@@ -6,8 +6,10 @@ from psycopg.errors import UndefinedTable
from psycopg.rows import dict_row
from pytest_mock import MockerFixture
from langgraph.checkpoint.postgres import PostgresSaver
DEFAULT_POSTGRES_URI = (
"postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
"postgres://postgres:postgres@localhost:5442/postgres?sslmode=disable"
)
@@ -27,6 +29,12 @@ async def conn():
yield conn
@pytest.fixture(scope="session", autouse=True)
def setup_before_all_tests():
with PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI) as checkpointer:
checkpointer.setup()
@pytest.fixture(scope="function", autouse=True)
async def clear_test_db(conn):
"""Delete all tables before each test."""