123 lines
4.7 KiB
Python
123 lines
4.7 KiB
Python
"""Opt-in end-to-end archive test for a disposable PostgreSQL database.
|
|
|
|
Run with DEERFLOW_TEST_POSTGRES_URL set to an isolated database. This test
|
|
archives, purges, transparently reads, and restores one uniquely named thread.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import uuid
|
|
from datetime import datetime
|
|
|
|
import pytest
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
from psycopg.rows import dict_row
|
|
from psycopg_pool import AsyncConnectionPool
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.routers.chat_archive import ChatArchiveRunner, _save_job, _upsert_month
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
from deerflow.runtime.checkpointer.chat_archive import archive_db_path, ensure_catalog
|
|
from deerflow.runtime.checkpointer.checkpoint_migration import DualReadCheckpointer
|
|
|
|
|
|
async def _checkpoint_ids(saver, thread_id: str) -> list[str]:
|
|
return [
|
|
item.config["configurable"]["checkpoint_id"]
|
|
async for item in saver.alist({"configurable": {"thread_id": thread_id}})
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_postgres_month_archive_purge_fallback_and_rollback(tmp_path):
|
|
conninfo = os.getenv("DEERFLOW_TEST_POSTGRES_URL")
|
|
if not conninfo:
|
|
pytest.skip("set DEERFLOW_TEST_POSTGRES_URL to an isolated PostgreSQL database")
|
|
|
|
thread_id = f"archive-qa-{uuid.uuid4().hex}"
|
|
month = "2025-02"
|
|
pool = AsyncConnectionPool(
|
|
conninfo=conninfo,
|
|
min_size=1,
|
|
max_size=3,
|
|
kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row},
|
|
open=False,
|
|
)
|
|
business_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'business.db'}")
|
|
try:
|
|
await pool.open(wait=True)
|
|
primary = AsyncPostgresSaver(pool)
|
|
await primary.setup()
|
|
|
|
async with business_engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all, tables=[ThreadMetaRow.__table__])
|
|
session_factory = async_sessionmaker(business_engine, expire_on_commit=False)
|
|
created_at = datetime(2025, 2, 12, 12, 0)
|
|
async with session_factory() as session:
|
|
session.add(ThreadMetaRow(
|
|
thread_id=thread_id,
|
|
assistant_id="archive-qa",
|
|
user_id="archive-qa",
|
|
display_name="isolated PostgreSQL archive test",
|
|
status="idle",
|
|
metadata_json={},
|
|
created_at=created_at,
|
|
updated_at=created_at,
|
|
))
|
|
await session.commit()
|
|
|
|
checkpoint = empty_checkpoint()
|
|
await primary.aput(
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
checkpoint,
|
|
{"source": "isolated-postgres-qa"},
|
|
checkpoint["channel_versions"],
|
|
)
|
|
assert await _checkpoint_ids(primary, thread_id) == [checkpoint["id"]]
|
|
|
|
root = tmp_path / "chat-archives"
|
|
ensure_catalog(root)
|
|
job_id = f"archive-qa-{uuid.uuid4().hex}"
|
|
now = "2026-10-25T00:00:00+08:00"
|
|
_save_job(root, {
|
|
"job_id": job_id,
|
|
"month": month,
|
|
"action": "archive",
|
|
"status": "running",
|
|
"total_threads": 1,
|
|
"processed_threads": 0,
|
|
"failed_threads": 0,
|
|
"current_thread_id": None,
|
|
"message": "",
|
|
"error": None,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
})
|
|
_upsert_month(root, month, status="copying", total=1, job_id=job_id)
|
|
runner = ChatArchiveRunner(root, primary, session_factory)
|
|
await runner._archive_month(job_id, month)
|
|
|
|
archived_path = archive_db_path(root, month)
|
|
async with AsyncSqliteSaver.from_conn_string(str(archived_path)) as archived:
|
|
await archived.setup()
|
|
assert await _checkpoint_ids(archived, thread_id) == [checkpoint["id"]]
|
|
|
|
await runner._purge_month(job_id, month)
|
|
assert await _checkpoint_ids(primary, thread_id) == []
|
|
|
|
proxy = DualReadCheckpointer(primary, [], archive_dir=root)
|
|
recovered = await proxy.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
|
|
assert recovered is not None
|
|
assert recovered.checkpoint["id"] == checkpoint["id"]
|
|
await proxy.aclose_sources()
|
|
|
|
await runner._rollback_month(job_id, month)
|
|
assert await _checkpoint_ids(primary, thread_id) == [checkpoint["id"]]
|
|
finally:
|
|
await business_engine.dispose()
|
|
await pool.close()
|