sync update from LAN repo (source dhk-1 @ 3b1ec5b4)

offline-backend-20260512/backend/tests/test_chat_archive_postgres_integration.py

@@ -0,0 +1,122 @@
+"""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()