deer-flow/backend/packages/harness/deerflow/config/checkpointer_config.py
Ryker_Feng 063d62c3c3
feat(persistence): support custom postgres schema (#3442)
* feat(persistence): support custom postgres schema

* fix(persistence): address CI lint/test failures and review feedback

- Map missing psycopg import to actionable POSTGRES_INSTALL guidance in
  sync/async schema-creation helpers
- Accept SQLAlchemy compound DSN schemes (postgresql+asyncpg) when
  injecting search_path, normalizing to a libpq-consumable DSN
- Guard keyword-DSN tests with importorskip so they skip without psycopg
- Set database=None in sync checkpointer none-fix test to avoid MagicMock
  backend resolution
- Apply ruff import sort and format

* fix(persistence): address pg-schema review feedback

- Restrict postgres_schema regex to lowercase-only so the quoted CREATE
  SCHEMA matches the unquoted search_path (PG case-folds it), fixing the
  mixed-case bug where tables silently fell back to public.
- Replace shlex.join/split with libpq-correct backslash escaping for the
  options parameter so values containing spaces survive intact.
- Add normalize_libpq_dsn() and route the async checkpointer pool through
  dsn_with_search_path() so a +asyncpg suffix is stripped and existing DSN
  options (e.g. statement_timeout) are merged instead of overridden.
- Extract shared ensure_postgres_schema()/ensure_postgres_schema_async()
  helpers (mapping missing psycopg to the install hint) used by all four
  provider sites.
- Tests: reject mixed-case schemas, preserve space-containing libpq option,
  cover normalize_libpq_dsn, and assert pool search_path via DSN.

* fix(persistence): align pg-schema test with merged store API

The main merge moved the sync Store factory to the single-path
_resolve_store_config/_sync_store_cm design, dropping the PR's
_sync_store_from_database helper. The integration test still imported
the removed symbol, breaking test collection (backend-unit-tests).
Resolve the store config from a DatabaseConfig and drive it through
_sync_store_cm instead.

* fix(persistence): address pg-schema review feedback

- reject trailing/leading whitespace in postgres_schema via re.fullmatch
  (a $-anchored re.match let "deerflow\n" through, silently landing tables
  in public)
- re-escape all whitespace (TAB/CR/LF) when re-joining libpq options so a
  caller's pre-existing options value round-trips losslessly
- re-validate the identifier inside create_schema_sql as defense-in-depth
  at the SQL-emitting boundary
- accept the postgres:// short scheme in the alembic search_path injection
- close the sync psycopg connection explicitly (psycopg3 __exit__ does not
  close()), mirroring the async path
- drop the partial checkpointer/store reset on a database config change;
  database is restart-required and the ORM engine is not rebuilt, so a
  partial reset would half-migrate the deployment

* docs(config): complete the postgres_schema migration checklist

Address PR review (P1): the documented `public`->schema migration only
moved runs, run_events, threads_meta, feedback, and users. That strands
every other DeerFlow-owned table -- the four channel_* tables, both
scheduled_* tables, agents, and (critically) alembic_version -- in
`public`. On restart bootstrap treats the partially-populated target
schema as unversioned, re-baselines it, and replays migrations while the
real rows stay invisible in `public`.

List the full owned set explicitly, call out alembic_version as required,
and keep the "discover the rest" query for version-drift safety.

* refactor(checkpointer): drop test-only _sync_checkpointer_from_database

Address PR review: the helper was only reached by the env-gated
integration test and re-implemented the DatabaseConfig->CheckpointerConfig
backend resolution that _resolve_checkpointer_config already owns, so a
future backend added there would silently miss this path. Mirror the store
side of the same test, which reuses the production path directly:
_resolve_checkpointer_config(...) + _sync_checkpointer_cm(...).
2026-07-30 13:51:17 +08:00

76 lines
2.7 KiB
Python

"""Configuration for LangGraph checkpointer."""
from typing import Literal
from pydantic import BaseModel, Field, field_validator
from deerflow.config.postgres_schema import POSTGRES_SCHEMA_PATTERN, validate_postgres_schema
CheckpointerType = Literal["memory", "sqlite", "postgres"]
class CheckpointerConfig(BaseModel):
"""Configuration for LangGraph state persistence checkpointer."""
type: CheckpointerType = Field(
description="Checkpointer backend type. "
"'memory' is in-process only (lost on restart). "
"'sqlite' persists to a local file (requires langgraph-checkpoint-sqlite). "
"'postgres' persists to PostgreSQL (install with deerflow-harness[postgres])."
)
connection_string: str | None = Field(
default=None,
description="Connection string for sqlite (file path) or postgres (DSN). "
"Optional for sqlite and defaults to 'store.db' when omitted. "
"Required for postgres. "
"For sqlite, use a file path like '.deer-flow/checkpoints.db' or ':memory:' for in-memory. "
"For postgres, use a DSN like 'postgresql://user:pass@localhost:5432/db'.",
)
postgres_schema: str = Field(
default="",
description=(f"PostgreSQL schema for legacy checkpointer/store tables (postgres only). Empty string keeps the server default search_path (usually 'public'). Only plain identifiers are allowed: {POSTGRES_SCHEMA_PATTERN}."),
)
@field_validator("postgres_schema")
@classmethod
def _validate_postgres_schema(cls, value: str) -> str:
return validate_postgres_schema(value)
# Global configuration instance — None means no checkpointer is configured.
_checkpointer_config: CheckpointerConfig | None = None
def get_checkpointer_config() -> CheckpointerConfig | None:
"""Get the current checkpointer configuration, or None if not configured."""
return _checkpointer_config
def set_checkpointer_config(config: CheckpointerConfig | None) -> None:
"""Set the checkpointer configuration."""
global _checkpointer_config
_checkpointer_config = config
def ensure_config_loaded() -> None:
"""Lazily load app config when checkpointer config has not been initialized."""
from deerflow.config.app_config import _app_config, get_app_config
config = get_checkpointer_config()
if config is not None or _app_config is not None:
return
try:
get_app_config()
except FileNotFoundError:
pass
def load_checkpointer_config_from_dict(config_dict: dict | None) -> None:
"""Load checkpointer configuration from a dictionary."""
global _checkpointer_config
if config_dict is None:
_checkpointer_config = None
return
_checkpointer_config = CheckpointerConfig(**config_dict)