399 lines
15 KiB
Python
399 lines
15 KiB
Python
"""SQLite-to-Postgres checkpoint migration and read fallback helpers.
|
|
|
|
The live gateway can wrap a Postgres checkpointer with ``DualReadCheckpointer``.
|
|
New writes go to Postgres, reads prefer Postgres and fall back to legacy SQLite.
|
|
Manual migration scripts reuse the copy helpers in this module.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import sqlite3
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from deerflow.runtime.checkpointer.chat_archive import (
|
|
archive_db_path,
|
|
archived_month_for_thread,
|
|
catalog_path,
|
|
remove_archived_thread,
|
|
thread_archive_lock,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class SQLiteCheckpointSource:
|
|
path: str
|
|
saver: Any
|
|
context: Any
|
|
progress_path: Path
|
|
closed: bool = False
|
|
|
|
async def close(self) -> None:
|
|
if self.closed:
|
|
return
|
|
self.closed = True
|
|
try:
|
|
await self.context.__aexit__(None, None, None)
|
|
except Exception:
|
|
logger.warning("Failed to close legacy SQLite checkpointer: %s", self.path, exc_info=True)
|
|
|
|
|
|
def _sqlite_db_exists(path: str) -> bool:
|
|
return Path(path).exists()
|
|
|
|
|
|
def _sqlite_threads(path: str) -> list[str]:
|
|
with sqlite3.connect(path) as conn:
|
|
rows = conn.execute("SELECT DISTINCT thread_id FROM checkpoints ORDER BY thread_id").fetchall()
|
|
return [str(row[0]) for row in rows if row and row[0]]
|
|
|
|
|
|
def _sqlite_has_table(path: str, table: str) -> bool:
|
|
with sqlite3.connect(path) as conn:
|
|
row = conn.execute(
|
|
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?",
|
|
(table,),
|
|
).fetchone()
|
|
return row is not None
|
|
|
|
|
|
def sqlite_recent_threads(path: str, *, since: str | None = None, until: str | None = None) -> list[str]:
|
|
"""Return threads updated since an ISO timestamp when ``threads_meta`` exists."""
|
|
|
|
if not _sqlite_has_table(path, "threads_meta"):
|
|
return []
|
|
query = "SELECT thread_id FROM threads_meta"
|
|
where: list[str] = []
|
|
params: list[str] = []
|
|
if since:
|
|
where.append("updated_at >= ?")
|
|
params.append(since)
|
|
if until:
|
|
where.append("updated_at <= ?")
|
|
params.append(until)
|
|
if where:
|
|
query += " WHERE " + " AND ".join(where)
|
|
query += " ORDER BY updated_at DESC"
|
|
with sqlite3.connect(path) as conn:
|
|
rows = conn.execute(query, tuple(params)).fetchall()
|
|
return [str(row[0]) for row in rows if row and row[0]]
|
|
|
|
|
|
def _sqlite_thread_checkpoint_count(path: str, thread_id: str) -> int:
|
|
with sqlite3.connect(path) as conn:
|
|
row = conn.execute(
|
|
"SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?",
|
|
(thread_id,),
|
|
).fetchone()
|
|
return int(row[0] if row else 0)
|
|
|
|
|
|
def _load_progress(path: Path) -> set[str]:
|
|
try:
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
except FileNotFoundError:
|
|
return set()
|
|
except Exception:
|
|
logger.warning("Ignoring unreadable checkpoint migration progress file: %s", path, exc_info=True)
|
|
return set()
|
|
migrated = payload.get("migrated_threads", [])
|
|
return {str(x) for x in migrated if x}
|
|
|
|
|
|
def _save_progress(path: Path, migrated: set[str]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_suffix(path.suffix + ".tmp")
|
|
tmp.write_text(
|
|
json.dumps({"migrated_threads": sorted(migrated)}, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
tmp.replace(path)
|
|
|
|
|
|
async def _copy_thread_state(source, target, thread_id: str) -> bool:
|
|
"""Copy one thread through the public LangGraph checkpointer API."""
|
|
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
tuples = [t async for t in source.alist(config)]
|
|
if not tuples:
|
|
return False
|
|
|
|
for tup in reversed(tuples):
|
|
configurable = (tup.config or {}).get("configurable", {}) or {}
|
|
checkpoint_ns = configurable.get("checkpoint_ns", "") or ""
|
|
parent_id = ((tup.parent_config or {}).get("configurable", {}) or {}).get("checkpoint_id")
|
|
|
|
put_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}}
|
|
if parent_id:
|
|
put_config["configurable"]["checkpoint_id"] = parent_id
|
|
|
|
new_versions = (tup.checkpoint or {}).get("channel_versions", {}) or {}
|
|
await target.aput(put_config, tup.checkpoint, tup.metadata or {}, new_versions)
|
|
|
|
if tup.pending_writes:
|
|
by_task: dict[str, list] = {}
|
|
for task_id, channel, value in tup.pending_writes:
|
|
by_task.setdefault(task_id, []).append((channel, value))
|
|
checkpoint_id = (tup.checkpoint or {}).get("id")
|
|
for task_id, writes in by_task.items():
|
|
writes_config = {
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": checkpoint_id,
|
|
}
|
|
}
|
|
await target.aput_writes(writes_config, writes, task_id)
|
|
return True
|
|
|
|
|
|
async def copy_thread_state(source, target, thread_id: str) -> bool:
|
|
return await _copy_thread_state(source, target, thread_id)
|
|
|
|
|
|
async def _target_thread_checkpoint_count(target, thread_id: str) -> int:
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
count = 0
|
|
async for _ in target.alist(config):
|
|
count += 1
|
|
return count
|
|
|
|
|
|
async def target_has_thread(target, thread_id: str) -> bool:
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
async for _ in target.alist(config, limit=1):
|
|
return True
|
|
return False
|
|
|
|
|
|
class DualReadCheckpointer:
|
|
"""Postgres-first checkpointer with legacy and monthly archive fallbacks."""
|
|
|
|
_PRIMARY_ONLY_MAX = 8192
|
|
|
|
def __init__(
|
|
self,
|
|
primary,
|
|
sources: list[SQLiteCheckpointSource],
|
|
*,
|
|
archive_dir: Path | None = None,
|
|
sleep_seconds: float = 0.05,
|
|
):
|
|
self.primary = primary
|
|
self.sources = sources
|
|
self.archive_dir = archive_dir
|
|
self.sleep_seconds = max(0.0, sleep_seconds)
|
|
self._closed = False
|
|
self._archive_sources: dict[str, SQLiteCheckpointSource] = {}
|
|
self._archive_sources_lock = asyncio.Lock()
|
|
# Threads created or written after dual-read started cannot exist in
|
|
# the legacy SQLite files. Skipping that fallback avoids opening
|
|
# multi-GB sqlite sources on every empty /state or /history read.
|
|
self._primary_only: set[str] = set()
|
|
|
|
def __getattr__(self, name: str):
|
|
return getattr(self.primary, name)
|
|
|
|
def remember_primary_thread(self, thread_id: str | None) -> None:
|
|
"""Mark *thread_id* as primary-only so empty reads skip SQLite."""
|
|
if not thread_id:
|
|
return
|
|
self._primary_only.add(str(thread_id))
|
|
overflow = len(self._primary_only) - self._PRIMARY_ONLY_MAX
|
|
if overflow > 0:
|
|
for stale in list(self._primary_only)[:overflow]:
|
|
self._primary_only.discard(stale)
|
|
|
|
def _skip_legacy_fallback(self, config) -> bool:
|
|
thread_id = ((config or {}).get("configurable") or {}).get("thread_id")
|
|
return bool(thread_id) and str(thread_id) in self._primary_only
|
|
|
|
async def aclose_sources(self) -> None:
|
|
self._closed = True
|
|
for source in [*self.sources, *self._archive_sources.values()]:
|
|
await source.close()
|
|
self._archive_sources.clear()
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def thread_write_lock(self, thread_id: str, *, exclusive: bool = False):
|
|
"""Coordinate archive deletion with writes across all PG-backed workers."""
|
|
if not self.archive_dir or not catalog_path(self.archive_dir).is_file():
|
|
yield
|
|
return
|
|
async with thread_archive_lock(self.archive_dir, str(thread_id), exclusive=exclusive):
|
|
yield
|
|
|
|
async def _archive_source_for_thread(self, thread_id: str | None) -> SQLiteCheckpointSource | None:
|
|
if not self.archive_dir or not thread_id or self._closed:
|
|
return None
|
|
month = await asyncio.to_thread(archived_month_for_thread, self.archive_dir, thread_id)
|
|
if not month:
|
|
return None
|
|
cached = self._archive_sources.get(month)
|
|
if cached is not None and not cached.closed:
|
|
return cached
|
|
async with self._archive_sources_lock:
|
|
cached = self._archive_sources.get(month)
|
|
if cached is not None and not cached.closed:
|
|
return cached
|
|
path = archive_db_path(self.archive_dir, month)
|
|
if not path.is_file():
|
|
return None
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
|
|
context = AsyncSqliteSaver.from_conn_string(str(path))
|
|
saver = await context.__aenter__()
|
|
await saver.conn.execute("PRAGMA busy_timeout=30000")
|
|
if not await asyncio.to_thread(_sqlite_has_table, str(path), "checkpoints"):
|
|
await saver.setup()
|
|
source = SQLiteCheckpointSource(
|
|
path=str(path),
|
|
saver=saver,
|
|
context=context,
|
|
progress_path=path.with_suffix(path.suffix + ".archive-progress.json"),
|
|
)
|
|
self._archive_sources[month] = source
|
|
return source
|
|
|
|
async def aget_tuple(self, config):
|
|
found = await self.primary.aget_tuple(config)
|
|
if found is not None or self._skip_legacy_fallback(config):
|
|
return found
|
|
for source in self.sources:
|
|
if source.closed:
|
|
continue
|
|
found = await source.saver.aget_tuple(config)
|
|
if found is not None:
|
|
return found
|
|
configurable = (config or {}).get("configurable", {}) or {}
|
|
archive = await self._archive_source_for_thread(configurable.get("thread_id"))
|
|
if archive is not None:
|
|
found = await archive.saver.aget_tuple(config)
|
|
if found is not None:
|
|
return found
|
|
return None
|
|
|
|
async def aget(self, config):
|
|
found = await self.primary.aget(config)
|
|
if found is not None or self._skip_legacy_fallback(config):
|
|
return found
|
|
for source in self.sources:
|
|
if source.closed:
|
|
continue
|
|
found = await source.saver.aget(config)
|
|
if found is not None:
|
|
return found
|
|
configurable = (config or {}).get("configurable", {}) or {}
|
|
archive = await self._archive_source_for_thread(configurable.get("thread_id"))
|
|
if archive is not None:
|
|
found = await archive.saver.aget(config)
|
|
if found is not None:
|
|
return found
|
|
return None
|
|
|
|
async def alist(self, config, *, filter=None, before=None, limit=None):
|
|
seen: set[tuple[str, str, str]] = set()
|
|
yielded = 0
|
|
|
|
async for tup in self.primary.alist(config, filter=filter, before=before, limit=limit):
|
|
key = self._tuple_key(tup)
|
|
seen.add(key)
|
|
yield tup
|
|
yielded += 1
|
|
if limit is not None and yielded >= limit:
|
|
return
|
|
|
|
if self._skip_legacy_fallback(config):
|
|
return
|
|
|
|
remaining = None if limit is None else max(0, limit - yielded)
|
|
for source in self.sources:
|
|
if source.closed:
|
|
continue
|
|
async for tup in source.saver.alist(config, filter=filter, before=before, limit=remaining):
|
|
key = self._tuple_key(tup)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
yield tup
|
|
yielded += 1
|
|
if limit is not None:
|
|
remaining = max(0, limit - yielded)
|
|
if remaining <= 0:
|
|
return
|
|
configurable = (config or {}).get("configurable", {}) or {}
|
|
archive = await self._archive_source_for_thread(configurable.get("thread_id"))
|
|
if archive is not None and not archive.closed:
|
|
async for tup in archive.saver.alist(config, filter=filter, before=before, limit=remaining):
|
|
key = self._tuple_key(tup)
|
|
if key in seen:
|
|
continue
|
|
yield tup
|
|
yielded += 1
|
|
if limit is not None and yielded >= limit:
|
|
return
|
|
|
|
async def aput(self, *args, **kwargs):
|
|
config = args[0] if args else kwargs.get("config")
|
|
if isinstance(config, dict):
|
|
self.remember_primary_thread(((config.get("configurable") or {}).get("thread_id")))
|
|
return await self.primary.aput(*args, **kwargs)
|
|
|
|
async def aput_writes(self, *args, **kwargs):
|
|
config = args[0] if args else kwargs.get("config")
|
|
thread_id = ((config or {}).get("configurable", {}) or {}).get("thread_id")
|
|
if not thread_id:
|
|
return await self.primary.aput_writes(*args, **kwargs)
|
|
async with self.thread_write_lock(str(thread_id)):
|
|
return await self.primary.aput_writes(*args, **kwargs)
|
|
|
|
async def adelete_thread(self, thread_id: str) -> None:
|
|
async with self.thread_write_lock(str(thread_id), exclusive=True):
|
|
delete = getattr(self.primary, "adelete_thread", None)
|
|
if delete is not None:
|
|
await delete(thread_id)
|
|
for source in self.sources:
|
|
if source.closed:
|
|
continue
|
|
delete = getattr(source.saver, "adelete_thread", None)
|
|
if delete is not None:
|
|
await delete(thread_id)
|
|
archive = await self._archive_source_for_thread(str(thread_id))
|
|
if archive is not None:
|
|
delete = getattr(archive.saver, "adelete_thread", None)
|
|
if delete is not None:
|
|
await delete(thread_id)
|
|
if self.archive_dir and catalog_path(self.archive_dir).is_file():
|
|
await asyncio.to_thread(remove_archived_thread, self.archive_dir, str(thread_id))
|
|
|
|
@staticmethod
|
|
def _tuple_key(tup) -> tuple[str, str, str]:
|
|
configurable = (tup.config or {}).get("configurable", {}) or {}
|
|
return (
|
|
str(configurable.get("thread_id", "")),
|
|
str(configurable.get("checkpoint_ns", "") or ""),
|
|
str(configurable.get("checkpoint_id", "") or (tup.checkpoint or {}).get("id", "")),
|
|
)
|
|
|
|
async def _verify_source(self, source: SQLiteCheckpointSource, threads: list[str]) -> bool:
|
|
for thread_id in threads:
|
|
source_count = await asyncio.to_thread(_sqlite_thread_checkpoint_count, source.path, thread_id)
|
|
target_count = await _target_thread_checkpoint_count(self.primary, thread_id)
|
|
if target_count < source_count:
|
|
logger.warning(
|
|
"Checkpoint migration verification failed: source=%s thread_id=%s source_count=%d target_count=%d",
|
|
source.path,
|
|
thread_id,
|
|
source_count,
|
|
target_count,
|
|
)
|
|
return False
|
|
return True
|