223 lines
8.2 KiB
Python
223 lines
8.2 KiB
Python
"""Persistent local catalog for monthly PostgreSQL chat archives.
|
|
|
|
The archive directory is deliberately independent from the configured
|
|
checkpointer SQLite file: existing SQLite deployments and legacy dual-read
|
|
paths keep their current behavior.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import errno
|
|
import os
|
|
import re
|
|
import sqlite3
|
|
import threading
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import BinaryIO
|
|
|
|
from deerflow.config.runtime_paths import runtime_home
|
|
|
|
_MONTH_RE = re.compile(r"^\d{4}-(0[1-9]|1[0-2])$")
|
|
_WINDOWS_LOCKS: dict[str, threading.Lock] = {}
|
|
_WINDOWS_LOCKS_GUARD = threading.Lock()
|
|
_WINDOWS_LOCK_TYPE = type(threading.Lock())
|
|
|
|
|
|
def archive_root() -> Path:
|
|
"""Return the persistent archive directory used by the gateway/workers."""
|
|
|
|
configured = os.getenv("DEER_FLOW_CHAT_ARCHIVE_DIR", "").strip()
|
|
if configured:
|
|
return Path(configured).expanduser().resolve()
|
|
# The standard deployment keeps ``<DEER_FLOW_HOME>/data`` on a shared,
|
|
# persistent Docker volume. Workers and the gateway therefore see the
|
|
# same archive files and cross-process lock files after restarts.
|
|
return runtime_home() / "data" / "chat-archives"
|
|
|
|
|
|
def archive_db_path(root: Path, month: str) -> Path:
|
|
if not _MONTH_RE.fullmatch(month):
|
|
raise ValueError(f"Invalid archive month: {month!r}")
|
|
return root / f"chat-{month}.db"
|
|
|
|
|
|
def catalog_path(root: Path) -> Path:
|
|
return root / "catalog.db"
|
|
|
|
|
|
def try_acquire_archive_job_lock(root: Path) -> BinaryIO | None:
|
|
"""Hold one archive operation across gateway processes and worker containers.
|
|
|
|
The kernel releases this lock on process exit. Recovery can therefore
|
|
distinguish a dead runner from an older but still active gateway process.
|
|
"""
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
handle = (root / ".archive-operation.lock").open("a+b")
|
|
try:
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
if handle.seek(0, os.SEEK_END) == 0:
|
|
handle.write(b"\0")
|
|
handle.flush()
|
|
handle.seek(0)
|
|
msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
|
|
else:
|
|
import fcntl
|
|
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
return handle
|
|
except OSError as exc:
|
|
handle.close()
|
|
if exc.errno in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}:
|
|
return None
|
|
raise
|
|
except BaseException:
|
|
handle.close()
|
|
raise
|
|
|
|
|
|
def ensure_catalog(root: Path) -> Path:
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
path = catalog_path(root)
|
|
with sqlite3.connect(path, timeout=30) as conn:
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA busy_timeout=30000")
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS archive_months (
|
|
month TEXT PRIMARY KEY,
|
|
status TEXT NOT NULL DEFAULT 'empty',
|
|
total_threads INTEGER NOT NULL DEFAULT 0,
|
|
copied_threads INTEGER NOT NULL DEFAULT 0,
|
|
deleted_threads INTEGER NOT NULL DEFAULT 0,
|
|
restored_threads INTEGER NOT NULL DEFAULT 0,
|
|
current_job_id TEXT,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS archive_threads (
|
|
thread_id TEXT PRIMARY KEY,
|
|
month TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
checkpoint_count INTEGER NOT NULL DEFAULT 0,
|
|
archived_at TEXT NOT NULL,
|
|
deleted_at TEXT,
|
|
restored_at TEXT
|
|
);
|
|
CREATE INDEX IF NOT EXISTS ix_archive_threads_month_status
|
|
ON archive_threads(month, status);
|
|
CREATE TABLE IF NOT EXISTS archive_jobs (
|
|
job_id TEXT PRIMARY KEY,
|
|
month TEXT NOT NULL,
|
|
action TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
total_threads INTEGER NOT NULL DEFAULT 0,
|
|
processed_threads INTEGER NOT NULL DEFAULT 0,
|
|
failed_threads INTEGER NOT NULL DEFAULT 0,
|
|
current_thread_id TEXT,
|
|
snapshot_at TEXT NOT NULL DEFAULT '',
|
|
processed_through TEXT,
|
|
message TEXT NOT NULL DEFAULT '',
|
|
error TEXT,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS ix_archive_jobs_month_created
|
|
ON archive_jobs(month, created_at DESC);
|
|
"""
|
|
)
|
|
# Upgrade catalogs created by earlier releases without replacing data.
|
|
columns = {str(row[1]) for row in conn.execute("PRAGMA table_info(archive_jobs)")}
|
|
if "snapshot_at" not in columns:
|
|
conn.execute("ALTER TABLE archive_jobs ADD COLUMN snapshot_at TEXT NOT NULL DEFAULT ''")
|
|
if "processed_through" not in columns:
|
|
conn.execute("ALTER TABLE archive_jobs ADD COLUMN processed_through TEXT")
|
|
return path
|
|
|
|
|
|
def archived_month_for_thread(root: Path, thread_id: str) -> str | None:
|
|
path = catalog_path(root)
|
|
if not path.is_file():
|
|
return None
|
|
try:
|
|
with sqlite3.connect(f"file:{path.as_posix()}?mode=ro", uri=True, timeout=5) as conn:
|
|
row = conn.execute(
|
|
"SELECT month FROM archive_threads WHERE thread_id = ? AND status IN ('copied', 'deleted', 'restored')",
|
|
(thread_id,),
|
|
).fetchone()
|
|
except sqlite3.Error:
|
|
return None
|
|
return str(row[0]) if row else None
|
|
|
|
|
|
def _acquire_thread_file_lock(root: Path, thread_id: str, exclusive: bool) -> int | _WINDOWS_LOCK_TYPE:
|
|
lock_dir = root / ".locks"
|
|
lock_dir.mkdir(parents=True, exist_ok=True)
|
|
import hashlib
|
|
|
|
# Stripe locks keep the lock directory bounded instead of creating one
|
|
# permanent inode for every conversation ever written.
|
|
name = hashlib.sha256(thread_id.encode("utf-8")).hexdigest()[:3] + ".lock"
|
|
path = lock_dir / name
|
|
try:
|
|
import fcntl
|
|
except ImportError: # pragma: no cover - Windows test/development fallback
|
|
key = str(path.resolve())
|
|
with _WINDOWS_LOCKS_GUARD:
|
|
lock = _WINDOWS_LOCKS.setdefault(key, threading.Lock())
|
|
lock.acquire()
|
|
return lock
|
|
descriptor = os.open(path, os.O_CREAT | os.O_RDWR, 0o600)
|
|
fcntl.flock(descriptor, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
|
|
return descriptor
|
|
|
|
|
|
def _release_thread_file_lock(lock: int | _WINDOWS_LOCK_TYPE) -> None:
|
|
if isinstance(lock, int):
|
|
import fcntl
|
|
|
|
fcntl.flock(lock, fcntl.LOCK_UN)
|
|
os.close(lock)
|
|
else: # pragma: no cover - Windows test/development fallback
|
|
lock.release()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def thread_archive_lock(root: Path, thread_id: str, *, exclusive: bool = False) -> AsyncIterator[None]:
|
|
"""Cross-process read/write gate shared by the main service and workers.
|
|
|
|
Linux ``flock`` readers do not serialize one another. A purge acquires an
|
|
exclusive lock, waits for current checkpoint writes to finish, and holds
|
|
off new writes only for the short copy/verify/delete section. Kernel locks
|
|
are released automatically if a process exits unexpectedly.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
lock = await asyncio.to_thread(_acquire_thread_file_lock, root, thread_id, exclusive)
|
|
try:
|
|
yield
|
|
finally:
|
|
await asyncio.to_thread(_release_thread_file_lock, lock)
|
|
|
|
|
|
def remove_archived_thread(root: Path, thread_id: str) -> None:
|
|
path = catalog_path(root)
|
|
if not path.is_file():
|
|
return
|
|
with sqlite3.connect(path, timeout=30) as conn:
|
|
conn.execute("PRAGMA busy_timeout=30000")
|
|
row = conn.execute("SELECT month, status FROM archive_threads WHERE thread_id = ?", (thread_id,)).fetchone()
|
|
if not row:
|
|
return
|
|
month = str(row[0])
|
|
status = str(row[1])
|
|
conn.execute("DELETE FROM archive_threads WHERE thread_id = ?", (thread_id,))
|
|
conn.execute(
|
|
"UPDATE archive_months SET copied_threads = MAX(0, copied_threads - 1), "
|
|
"deleted_threads = MAX(0, deleted_threads - ?), updated_at = CURRENT_TIMESTAMP WHERE month = ?",
|
|
(1 if status == "deleted" else 0, month),
|
|
)
|