mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-31 10:26:04 +00:00
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(...).
This commit is contained in:
Ryker_Feng
GitHub
parent
904cee4a72
commit
063d62c3c3
@ -446,6 +446,16 @@ class AppConfig(BaseModel):
|
||||
if previous_checkpointer_config != config.checkpointer:
|
||||
# These runtime singletons derive their backend from checkpointer config.
|
||||
# Keep imports local to avoid cycles: both providers import get_app_config.
|
||||
#
|
||||
# The unified ``database`` section is intentionally NOT handled here.
|
||||
# ``database`` is a restart-required field (reload_boundary.STARTUP_ONLY_FIELDS):
|
||||
# ``init_engine_from_config()`` builds the ORM engine once at startup and
|
||||
# never rebuilds it on a config.yaml edit. Resetting only the sync
|
||||
# checkpointer/store singletons on a live ``database``/``postgres_schema``
|
||||
# change would half-migrate the deployment -- new checkpoint/store tables
|
||||
# would land in the new schema while ORM rows keep landing in the old one,
|
||||
# with no error surfaced. Requiring the documented restart keeps the
|
||||
# deployment self-consistent.
|
||||
from deerflow.runtime.checkpointer import reset_checkpointer
|
||||
from deerflow.runtime.store import reset_store
|
||||
|
||||
|
||||
@ -2,7 +2,9 @@
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
from deerflow.config.postgres_schema import POSTGRES_SCHEMA_PATTERN, validate_postgres_schema
|
||||
|
||||
CheckpointerType = Literal["memory", "sqlite", "postgres"]
|
||||
|
||||
@ -24,6 +26,15 @@ class CheckpointerConfig(BaseModel):
|
||||
"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.
|
||||
|
||||
@ -35,7 +35,9 @@ import logging
|
||||
import os
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
from deerflow.config.postgres_schema import POSTGRES_SCHEMA_PATTERN, validate_postgres_schema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -151,6 +153,22 @@ class DatabaseConfig(BaseModel):
|
||||
gt=0,
|
||||
description="Timeout in seconds for app ORM PostgreSQL commands. Set to null to disable the command timeout.",
|
||||
)
|
||||
postgres_schema: str = Field(
|
||||
default="",
|
||||
description=(
|
||||
"PostgreSQL schema for both app ORM tables and LangGraph "
|
||||
"checkpointer/store tables (postgres only). Empty string keeps "
|
||||
"the server default search_path (usually 'public'). When set, "
|
||||
"the schema is created automatically at startup and applied via "
|
||||
"connection-level search_path. Only plain identifiers are "
|
||||
f"allowed: {POSTGRES_SCHEMA_PATTERN}."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("postgres_schema")
|
||||
@classmethod
|
||||
def _validate_postgres_schema(cls, value: str) -> str:
|
||||
return validate_postgres_schema(value)
|
||||
|
||||
# -- Legacy key migration (not user-configured) --
|
||||
|
||||
|
||||
27
backend/packages/harness/deerflow/config/postgres_schema.py
Normal file
27
backend/packages/harness/deerflow/config/postgres_schema.py
Normal file
@ -0,0 +1,27 @@
|
||||
"""Shared validation for PostgreSQL schema names."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
# Lowercase-only on purpose. The schema is created *quoted*
|
||||
# (``CREATE SCHEMA IF NOT EXISTS "<schema>"``, case-preserved) but pinned via an
|
||||
# *unquoted* ``search_path`` token, which PostgreSQL folds to lowercase. Allowing
|
||||
# uppercase here would let the two diverge so tables silently land in ``public``.
|
||||
# No anchors: validation uses ``re.fullmatch`` so the whole value must match.
|
||||
# A ``$``-anchored ``re.match`` would accept a trailing newline (``"deerflow\n"``):
|
||||
# Python's ``$`` matches just before a single trailing ``\n``, which then
|
||||
# creates a *quoted* schema literally named ``deerflow\n`` while the *unquoted*
|
||||
# ``search_path`` folds to ``deerflow`` and misses it, silently landing tables
|
||||
# in ``public``.
|
||||
POSTGRES_SCHEMA_PATTERN = r"[a-z_][a-z0-9_]{0,62}"
|
||||
_POSTGRES_SCHEMA_RE = re.compile(POSTGRES_SCHEMA_PATTERN)
|
||||
|
||||
|
||||
def validate_postgres_schema(value: str) -> str:
|
||||
"""Validate the v1 plain-identifier PostgreSQL schema contract."""
|
||||
if value == "":
|
||||
return value
|
||||
if not _POSTGRES_SCHEMA_RE.fullmatch(value):
|
||||
raise ValueError(f"postgres_schema must be a plain lowercase PostgreSQL identifier matching {POSTGRES_SCHEMA_PATTERN}; got {value!r}. Mixed-case and quoted identifiers are not supported.")
|
||||
return value
|
||||
@ -234,16 +234,25 @@ def _alembic_safe_url(engine: AsyncEngine) -> str:
|
||||
return _escape_url_for_alembic(rendered)
|
||||
|
||||
|
||||
def _get_alembic_config(engine: AsyncEngine) -> AlembicConfig:
|
||||
def _get_alembic_config(engine: AsyncEngine, *, postgres_schema: str = "") -> AlembicConfig:
|
||||
"""Build an in-process alembic config pointing at our migrations dir.
|
||||
|
||||
Avoids reading ``alembic.ini`` from disk so the production runtime doesn't
|
||||
depend on a working-directory-relative file lookup. The ``script_location``
|
||||
is anchored at the package path on disk.
|
||||
|
||||
When *postgres_schema* is set it is forwarded as the ``deerflow_pg_schema``
|
||||
main option so ``env.py`` can pin its alembic-spawned engine's
|
||||
``search_path`` to the same schema the app engine uses. Without it,
|
||||
alembic's own engine -- built from the bare URL -- would create
|
||||
``alembic_version`` and all migration DDL in the default (``public``)
|
||||
schema while the app tables land in the custom schema.
|
||||
"""
|
||||
cfg = AlembicConfig()
|
||||
cfg.set_main_option("script_location", str(_MIGRATIONS_DIR))
|
||||
cfg.set_main_option("sqlalchemy.url", _alembic_safe_url(engine))
|
||||
if postgres_schema:
|
||||
cfg.set_main_option("deerflow_pg_schema", postgres_schema)
|
||||
return cfg
|
||||
|
||||
|
||||
@ -470,7 +479,7 @@ def _bootstrap_lock(engine: AsyncEngine, *, backend: str):
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def bootstrap_schema(engine: AsyncEngine, *, backend: str) -> None:
|
||||
async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema: str = "") -> None:
|
||||
"""Bring the DB schema to head.
|
||||
|
||||
Postgres calls are serialised across processes with an advisory lock.
|
||||
@ -480,9 +489,14 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str) -> None:
|
||||
Branch dispatch is documented at module top. ``alembic.command.stamp`` and
|
||||
``alembic.command.upgrade`` are synchronous and would block the event
|
||||
loop; both are wrapped in ``asyncio.to_thread``.
|
||||
|
||||
*postgres_schema*, when set, is forwarded to the alembic config so the
|
||||
alembic-spawned engine pins its ``search_path`` to that schema. The target
|
||||
schema must already exist (``init_engine`` issues ``CREATE SCHEMA`` before
|
||||
calling this). Ignored for non-postgres backends.
|
||||
"""
|
||||
head = _get_head_revision()
|
||||
cfg = _get_alembic_config(engine)
|
||||
cfg = _get_alembic_config(engine, postgres_schema=postgres_schema if backend == "postgres" else "")
|
||||
|
||||
async with _bootstrap_lock(engine, backend=backend):
|
||||
async with engine.connect() as conn:
|
||||
|
||||
@ -33,17 +33,18 @@ def _postgres_engine_kwargs(
|
||||
pool_size: int,
|
||||
pool_recycle: int = POSTGRES_POOL_RECYCLE_SECONDS,
|
||||
command_timeout: float | None = POSTGRES_COMMAND_TIMEOUT_SECONDS,
|
||||
connect_args: dict[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""Build the shared SQLAlchemy engine options for PostgreSQL."""
|
||||
connect_args = {}
|
||||
merged_connect_args = dict(connect_args or {})
|
||||
if command_timeout is not None:
|
||||
connect_args["command_timeout"] = command_timeout
|
||||
merged_connect_args["command_timeout"] = command_timeout
|
||||
return {
|
||||
"echo": echo,
|
||||
"pool_size": pool_size,
|
||||
"pool_pre_ping": True,
|
||||
"pool_recycle": pool_recycle,
|
||||
"connect_args": connect_args,
|
||||
"connect_args": merged_connect_args,
|
||||
"json_serializer": _json_serializer,
|
||||
}
|
||||
|
||||
@ -90,6 +91,7 @@ async def init_engine(
|
||||
pool_recycle: int = POSTGRES_POOL_RECYCLE_SECONDS,
|
||||
command_timeout: float | None = POSTGRES_COMMAND_TIMEOUT_SECONDS,
|
||||
sqlite_dir: str = "",
|
||||
postgres_schema: str = "",
|
||||
) -> None:
|
||||
"""Create the async engine and session factory, then auto-create tables.
|
||||
|
||||
@ -101,6 +103,10 @@ async def init_engine(
|
||||
pool_recycle: Seconds before Postgres connections are recycled.
|
||||
command_timeout: Timeout in seconds for app ORM Postgres commands, or None to disable.
|
||||
sqlite_dir: Directory to create for SQLite (ensured to exist).
|
||||
postgres_schema: Target PostgreSQL schema. When set, the engine
|
||||
pins the connection ``search_path`` to it via asyncpg
|
||||
``server_settings`` and the schema is created (if missing)
|
||||
before tables are auto-created. Ignored for non-postgres.
|
||||
"""
|
||||
global _engine, _session_factory
|
||||
|
||||
@ -161,6 +167,9 @@ async def init_engine(
|
||||
finally:
|
||||
cursor.close()
|
||||
elif backend == "postgres":
|
||||
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
|
||||
|
||||
pg_connect_args = build_asyncpg_connect_args(postgres_schema)
|
||||
_engine = create_async_engine(
|
||||
url,
|
||||
**_postgres_engine_kwargs(
|
||||
@ -168,6 +177,7 @@ async def init_engine(
|
||||
pool_size=pool_size,
|
||||
pool_recycle=pool_recycle,
|
||||
command_timeout=command_timeout,
|
||||
connect_args=pg_connect_args,
|
||||
),
|
||||
)
|
||||
else:
|
||||
@ -186,13 +196,28 @@ async def init_engine(
|
||||
# See deerflow.persistence.bootstrap for the full state machine.
|
||||
from deerflow.persistence.bootstrap import bootstrap_schema
|
||||
|
||||
async def _ensure_postgres_schema() -> None:
|
||||
# CREATE SCHEMA is DDL and is unaffected by search_path, so it is
|
||||
# safe even though the connection's search_path already points at
|
||||
# the (not-yet-existing) target schema. It must run before
|
||||
# ``bootstrap_schema`` so the subsequent ``create_all`` / alembic
|
||||
# DDL lands in the target schema instead of failing on a missing one.
|
||||
if backend == "postgres" and postgres_schema:
|
||||
from sqlalchemy.schema import CreateSchema
|
||||
|
||||
async with _engine.begin() as conn:
|
||||
await conn.execute(CreateSchema(postgres_schema, if_not_exists=True))
|
||||
|
||||
try:
|
||||
await bootstrap_schema(_engine, backend=backend)
|
||||
await _ensure_postgres_schema()
|
||||
await bootstrap_schema(_engine, backend=backend, postgres_schema=postgres_schema)
|
||||
except Exception as exc:
|
||||
if backend == "postgres" and "does not exist" in str(exc):
|
||||
# Database not yet created -- attempt to auto-create it, then retry.
|
||||
await _auto_create_postgres_db(url)
|
||||
# Rebuild engine against the now-existing database
|
||||
# Rebuild engine against the now-existing database. The rebuilt
|
||||
# engine MUST keep the same connect_args so the retried bootstrap
|
||||
# lands in the target schema, not the default one.
|
||||
await _engine.dispose()
|
||||
_engine = create_async_engine(
|
||||
url,
|
||||
@ -201,10 +226,12 @@ async def init_engine(
|
||||
pool_size=pool_size,
|
||||
pool_recycle=pool_recycle,
|
||||
command_timeout=command_timeout,
|
||||
connect_args=pg_connect_args,
|
||||
),
|
||||
)
|
||||
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
||||
await bootstrap_schema(_engine, backend=backend)
|
||||
await _ensure_postgres_schema()
|
||||
await bootstrap_schema(_engine, backend=backend, postgres_schema=postgres_schema)
|
||||
else:
|
||||
raise
|
||||
|
||||
@ -224,6 +251,7 @@ async def init_engine_from_config(config) -> None:
|
||||
pool_recycle=config.pool_recycle,
|
||||
command_timeout=config.command_timeout,
|
||||
sqlite_dir=config.sqlite_dir if config.backend == "sqlite" else "",
|
||||
postgres_schema=config.postgres_schema if config.backend == "postgres" else "",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@ -72,7 +72,26 @@ def do_run_migrations(connection):
|
||||
|
||||
|
||||
async def run_migrations_online() -> None:
|
||||
connectable = create_async_engine(config.get_main_option("sqlalchemy.url"))
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
# When a custom Postgres schema is configured, pin the alembic-spawned
|
||||
# engine's search_path to it. This engine is built from the bare URL and
|
||||
# does NOT inherit the asyncpg ``server_settings`` the app engine sets via
|
||||
# ``connect_args``, so without this both ``alembic_version`` and every
|
||||
# migration's DDL would land in the default (``public``) schema while the
|
||||
# ORM tables land in the custom schema. ``init_engine`` has already created
|
||||
# the schema (``CREATE SCHEMA IF NOT EXISTS``) before bootstrap runs.
|
||||
pg_schema = config.get_main_option("deerflow_pg_schema")
|
||||
connect_args: dict = {}
|
||||
# Accept both the canonical ``postgresql`` scheme and libpq's ``postgres``
|
||||
# short scheme (with or without a SQLAlchemy ``+driver`` suffix) so a
|
||||
# ``postgres://`` DSN still gets its search_path pinned instead of silently
|
||||
# writing ``alembic_version`` + migration DDL to the default schema.
|
||||
if pg_schema and url and url.split("+", 1)[0].split(":", 1)[0] in {"postgresql", "postgres"}:
|
||||
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
|
||||
|
||||
connect_args = build_asyncpg_connect_args(pg_schema)
|
||||
|
||||
connectable = create_async_engine(url, connect_args=connect_args)
|
||||
|
||||
# Cross-process bootstrap safety for SQLite: every connection alembic
|
||||
# opens needs a wide ``busy_timeout`` so that when another process holds
|
||||
|
||||
257
backend/packages/harness/deerflow/persistence/postgres_schema.py
Normal file
257
backend/packages/harness/deerflow/persistence/postgres_schema.py
Normal file
@ -0,0 +1,257 @@
|
||||
"""PostgreSQL schema helpers (Issue #3380).
|
||||
|
||||
Centralizes the driver-specific ways of pinning a connection's
|
||||
``search_path`` to a target schema. The two PostgreSQL drivers DeerFlow
|
||||
uses expect different mechanisms:
|
||||
|
||||
- **asyncpg** (app ORM engine): only honours ``server_settings`` passed
|
||||
via SQLAlchemy ``connect_args``. It does not understand libpq's
|
||||
``options=-c ...`` syntax.
|
||||
- **psycopg** (LangGraph checkpointer/store): uses the libpq
|
||||
``options=-c search_path=...`` connection parameter, either as a pool
|
||||
kwarg or encoded into the DSN query string.
|
||||
|
||||
Schema names are validated upstream by
|
||||
:class:`deerflow.config.database_config.DatabaseConfig` to be plain
|
||||
identifiers. SQL-emitting helpers re-validate at the boundary as
|
||||
defense-in-depth; connection-argument helpers only assemble driver payloads.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit
|
||||
|
||||
|
||||
def build_asyncpg_connect_args(schema: str) -> dict:
|
||||
"""Return SQLAlchemy ``connect_args`` that pin asyncpg's search_path.
|
||||
|
||||
Empty *schema* yields ``{}`` so the engine keeps the server default.
|
||||
"""
|
||||
if not schema:
|
||||
return {}
|
||||
return {"server_settings": {"search_path": schema}}
|
||||
|
||||
|
||||
def build_psycopg_options(schema: str) -> str | None:
|
||||
"""Return the libpq ``options`` value for psycopg pool kwargs.
|
||||
|
||||
Empty *schema* yields ``None`` so callers can skip setting the kwarg.
|
||||
"""
|
||||
if not schema:
|
||||
return None
|
||||
return f"-c search_path={schema}"
|
||||
|
||||
|
||||
def _split_libpq_options(options: str) -> list[str]:
|
||||
"""Tokenize a libpq ``options`` string.
|
||||
|
||||
libpq splits on unescaped whitespace; a backslash escapes the next
|
||||
character (so ``\\ `` is a literal space and ``\\\\`` a literal backslash).
|
||||
This is NOT POSIX shell quoting -- single/double quotes are literal here.
|
||||
"""
|
||||
tokens: list[str] = []
|
||||
current: list[str] = []
|
||||
in_token = False
|
||||
escaped = False
|
||||
for char in options:
|
||||
if escaped:
|
||||
current.append(char)
|
||||
escaped = False
|
||||
in_token = True
|
||||
continue
|
||||
if char == "\\":
|
||||
escaped = True
|
||||
in_token = True
|
||||
continue
|
||||
if char.isspace():
|
||||
if in_token:
|
||||
tokens.append("".join(current))
|
||||
current = []
|
||||
in_token = False
|
||||
continue
|
||||
current.append(char)
|
||||
in_token = True
|
||||
if in_token:
|
||||
tokens.append("".join(current))
|
||||
return tokens
|
||||
|
||||
|
||||
def _join_libpq_options(tokens: list[str]) -> str:
|
||||
"""Join tokens into a libpq ``options`` string.
|
||||
|
||||
Whitespace and backslashes inside a token are backslash-escaped so libpq
|
||||
keeps each token intact. ``shlex.join`` cannot be used: it emits POSIX
|
||||
shell quoting (single quotes), which libpq treats as literal characters.
|
||||
|
||||
All whitespace bytes are escaped, not just spaces: ``_split_libpq_options``
|
||||
preserves a backslash-escaped TAB/CR/LF as part of one token, so re-joining
|
||||
with a bare whitespace byte would let libpq re-tokenize on it and corrupt a
|
||||
caller's pre-existing ``options`` value.
|
||||
"""
|
||||
escaped = [re.sub(r"([\\\s])", r"\\\1", token) for token in tokens]
|
||||
return " ".join(escaped)
|
||||
|
||||
|
||||
def _merge_search_path_option(existing_options: str, schema: str) -> str:
|
||||
"""Return libpq options with search_path replaced while preserving others."""
|
||||
new_option = build_psycopg_options(schema)
|
||||
if not new_option:
|
||||
return existing_options
|
||||
|
||||
if not existing_options:
|
||||
return new_option
|
||||
|
||||
tokens = _split_libpq_options(existing_options)
|
||||
|
||||
merged: list[str] = []
|
||||
index = 0
|
||||
while index < len(tokens):
|
||||
token = tokens[index]
|
||||
if token == "-c" and index + 1 < len(tokens):
|
||||
setting = tokens[index + 1]
|
||||
if setting.split("=", 1)[0] == "search_path":
|
||||
index += 2
|
||||
continue
|
||||
merged.extend([token, setting])
|
||||
index += 2
|
||||
continue
|
||||
if token.startswith("-csearch_path="):
|
||||
index += 1
|
||||
continue
|
||||
merged.append(token)
|
||||
index += 1
|
||||
|
||||
merged.extend(_split_libpq_options(new_option))
|
||||
return _join_libpq_options(merged)
|
||||
|
||||
|
||||
def create_schema_sql(schema: str) -> str | None:
|
||||
"""Return a safe CREATE SCHEMA statement for a validated plain identifier.
|
||||
|
||||
Defense-in-depth: the identifier is re-validated here rather than trusting
|
||||
the distant pydantic validator. ``create_schema_sql`` is publicly exported
|
||||
and psycopg accepts multiple ``;``-separated statements, so a future caller
|
||||
that bypasses ``DatabaseConfig``/``CheckpointerConfig`` (e.g. a test helper)
|
||||
must not be able to inject SQL through this f-string boundary.
|
||||
"""
|
||||
if not schema:
|
||||
return None
|
||||
from deerflow.config.postgres_schema import validate_postgres_schema
|
||||
|
||||
validate_postgres_schema(schema)
|
||||
return f'CREATE SCHEMA IF NOT EXISTS "{schema}"'
|
||||
|
||||
|
||||
def normalize_libpq_dsn(dsn: str) -> str:
|
||||
"""Return *dsn* with any SQLAlchemy ``+driver`` suffix dropped.
|
||||
|
||||
``DatabaseConfig.postgres_url`` may carry a SQLAlchemy driver suffix such
|
||||
as ``postgresql+asyncpg://``. psycopg's libpq only understands the bare
|
||||
``postgres``/``postgresql`` scheme, so a raw ``+asyncpg`` DSN handed to
|
||||
``psycopg.connect`` raises an opaque parse error. Keyword/DSN strings
|
||||
without a URL scheme (``host=... dbname=...``) are returned unchanged.
|
||||
|
||||
Raises ``ValueError`` for URL schemes that are not a PostgreSQL variant.
|
||||
"""
|
||||
parts = urlsplit(dsn)
|
||||
if not parts.scheme:
|
||||
return dsn
|
||||
scheme_base = parts.scheme.split("+", 1)[0]
|
||||
if scheme_base not in {"postgres", "postgresql"}:
|
||||
raise ValueError(f"Unsupported PostgreSQL DSN scheme for schema injection: {parts.scheme!r}")
|
||||
if scheme_base == parts.scheme:
|
||||
return dsn
|
||||
return urlunsplit((scheme_base, parts.netloc, parts.path, parts.query, parts.fragment))
|
||||
|
||||
|
||||
def dsn_with_search_path(dsn: str, schema: str) -> str:
|
||||
"""Return *dsn* with an ``options=-c search_path=<schema>`` query param.
|
||||
|
||||
Used for psycopg ``from_conn_string`` call sites that take a DSN
|
||||
string rather than pool kwargs. The value contains a space and ``=``;
|
||||
both are percent-encoded so libpq parses the URL correctly.
|
||||
|
||||
libpq only recognizes ``%XX`` percent-encoding in URI query values; it
|
||||
does NOT treat ``+`` as a space (that is an HTML-form convention). So
|
||||
the space MUST be encoded as ``%20`` rather than ``+``, otherwise libpq
|
||||
sees a single broken token ``-c+search_path=...`` and the search_path is
|
||||
never applied. Existing query parameters are preserved. Empty *schema*
|
||||
returns *dsn* unchanged.
|
||||
"""
|
||||
if not schema:
|
||||
return dsn
|
||||
parts = urlsplit(dsn)
|
||||
|
||||
if not parts.scheme:
|
||||
from psycopg.conninfo import conninfo_to_dict, make_conninfo
|
||||
|
||||
params = conninfo_to_dict(dsn)
|
||||
params["options"] = _merge_search_path_option(params.get("options", ""), schema)
|
||||
return make_conninfo(**params)
|
||||
|
||||
# DatabaseConfig.postgres_url may carry a SQLAlchemy driver suffix such as
|
||||
# ``postgresql+asyncpg://``. psycopg's libpq only understands the bare
|
||||
# ``postgres``/``postgresql`` scheme, so accept the compound form but emit
|
||||
# a psycopg-consumable DSN by dropping the ``+driver`` part.
|
||||
scheme_base = parts.scheme.split("+", 1)[0]
|
||||
if scheme_base not in {"postgres", "postgresql"}:
|
||||
raise ValueError(f"Unsupported PostgreSQL DSN scheme for schema injection: {parts.scheme!r}")
|
||||
|
||||
options_values: list[str] = []
|
||||
query_pairs = []
|
||||
for key, value in parse_qsl(parts.query, keep_blank_values=True):
|
||||
if key == "options":
|
||||
options_values.append(value)
|
||||
else:
|
||||
query_pairs.append((key, value))
|
||||
|
||||
options = _merge_search_path_option(" ".join(options_values), schema)
|
||||
query_pairs.append(("options", options))
|
||||
# quote_via=quote encodes space as %20 (libpq-safe), not + (form-style).
|
||||
query = urlencode(query_pairs, quote_via=quote)
|
||||
return urlunsplit((scheme_base, parts.netloc, parts.path, query, parts.fragment))
|
||||
|
||||
|
||||
def ensure_postgres_schema(conn_string: str, schema: str, *, install_hint: str) -> None:
|
||||
"""Create *schema* over a fresh sync psycopg connection.
|
||||
|
||||
No-op when *schema* is empty. A missing ``psycopg`` dependency is mapped to
|
||||
*install_hint* so callers surface the same actionable message they use for
|
||||
the rest of the backend. The DSN is normalized so a SQLAlchemy ``+driver``
|
||||
suffix does not reach libpq.
|
||||
"""
|
||||
statement = create_schema_sql(schema)
|
||||
if statement is None:
|
||||
return
|
||||
try:
|
||||
import psycopg
|
||||
except ImportError as exc:
|
||||
raise ImportError(install_hint) from exc
|
||||
|
||||
# psycopg 3's ``Connection.__exit__`` only commits/rolls back -- it does NOT
|
||||
# close the connection (a documented psycopg2->3 change). Use try/finally so
|
||||
# the libpq connection is released deterministically, mirroring the async
|
||||
# counterpart, instead of leaking it until GC.
|
||||
conn = psycopg.connect(normalize_libpq_dsn(conn_string), autocommit=True)
|
||||
try:
|
||||
conn.execute(statement)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
async def ensure_postgres_schema_async(conn_string: str, schema: str, *, install_hint: str) -> None:
|
||||
"""Async counterpart of :func:`ensure_postgres_schema`."""
|
||||
statement = create_schema_sql(schema)
|
||||
if statement is None:
|
||||
return
|
||||
try:
|
||||
import psycopg
|
||||
except ImportError as exc:
|
||||
raise ImportError(install_hint) from exc
|
||||
|
||||
conn = await psycopg.AsyncConnection.connect(normalize_libpq_dsn(conn_string), autocommit=True)
|
||||
try:
|
||||
await conn.execute(statement)
|
||||
finally:
|
||||
await conn.close()
|
||||
@ -25,6 +25,7 @@ from collections.abc import AsyncIterator
|
||||
from langgraph.types import Checkpointer
|
||||
|
||||
from deerflow.config.app_config import AppConfig, get_app_config
|
||||
from deerflow.persistence.postgres_schema import create_schema_sql, dsn_with_search_path, normalize_libpq_dsn
|
||||
from deerflow.runtime.checkpointer.provider import (
|
||||
POSTGRES_CONN_REQUIRED,
|
||||
POSTGRES_INSTALL,
|
||||
@ -47,26 +48,43 @@ def _prepare_database_sqlite_checkpointer_path(db_config) -> str:
|
||||
return conn_str
|
||||
|
||||
|
||||
def _build_postgres_pool(conn_string: str):
|
||||
def _build_postgres_pool(conn_string: str, schema: str = ""):
|
||||
"""Build an AsyncConnectionPool with TCP keepalive and connection checking."""
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
kwargs = {
|
||||
"autocommit": True,
|
||||
"prepare_threshold": 0,
|
||||
"row_factory": dict_row,
|
||||
"keepalives": 1,
|
||||
"keepalives_idle": 60,
|
||||
"keepalives_interval": 10,
|
||||
"keepalives_count": 6,
|
||||
}
|
||||
# Inject search_path into the DSN (merging with any libpq options already in
|
||||
# the conn string) rather than via kwargs["options"], which psycopg applies
|
||||
# *on top of* the conninfo and would silently drop a DSN-supplied option
|
||||
# such as statement_timeout. This also strips a SQLAlchemy ``+driver``
|
||||
# suffix so libpq can parse the DSN. Matches the sync/DSN paths.
|
||||
dsn = dsn_with_search_path(normalize_libpq_dsn(conn_string), schema)
|
||||
|
||||
return AsyncConnectionPool(
|
||||
conn_string,
|
||||
kwargs={
|
||||
"autocommit": True,
|
||||
"prepare_threshold": 0,
|
||||
"row_factory": dict_row,
|
||||
"keepalives": 1,
|
||||
"keepalives_idle": 60,
|
||||
"keepalives_interval": 10,
|
||||
"keepalives_count": 6,
|
||||
},
|
||||
dsn,
|
||||
kwargs=kwargs,
|
||||
check=AsyncConnectionPool.check_connection,
|
||||
)
|
||||
|
||||
|
||||
async def _ensure_postgres_schema_with_pool(pool, schema: str) -> None:
|
||||
"""Create the configured schema before LangGraph creates its tables."""
|
||||
statement = create_schema_sql(schema)
|
||||
if statement is None:
|
||||
return
|
||||
async with pool.connection() as conn:
|
||||
await conn.execute(statement)
|
||||
|
||||
|
||||
def _ensure_postgres_imports():
|
||||
"""Import and return (AsyncPostgresSaver, AsyncConnectionPool), raising ImportError on failure."""
|
||||
try:
|
||||
@ -113,8 +131,9 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
|
||||
raise ValueError(POSTGRES_CONN_REQUIRED)
|
||||
|
||||
AsyncPostgresSaver, _ = _ensure_postgres_imports()
|
||||
pool = _build_postgres_pool(config.connection_string)
|
||||
pool = _build_postgres_pool(config.connection_string, config.postgres_schema)
|
||||
async with pool:
|
||||
await _ensure_postgres_schema_with_pool(pool, config.postgres_schema)
|
||||
saver = AsyncPostgresSaver(conn=pool)
|
||||
await saver.setup()
|
||||
yield saver
|
||||
@ -154,8 +173,9 @@ async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpoi
|
||||
raise ValueError("database.postgres_url is required for the postgres backend")
|
||||
|
||||
AsyncPostgresSaver, _ = _ensure_postgres_imports()
|
||||
pool = _build_postgres_pool(db_config.postgres_url)
|
||||
pool = _build_postgres_pool(db_config.postgres_url, db_config.postgres_schema)
|
||||
async with pool:
|
||||
await _ensure_postgres_schema_with_pool(pool, db_config.postgres_schema)
|
||||
saver = AsyncPostgresSaver(conn=pool)
|
||||
await saver.setup()
|
||||
yield saver
|
||||
|
||||
@ -28,6 +28,7 @@ from langgraph.types import Checkpointer
|
||||
|
||||
from deerflow.config.app_config import AppConfig, get_app_config
|
||||
from deerflow.config.checkpointer_config import CheckpointerConfig, ensure_config_loaded, get_checkpointer_config
|
||||
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema
|
||||
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -43,6 +44,11 @@ POSTGRES_INSTALL = (
|
||||
POSTGRES_CONN_REQUIRED = "checkpointer.connection_string is required for the postgres backend"
|
||||
|
||||
|
||||
def _ensure_postgres_schema(conn_string: str, schema: str) -> None:
|
||||
"""Create the configured schema before LangGraph creates its tables."""
|
||||
ensure_postgres_schema(conn_string, schema, install_hint=POSTGRES_INSTALL)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -68,7 +74,7 @@ def _resolve_checkpointer_config(app_config: AppConfig) -> CheckpointerConfig:
|
||||
if database.backend == "postgres":
|
||||
if not database.postgres_url:
|
||||
raise ValueError("database.postgres_url is required for the postgres backend")
|
||||
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url)
|
||||
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url, postgres_schema=database.postgres_schema)
|
||||
raise ValueError(f"Unknown database backend: {database.backend!r}")
|
||||
|
||||
|
||||
@ -131,7 +137,9 @@ def _sync_checkpointer_cm(config: CheckpointerConfig) -> Iterator[Checkpointer]:
|
||||
if not config.connection_string:
|
||||
raise ValueError(POSTGRES_CONN_REQUIRED)
|
||||
|
||||
with PostgresSaver.from_conn_string(config.connection_string) as saver:
|
||||
_ensure_postgres_schema(config.connection_string, config.postgres_schema)
|
||||
conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
|
||||
with PostgresSaver.from_conn_string(conn_string) as saver:
|
||||
saver.setup()
|
||||
logger.info("Checkpointer: using PostgresSaver")
|
||||
yield saver
|
||||
|
||||
@ -25,6 +25,7 @@ from collections.abc import AsyncIterator
|
||||
from langgraph.store.base import BaseStore
|
||||
|
||||
from deerflow.config.app_config import AppConfig, get_app_config
|
||||
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema_async
|
||||
from deerflow.runtime.store.provider import (
|
||||
POSTGRES_CONN_REQUIRED,
|
||||
POSTGRES_STORE_INSTALL,
|
||||
@ -36,6 +37,12 @@ from deerflow.runtime.store.provider import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _ensure_postgres_schema(conn_string: str, schema: str) -> None:
|
||||
"""Create the configured schema before LangGraph creates its store tables."""
|
||||
await ensure_postgres_schema_async(conn_string, schema, install_hint=POSTGRES_STORE_INSTALL)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal backend factory
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -79,7 +86,9 @@ async def _async_store(config) -> AsyncIterator[BaseStore]:
|
||||
if not config.connection_string:
|
||||
raise ValueError(POSTGRES_CONN_REQUIRED)
|
||||
|
||||
async with AsyncPostgresStore.from_conn_string(config.connection_string) as store:
|
||||
await _ensure_postgres_schema(config.connection_string, config.postgres_schema)
|
||||
conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
|
||||
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
|
||||
await store.setup()
|
||||
logger.info("Store: using AsyncPostgresStore")
|
||||
yield store
|
||||
|
||||
@ -30,6 +30,7 @@ from langgraph.store.base import BaseStore
|
||||
|
||||
from deerflow.config.app_config import AppConfig, get_app_config
|
||||
from deerflow.config.checkpointer_config import CheckpointerConfig, ensure_config_loaded, get_checkpointer_config
|
||||
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema
|
||||
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -45,12 +46,19 @@ POSTGRES_STORE_INSTALL = (
|
||||
POSTGRES_CONN_REQUIRED = "checkpointer.connection_string is required for the postgres backend"
|
||||
|
||||
|
||||
def _ensure_postgres_schema(conn_string: str, schema: str) -> None:
|
||||
"""Create the configured schema before LangGraph creates its store tables."""
|
||||
ensure_postgres_schema(conn_string, schema, install_hint=POSTGRES_STORE_INSTALL)
|
||||
|
||||
|
||||
def _resolve_store_config(app_config: AppConfig) -> CheckpointerConfig:
|
||||
"""Resolve the Store backend from legacy or unified application config.
|
||||
|
||||
The legacy ``checkpointer`` section remains authoritative when present so
|
||||
Store and Checkpointer continue to use the same backend. Otherwise the
|
||||
unified ``database`` section drives the Store as documented.
|
||||
unified ``database`` section drives the Store as documented. The unified
|
||||
``postgres_schema`` is forwarded so Store tables land in the configured
|
||||
schema alongside the checkpointer and app tables.
|
||||
"""
|
||||
if app_config.checkpointer is not None:
|
||||
return app_config.checkpointer
|
||||
@ -63,7 +71,7 @@ def _resolve_store_config(app_config: AppConfig) -> CheckpointerConfig:
|
||||
if database.backend == "postgres":
|
||||
if not database.postgres_url:
|
||||
raise ValueError("database.postgres_url is required for the postgres backend")
|
||||
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url)
|
||||
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url, postgres_schema=database.postgres_schema)
|
||||
raise ValueError(f"Unknown database backend: {database.backend!r}")
|
||||
|
||||
|
||||
@ -126,7 +134,9 @@ def _sync_store_cm(config) -> Iterator[BaseStore]:
|
||||
if not config.connection_string:
|
||||
raise ValueError(POSTGRES_CONN_REQUIRED)
|
||||
|
||||
with PostgresStore.from_conn_string(config.connection_string) as store:
|
||||
_ensure_postgres_schema(config.connection_string, config.postgres_schema)
|
||||
conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
|
||||
with PostgresStore.from_conn_string(conn_string) as store:
|
||||
store.setup()
|
||||
logger.info("Store: using PostgresStore")
|
||||
yield store
|
||||
|
||||
@ -646,6 +646,56 @@ def test_get_app_config_keeps_persistence_runtime_singletons_when_checkpointer_u
|
||||
_reset_config_singletons()
|
||||
|
||||
|
||||
def test_get_app_config_does_not_reset_persistence_singletons_when_database_changes(tmp_path, monkeypatch):
|
||||
# ``database`` is a restart-required field: the ORM engine is built once at
|
||||
# startup and never rebuilt on a config.yaml edit. Resetting only the sync
|
||||
# checkpointer/store singletons on a live ``postgres_schema`` change would
|
||||
# half-migrate the deployment (new checkpoint/store tables in the new schema,
|
||||
# ORM rows still in the old one). So a ``database`` change must NOT trigger a
|
||||
# partial reset -- the operator must restart.
|
||||
config_path = tmp_path / "config.yaml"
|
||||
extensions_path = tmp_path / "extensions_config.json"
|
||||
_write_extensions_config(extensions_path)
|
||||
_write_config_with_sections(
|
||||
config_path,
|
||||
{"database": {"backend": "postgres", "postgres_url": "postgresql://localhost/db", "postgres_schema": "schema_a"}},
|
||||
)
|
||||
|
||||
monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(config_path))
|
||||
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(extensions_path))
|
||||
_reset_config_singletons()
|
||||
|
||||
reset_calls = {"checkpointer": 0, "store": 0}
|
||||
|
||||
def _reset_checkpointer() -> None:
|
||||
reset_calls["checkpointer"] += 1
|
||||
|
||||
def _reset_store() -> None:
|
||||
reset_calls["store"] += 1
|
||||
|
||||
monkeypatch.setattr("deerflow.runtime.checkpointer.reset_checkpointer", _reset_checkpointer)
|
||||
monkeypatch.setattr("deerflow.runtime.store.reset_store", _reset_store)
|
||||
|
||||
try:
|
||||
get_app_config()
|
||||
reset_calls["checkpointer"] = 0
|
||||
reset_calls["store"] = 0
|
||||
|
||||
_write_config_with_sections(
|
||||
config_path,
|
||||
{"database": {"backend": "postgres", "postgres_url": "postgresql://localhost/db", "postgres_schema": "schema_b"}},
|
||||
)
|
||||
next_mtime = config_path.stat().st_mtime + 5
|
||||
os.utime(config_path, (next_mtime, next_mtime))
|
||||
|
||||
get_app_config()
|
||||
|
||||
assert get_checkpointer_config() is None
|
||||
assert reset_calls == {"checkpointer": 0, "store": 0}
|
||||
finally:
|
||||
_reset_config_singletons()
|
||||
|
||||
|
||||
def test_get_app_config_does_not_mutate_singletons_when_reload_validation_fails(tmp_path, monkeypatch):
|
||||
config_path = tmp_path / "config.yaml"
|
||||
extensions_path = tmp_path / "extensions_config.json"
|
||||
|
||||
@ -144,6 +144,18 @@ class TestCheckpointerConfig:
|
||||
assert config is not None
|
||||
assert config.type == "postgres"
|
||||
assert config.connection_string == "postgresql://localhost/db"
|
||||
assert config.postgres_schema == ""
|
||||
|
||||
def test_postgres_schema_accepts_valid_identifier(self):
|
||||
config = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema="deerflow")
|
||||
assert config.postgres_schema == "deerflow"
|
||||
|
||||
@pytest.mark.parametrize("schema", ["1abc", "a b", "a;b", "a-b", "a" * 64, 'a"b', "MySchema", "Orders", "Public"])
|
||||
def test_postgres_schema_rejects_invalid_identifier(self, schema):
|
||||
from pydantic import ValidationError
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema=schema)
|
||||
|
||||
def test_default_connection_string_is_none(self):
|
||||
config = CheckpointerConfig(type="memory")
|
||||
@ -393,6 +405,76 @@ class TestGetCheckpointer:
|
||||
mock_saver_cls.from_conn_string.assert_called_once_with("postgresql://localhost/db")
|
||||
mock_saver_instance.setup.assert_called_once()
|
||||
|
||||
def test_postgres_schema_creates_schema_and_sets_search_path(self):
|
||||
"""Sync Postgres checkpointer should create schema before setup."""
|
||||
load_checkpointer_config_from_dict({"type": "postgres", "connection_string": "postgresql://localhost/db", "postgres_schema": "deerflow"})
|
||||
|
||||
mock_saver_instance = MagicMock()
|
||||
mock_cm = MagicMock()
|
||||
mock_cm.__enter__ = MagicMock(return_value=mock_saver_instance)
|
||||
mock_cm.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
mock_saver_cls = MagicMock()
|
||||
mock_saver_cls.from_conn_string = MagicMock(return_value=mock_cm)
|
||||
|
||||
mock_pg_module = MagicMock()
|
||||
mock_pg_module.PostgresSaver = mock_saver_cls
|
||||
|
||||
mock_conn = MagicMock()
|
||||
mock_psycopg = MagicMock()
|
||||
mock_psycopg.connect.return_value = mock_conn
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"langgraph.checkpoint.postgres": mock_pg_module}),
|
||||
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
|
||||
):
|
||||
reset_checkpointer()
|
||||
cp = get_checkpointer()
|
||||
|
||||
assert cp is mock_saver_instance
|
||||
mock_psycopg.connect.assert_called_once_with("postgresql://localhost/db", autocommit=True)
|
||||
mock_conn.execute.assert_called_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
|
||||
# psycopg 3 __exit__ does not close(); the sync path must close explicitly.
|
||||
mock_conn.close.assert_called_once_with()
|
||||
called_dsn = mock_saver_cls.from_conn_string.call_args.args[0]
|
||||
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
|
||||
mock_saver_instance.setup.assert_called_once()
|
||||
|
||||
def test_store_postgres_schema_creates_schema_and_sets_search_path(self):
|
||||
"""Sync Postgres store should use the legacy checkpointer schema."""
|
||||
load_checkpointer_config_from_dict({"type": "postgres", "connection_string": "postgresql://localhost/db", "postgres_schema": "deerflow"})
|
||||
|
||||
mock_store_instance = MagicMock()
|
||||
mock_cm = MagicMock()
|
||||
mock_cm.__enter__ = MagicMock(return_value=mock_store_instance)
|
||||
mock_cm.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
mock_store_cls = MagicMock()
|
||||
mock_store_cls.from_conn_string = MagicMock(return_value=mock_cm)
|
||||
|
||||
mock_pg_module = MagicMock()
|
||||
mock_pg_module.PostgresStore = mock_store_cls
|
||||
|
||||
mock_conn = MagicMock()
|
||||
mock_psycopg = MagicMock()
|
||||
mock_psycopg.connect.return_value = mock_conn
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"langgraph.store.postgres": mock_pg_module}),
|
||||
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
|
||||
):
|
||||
reset_store()
|
||||
store = get_store()
|
||||
|
||||
assert store is mock_store_instance
|
||||
mock_psycopg.connect.assert_called_once_with("postgresql://localhost/db", autocommit=True)
|
||||
mock_conn.execute.assert_called_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
|
||||
# psycopg 3 __exit__ does not close(); the sync path must close explicitly.
|
||||
mock_conn.close.assert_called_once_with()
|
||||
called_dsn = mock_store_cls.from_conn_string.call_args.args[0]
|
||||
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
|
||||
mock_store_instance.setup.assert_called_once()
|
||||
|
||||
|
||||
class TestSyncSingletonThreadSafety:
|
||||
def test_store_reset_clears_singleton(self):
|
||||
@ -576,7 +658,7 @@ class TestAsyncCheckpointer:
|
||||
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.checkpointer = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db")
|
||||
mock_config.checkpointer = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema="deerflow")
|
||||
|
||||
mock_saver = AsyncMock()
|
||||
|
||||
@ -585,6 +667,11 @@ class TestAsyncCheckpointer:
|
||||
mock_pool_instance = AsyncMock()
|
||||
mock_pool_instance.__aenter__.return_value = mock_pool_instance
|
||||
mock_pool_instance.__aexit__.return_value = False
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn_cm = AsyncMock()
|
||||
mock_conn_cm.__aenter__.return_value = mock_conn
|
||||
mock_conn_cm.__aexit__.return_value = False
|
||||
mock_pool_instance.connection = MagicMock(return_value=mock_conn_cm)
|
||||
|
||||
mock_pool_cls = MagicMock(return_value=mock_pool_instance)
|
||||
mock_pool_cls.check_connection = AsyncMock()
|
||||
@ -610,8 +697,12 @@ class TestAsyncCheckpointer:
|
||||
# Verify the pool was constructed with check Connection
|
||||
mock_pool_cls.assert_called_once()
|
||||
call_kwargs = mock_pool_cls.call_args
|
||||
assert call_kwargs[0][0] == "postgresql://localhost/db"
|
||||
# search_path is injected into the DSN (merged with any existing libpq
|
||||
# options), not via kwargs["options"] which would clobber DSN options.
|
||||
assert "options=-c%20search_path%3Ddeerflow" in call_kwargs[0][0]
|
||||
assert call_kwargs[1]["check"] is mock_pool_cls.check_connection
|
||||
assert "options" not in call_kwargs[1]["kwargs"]
|
||||
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
|
||||
|
||||
# Verify saver was constructed with the pool (not via from_conn_string)
|
||||
mock_saver_cls.assert_called_once_with(conn=mock_pool_instance)
|
||||
@ -623,7 +714,7 @@ class TestAsyncCheckpointer:
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
|
||||
|
||||
db_config = DatabaseConfig(backend="postgres", postgres_url="postgresql://localhost/db")
|
||||
db_config = DatabaseConfig(backend="postgres", postgres_url="postgresql://localhost/db", postgres_schema="deerflow")
|
||||
mock_config = MagicMock()
|
||||
mock_config.checkpointer = None
|
||||
mock_config.database = db_config
|
||||
@ -635,6 +726,11 @@ class TestAsyncCheckpointer:
|
||||
mock_pool_instance = AsyncMock()
|
||||
mock_pool_instance.__aenter__.return_value = mock_pool_instance
|
||||
mock_pool_instance.__aexit__.return_value = False
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn_cm = AsyncMock()
|
||||
mock_conn_cm.__aenter__.return_value = mock_conn
|
||||
mock_conn_cm.__aexit__.return_value = False
|
||||
mock_pool_instance.connection = MagicMock(return_value=mock_conn_cm)
|
||||
|
||||
mock_pool_cls = MagicMock(return_value=mock_pool_instance)
|
||||
mock_pool_cls.check_connection = AsyncMock()
|
||||
@ -657,8 +753,10 @@ class TestAsyncCheckpointer:
|
||||
|
||||
mock_pool_cls.assert_called_once()
|
||||
call_kwargs = mock_pool_cls.call_args
|
||||
assert call_kwargs[0][0] == "postgresql://localhost/db"
|
||||
assert "options=-c%20search_path%3Ddeerflow" in call_kwargs[0][0]
|
||||
assert call_kwargs[1]["check"] is mock_pool_cls.check_connection
|
||||
assert "options" not in call_kwargs[1]["kwargs"]
|
||||
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
|
||||
|
||||
mock_saver_cls.assert_called_once_with(conn=mock_pool_instance)
|
||||
mock_saver.setup.assert_awaited_once()
|
||||
@ -705,6 +803,88 @@ class TestAsyncCheckpointer:
|
||||
mock_saver.setup.assert_awaited_once()
|
||||
|
||||
|
||||
class TestAsyncStore:
|
||||
@pytest.mark.anyio
|
||||
async def test_postgres_schema_creates_schema_and_sets_search_path(self):
|
||||
"""Async Postgres store should use the legacy checkpointer schema."""
|
||||
from deerflow.runtime.store.async_provider import make_store
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.checkpointer = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema="deerflow")
|
||||
|
||||
mock_store = AsyncMock()
|
||||
mock_cm = AsyncMock()
|
||||
mock_cm.__aenter__.return_value = mock_store
|
||||
mock_cm.__aexit__.return_value = False
|
||||
|
||||
mock_store_cls = MagicMock()
|
||||
mock_store_cls.from_conn_string.return_value = mock_cm
|
||||
|
||||
mock_pg_module = MagicMock()
|
||||
mock_pg_module.AsyncPostgresStore = mock_store_cls
|
||||
|
||||
mock_conn = AsyncMock()
|
||||
mock_async_connection = MagicMock()
|
||||
mock_async_connection.connect = AsyncMock(return_value=mock_conn)
|
||||
mock_psycopg = MagicMock()
|
||||
mock_psycopg.AsyncConnection = mock_async_connection
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"langgraph.store.postgres.aio": mock_pg_module}),
|
||||
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
|
||||
):
|
||||
async with make_store(mock_config) as store:
|
||||
assert store is mock_store
|
||||
|
||||
mock_async_connection.connect.assert_awaited_once_with("postgresql://localhost/db", autocommit=True)
|
||||
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
|
||||
mock_conn.close.assert_awaited_once()
|
||||
called_dsn = mock_store_cls.from_conn_string.call_args.args[0]
|
||||
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
|
||||
mock_store.setup.assert_awaited_once()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_database_postgres_schema_creates_schema_and_sets_search_path(self):
|
||||
"""Unified database postgres store should use database.postgres_schema."""
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.runtime.store.async_provider import make_store
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.checkpointer = None
|
||||
mock_config.database = DatabaseConfig(backend="postgres", postgres_url="postgresql://localhost/db", postgres_schema="deerflow")
|
||||
|
||||
mock_store = AsyncMock()
|
||||
mock_cm = AsyncMock()
|
||||
mock_cm.__aenter__.return_value = mock_store
|
||||
mock_cm.__aexit__.return_value = False
|
||||
|
||||
mock_store_cls = MagicMock()
|
||||
mock_store_cls.from_conn_string.return_value = mock_cm
|
||||
|
||||
mock_pg_module = MagicMock()
|
||||
mock_pg_module.AsyncPostgresStore = mock_store_cls
|
||||
|
||||
mock_conn = AsyncMock()
|
||||
mock_async_connection = MagicMock()
|
||||
mock_async_connection.connect = AsyncMock(return_value=mock_conn)
|
||||
mock_psycopg = MagicMock()
|
||||
mock_psycopg.AsyncConnection = mock_async_connection
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"langgraph.store.postgres.aio": mock_pg_module}),
|
||||
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
|
||||
):
|
||||
async with make_store(mock_config) as store:
|
||||
assert store is mock_store
|
||||
|
||||
mock_async_connection.connect.assert_awaited_once_with("postgresql://localhost/db", autocommit=True)
|
||||
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
|
||||
mock_conn.close.assert_awaited_once()
|
||||
called_dsn = mock_store_cls.from_conn_string.call_args.args[0]
|
||||
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
|
||||
mock_store.setup.assert_awaited_once()
|
||||
|
||||
|
||||
class TestCheckpointerDatabaseConfig:
|
||||
"""The sync checkpointer must follow the unified ``database`` section when no
|
||||
legacy ``checkpointer`` section is configured — matching the async
|
||||
|
||||
@ -69,3 +69,34 @@ def test_escape_url_for_alembic_doubles_only_percent_signs() -> None:
|
||||
# Idempotency is intentionally NOT a property -- doubling is one-way;
|
||||
# callers must escape exactly once on the way into set_main_option.
|
||||
assert _escape_url_for_alembic("a%%b") == "a%%%%b"
|
||||
|
||||
|
||||
def test_alembic_config_forwards_postgres_schema_option() -> None:
|
||||
# The custom schema must reach env.py so the alembic-spawned engine can
|
||||
# pin its search_path; otherwise alembic_version + migration DDL land in
|
||||
# ``public`` while the ORM tables land in the custom schema.
|
||||
engine = _fake_engine("postgresql://a:b@h/d")
|
||||
cfg = _get_alembic_config(engine, postgres_schema="deerflow")
|
||||
assert cfg.get_main_option("deerflow_pg_schema") == "deerflow"
|
||||
|
||||
|
||||
def test_alembic_config_omits_schema_option_when_unset() -> None:
|
||||
engine = _fake_engine("postgresql://a:b@h/d")
|
||||
cfg = _get_alembic_config(engine)
|
||||
assert cfg.get_main_option("deerflow_pg_schema") is None
|
||||
|
||||
|
||||
def test_env_module_pins_search_path_from_schema_option() -> None:
|
||||
"""env.py must read ``deerflow_pg_schema`` and pin the alembic-spawned
|
||||
engine's search_path. That engine is built from the bare URL and does not
|
||||
inherit the app engine's asyncpg ``server_settings``, so without this both
|
||||
``alembic_version`` and migration DDL would land in ``public`` while the
|
||||
ORM tables land in the custom schema. Source-parity check mirroring
|
||||
``test_env_module_wires_busy_timeout_for_sqlite``.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
env_path = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/persistence/migrations/env.py"
|
||||
src = env_path.read_text(encoding="utf-8")
|
||||
assert 'get_main_option("deerflow_pg_schema")' in src, "env.py must read the deerflow_pg_schema option set by _get_alembic_config"
|
||||
assert "build_asyncpg_connect_args" in src, "env.py must pin the alembic engine's search_path via build_asyncpg_connect_args"
|
||||
|
||||
@ -145,7 +145,7 @@ async def test_init_engine_postgres_uses_hardened_kwargs() -> None:
|
||||
await engine_mod.init_engine(backend="postgres", url=url, echo=True, pool_size=12)
|
||||
|
||||
create_engine.assert_called_once_with(url, **engine_mod._postgres_engine_kwargs(echo=True, pool_size=12))
|
||||
bootstrap_schema.assert_awaited_once_with(mock_engine, backend="postgres")
|
||||
bootstrap_schema.assert_awaited_once_with(mock_engine, backend="postgres", postgres_schema="")
|
||||
finally:
|
||||
await engine_mod.close_engine()
|
||||
|
||||
@ -174,7 +174,10 @@ async def test_init_engine_postgres_retry_uses_hardened_kwargs() -> None:
|
||||
assert create_engine.call_args_list == [call(url, **kwargs), call(url, **kwargs)]
|
||||
auto_create.assert_awaited_once_with(url)
|
||||
initial_engine.dispose.assert_awaited_once()
|
||||
assert bootstrap_schema.await_args_list == [call(initial_engine, backend="postgres"), call(retry_engine, backend="postgres")]
|
||||
assert bootstrap_schema.await_args_list == [
|
||||
call(initial_engine, backend="postgres", postgres_schema=""),
|
||||
call(retry_engine, backend="postgres", postgres_schema=""),
|
||||
]
|
||||
finally:
|
||||
await engine_mod.close_engine()
|
||||
|
||||
|
||||
@ -10,7 +10,7 @@ Tests:
|
||||
|
||||
import sys
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@ -71,6 +71,50 @@ class TestDatabaseConfig:
|
||||
with pytest.raises(ValueError, match="No SQLAlchemy URL"):
|
||||
_ = c.app_sqlalchemy_url
|
||||
|
||||
def test_postgres_schema_default_empty(self):
|
||||
c = DatabaseConfig()
|
||||
assert c.postgres_schema == ""
|
||||
|
||||
@pytest.mark.parametrize("schema", ["deerflow", "my_schema", "_private", "s", "a" * 63])
|
||||
def test_postgres_schema_accepts_valid_identifier(self, schema):
|
||||
c = DatabaseConfig(backend="postgres", postgres_url="postgresql://u:p@h:5432/db", postgres_schema=schema)
|
||||
assert c.postgres_schema == schema
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema",
|
||||
[
|
||||
"1abc",
|
||||
"a b",
|
||||
"a;b",
|
||||
"a-b",
|
||||
"a" * 64,
|
||||
'a"b',
|
||||
"MySchema",
|
||||
"Orders",
|
||||
"Public",
|
||||
# Trailing/leading whitespace must be rejected: a ``$``-anchored
|
||||
# ``re.match`` accepts a single trailing ``\n``, which would create a
|
||||
# quoted schema literally named ``deerflow\n`` while the unquoted
|
||||
# search_path folds to ``deerflow`` and misses it (tables land in
|
||||
# ``public``). ``re.fullmatch`` on an unanchored pattern rejects it.
|
||||
"deerflow\n",
|
||||
"deerflow\t",
|
||||
"\ndeerflow",
|
||||
"deerflow ",
|
||||
],
|
||||
)
|
||||
def test_postgres_schema_rejects_invalid_identifier(self, schema):
|
||||
from pydantic import ValidationError
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
DatabaseConfig(backend="postgres", postgres_url="postgresql://u:p@h:5432/db", postgres_schema=schema)
|
||||
|
||||
def test_postgres_schema_does_not_pollute_url(self):
|
||||
c = DatabaseConfig(backend="postgres", postgres_url="postgresql://u:p@h:5432/db", postgres_schema="deerflow")
|
||||
url = c.app_sqlalchemy_url
|
||||
assert "deerflow" not in url.replace("/db", "")
|
||||
assert url.startswith("postgresql+asyncpg://")
|
||||
|
||||
|
||||
# -- MemoryRunStore --
|
||||
|
||||
@ -276,20 +320,22 @@ class TestBaseToDictMixin:
|
||||
name: Mapped[str] = mapped_column(String(128))
|
||||
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'test.db'}")
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
async with sf() as session:
|
||||
session.add(_Tmp(id="1", name="hello"))
|
||||
await session.commit()
|
||||
obj = await session.get(_Tmp, "1")
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
async with sf() as session:
|
||||
session.add(_Tmp(id="1", name="hello"))
|
||||
await session.commit()
|
||||
obj = await session.get(_Tmp, "1")
|
||||
|
||||
assert obj.to_dict() == {"id": "1", "name": "hello"}
|
||||
assert obj.to_dict(exclude={"name"}) == {"id": "1"}
|
||||
assert "_Tmp" in repr(obj)
|
||||
|
||||
await engine.dispose()
|
||||
assert obj.to_dict() == {"id": "1", "name": "hello"}
|
||||
assert obj.to_dict(exclude={"name"}) == {"id": "1"}
|
||||
assert "_Tmp" in repr(obj)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
Base.metadata.remove(_Tmp.__table__)
|
||||
|
||||
|
||||
# -- Engine lifecycle --
|
||||
@ -327,3 +373,111 @@ class TestEngineLifecycle:
|
||||
pytest.raises(ImportError, match="uv sync --all-packages --extra postgres"),
|
||||
):
|
||||
await init_engine("postgres", url="postgresql+asyncpg://x:x@localhost/x")
|
||||
|
||||
|
||||
def _make_fake_pg_engine():
|
||||
"""Build a fake async engine whose begin()/dispose() are awaitable mocks.
|
||||
|
||||
Tracks ordering of conn.execute (CREATE SCHEMA) vs conn.run_sync
|
||||
(create_all) through a shared parent mock's ``mock_calls``.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
calls = MagicMock()
|
||||
conn = MagicMock()
|
||||
# Both are awaited by the engine code, so they must return awaitables.
|
||||
calls.execute = AsyncMock()
|
||||
calls.run_sync = AsyncMock()
|
||||
conn.execute = calls.execute
|
||||
conn.run_sync = calls.run_sync
|
||||
|
||||
begin_cm = AsyncMock()
|
||||
begin_cm.__aenter__.return_value = conn
|
||||
begin_cm.__aexit__.return_value = False
|
||||
|
||||
engine = MagicMock()
|
||||
engine.begin = MagicMock(return_value=begin_cm)
|
||||
engine.dispose = AsyncMock()
|
||||
return engine, calls
|
||||
|
||||
|
||||
class TestPostgresSchemaInit:
|
||||
@pytest.mark.anyio
|
||||
async def test_passes_search_path_connect_args(self, monkeypatch):
|
||||
import deerflow.persistence.engine as engine_module
|
||||
|
||||
monkeypatch.setitem(sys.modules, "asyncpg", object())
|
||||
fake_engine, _calls = _make_fake_pg_engine()
|
||||
captured = {}
|
||||
|
||||
def fake_create(_url, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return fake_engine
|
||||
|
||||
monkeypatch.setattr(engine_module, "create_async_engine", fake_create)
|
||||
monkeypatch.setattr("deerflow.persistence.bootstrap.bootstrap_schema", AsyncMock())
|
||||
|
||||
await engine_module.init_engine(
|
||||
"postgres",
|
||||
url="postgresql+asyncpg://u:p@h:5432/db",
|
||||
postgres_schema="deerflow",
|
||||
)
|
||||
|
||||
assert captured["connect_args"] == {
|
||||
"command_timeout": engine_module.POSTGRES_COMMAND_TIMEOUT_SECONDS,
|
||||
"server_settings": {"search_path": "deerflow"},
|
||||
}
|
||||
await engine_module.close_engine()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_creates_schema_before_bootstrap(self, monkeypatch):
|
||||
import deerflow.persistence.engine as engine_module
|
||||
|
||||
monkeypatch.setitem(sys.modules, "asyncpg", object())
|
||||
fake_engine, calls = _make_fake_pg_engine()
|
||||
monkeypatch.setattr(engine_module, "create_async_engine", lambda url, **kw: fake_engine)
|
||||
calls.attach_mock(AsyncMock(), "bootstrap_schema")
|
||||
monkeypatch.setattr("deerflow.persistence.bootstrap.bootstrap_schema", calls.bootstrap_schema)
|
||||
|
||||
await engine_module.init_engine(
|
||||
"postgres",
|
||||
url="postgresql+asyncpg://u:p@h:5432/db",
|
||||
postgres_schema="deerflow",
|
||||
)
|
||||
|
||||
names = [c[0] for c in calls.mock_calls]
|
||||
assert "execute" in names
|
||||
assert "bootstrap_schema" in names
|
||||
# CREATE SCHEMA must run before the alembic bootstrap so the
|
||||
# subsequent create_all / migration DDL lands in the target schema.
|
||||
assert names.index("execute") < names.index("bootstrap_schema")
|
||||
# The DDL passed to execute must be a CreateSchema for the target schema.
|
||||
execute_arg = calls.execute.call_args[0][0]
|
||||
assert "deerflow" in str(execute_arg)
|
||||
await engine_module.close_engine()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_empty_schema_skips_connect_args_and_ddl(self, monkeypatch):
|
||||
import deerflow.persistence.engine as engine_module
|
||||
|
||||
monkeypatch.setitem(sys.modules, "asyncpg", object())
|
||||
fake_engine, calls = _make_fake_pg_engine()
|
||||
captured = {}
|
||||
|
||||
def fake_create(_url, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return fake_engine
|
||||
|
||||
monkeypatch.setattr(engine_module, "create_async_engine", fake_create)
|
||||
calls.attach_mock(AsyncMock(), "bootstrap_schema")
|
||||
monkeypatch.setattr("deerflow.persistence.bootstrap.bootstrap_schema", calls.bootstrap_schema)
|
||||
|
||||
await engine_module.init_engine("postgres", url="postgresql+asyncpg://u:p@h:5432/db")
|
||||
|
||||
assert captured.get("connect_args", {}) == {
|
||||
"command_timeout": engine_module.POSTGRES_COMMAND_TIMEOUT_SECONDS,
|
||||
}
|
||||
names = [c[0] for c in calls.mock_calls]
|
||||
assert "execute" not in names # no CREATE SCHEMA
|
||||
assert "bootstrap_schema" in names # bootstrap still runs
|
||||
await engine_module.close_engine()
|
||||
|
||||
112
backend/tests/test_pg_schema_integration.py
Normal file
112
backend/tests/test_pg_schema_integration.py
Normal file
@ -0,0 +1,112 @@
|
||||
"""Optional live PostgreSQL schema integration tests for issue #3380."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import close_engine, get_engine, init_engine_from_config
|
||||
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
|
||||
from deerflow.runtime.checkpointer.provider import _resolve_checkpointer_config, _sync_checkpointer_cm
|
||||
from deerflow.runtime.store.async_provider import make_store
|
||||
from deerflow.runtime.store.provider import _resolve_store_config, _sync_store_cm
|
||||
|
||||
POSTGRES_URL = os.getenv("DEERFLOW_TEST_POSTGRES_URL")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not POSTGRES_URL,
|
||||
reason="set DEERFLOW_TEST_POSTGRES_URL to run live PostgreSQL schema integration tests",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_postgres_schema_places_orm_checkpointer_and_store_tables_together():
|
||||
"""Verify a real PostgreSQL backend places all persistence tables in one schema."""
|
||||
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
|
||||
db_config = DatabaseConfig(backend="postgres", postgres_url=POSTGRES_URL or "", postgres_schema=schema)
|
||||
app_config = SimpleNamespace(checkpointer=None, database=db_config)
|
||||
|
||||
await init_engine_from_config(db_config)
|
||||
engine = get_engine()
|
||||
assert engine is not None
|
||||
|
||||
try:
|
||||
async with make_checkpointer(app_config) as checkpointer:
|
||||
assert checkpointer is not None
|
||||
async with make_store(app_config) as store:
|
||||
assert store is not None
|
||||
|
||||
async with engine.begin() as conn:
|
||||
rows = (
|
||||
await conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT table_schema, table_name
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema IN (:schema, 'public')
|
||||
ORDER BY table_schema, table_name
|
||||
"""
|
||||
),
|
||||
{"schema": schema},
|
||||
)
|
||||
).all()
|
||||
|
||||
by_schema = {(row.table_schema, row.table_name) for row in rows}
|
||||
orm_tables = {"runs", "run_events", "threads_meta", "feedback", "users"}
|
||||
assert {("public", table) for table in orm_tables}.isdisjoint(by_schema)
|
||||
assert {(schema, table) for table in orm_tables}.issubset(by_schema)
|
||||
assert any(table_schema == schema and "checkpoint" in table_name for table_schema, table_name in by_schema)
|
||||
assert any(table_schema == schema and ("store" in table_name or "migration" in table_name) for table_schema, table_name in by_schema)
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await close_engine()
|
||||
|
||||
|
||||
def test_sync_postgres_schema_places_checkpointer_and_store_tables_together():
|
||||
"""Verify the sync (psycopg) path honours search_path via the DSN encoding.
|
||||
|
||||
This exercises ``dsn_with_search_path`` against a real psycopg connection,
|
||||
guarding against regression of the ``%20`` vs ``+`` libpq encoding bug.
|
||||
"""
|
||||
import psycopg
|
||||
|
||||
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
|
||||
db_config = DatabaseConfig(
|
||||
backend="postgres",
|
||||
postgres_url=POSTGRES_URL or "",
|
||||
postgres_schema=schema,
|
||||
)
|
||||
checkpointer_config = _resolve_checkpointer_config(SimpleNamespace(checkpointer=None, database=db_config))
|
||||
store_config = _resolve_store_config(SimpleNamespace(checkpointer=None, database=db_config))
|
||||
|
||||
try:
|
||||
with _sync_checkpointer_cm(checkpointer_config) as checkpointer:
|
||||
assert checkpointer is not None
|
||||
with _sync_store_cm(store_config) as store:
|
||||
assert store is not None
|
||||
|
||||
with psycopg.connect(POSTGRES_URL or "", autocommit=True) as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT table_schema, table_name
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema IN (%s, 'public')
|
||||
ORDER BY table_schema, table_name
|
||||
""",
|
||||
(schema,),
|
||||
).fetchall()
|
||||
|
||||
by_schema = {(table_schema, table_name) for table_schema, table_name in rows}
|
||||
assert any(table_schema == schema and "checkpoint" in table_name for table_schema, table_name in by_schema)
|
||||
assert any(table_schema == schema and ("store" in table_name or "migration" in table_name) for table_schema, table_name in by_schema)
|
||||
# The DeerFlow LangGraph tables must NOT leak into public.
|
||||
assert not any(table_schema == "public" and ("checkpoint" in table_name or table_name == "store") for table_schema, table_name in by_schema)
|
||||
finally:
|
||||
with psycopg.connect(POSTGRES_URL or "", autocommit=True) as conn:
|
||||
conn.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')
|
||||
178
backend/tests/test_postgres_schema_helper.py
Normal file
178
backend/tests/test_postgres_schema_helper.py
Normal file
@ -0,0 +1,178 @@
|
||||
"""Tests for the PostgreSQL schema helpers (Issue #3380)."""
|
||||
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.persistence.postgres_schema import (
|
||||
build_asyncpg_connect_args,
|
||||
build_psycopg_options,
|
||||
create_schema_sql,
|
||||
dsn_with_search_path,
|
||||
normalize_libpq_dsn,
|
||||
)
|
||||
|
||||
|
||||
class TestBuildAsyncpgConnectArgs:
|
||||
def test_sets_search_path_for_schema(self):
|
||||
assert build_asyncpg_connect_args("deerflow") == {"server_settings": {"search_path": "deerflow"}}
|
||||
|
||||
def test_empty_schema_returns_empty_dict(self):
|
||||
assert build_asyncpg_connect_args("") == {}
|
||||
|
||||
|
||||
class TestBuildPsycopgOptions:
|
||||
def test_builds_libpq_options(self):
|
||||
assert build_psycopg_options("deerflow") == "-c search_path=deerflow"
|
||||
|
||||
def test_empty_schema_returns_none(self):
|
||||
assert build_psycopg_options("") is None
|
||||
|
||||
|
||||
class TestCreateSchemaSql:
|
||||
def test_builds_create_schema_statement(self):
|
||||
assert create_schema_sql("deerflow") == 'CREATE SCHEMA IF NOT EXISTS "deerflow"'
|
||||
|
||||
def test_empty_schema_returns_none(self):
|
||||
assert create_schema_sql("") is None
|
||||
|
||||
@pytest.mark.parametrize("schema", ['a"; DROP SCHEMA public; --', "MySchema", "a b", "deerflow\n"])
|
||||
def test_rejects_non_plain_identifier(self, schema):
|
||||
# Defense-in-depth: the SQL-emitting boundary re-validates so a caller
|
||||
# that bypasses the pydantic config validator cannot inject.
|
||||
with pytest.raises(ValueError):
|
||||
create_schema_sql(schema)
|
||||
|
||||
|
||||
class TestDsnWithSearchPath:
|
||||
def test_empty_schema_returns_dsn_unchanged(self):
|
||||
dsn = "postgresql://u:p@h:5432/db"
|
||||
assert dsn_with_search_path(dsn, "") == dsn
|
||||
|
||||
def test_appends_options_query_encoded(self):
|
||||
dsn = "postgresql://u:p@h:5432/db"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
# libpq only decodes %XX in URI query values; '+' is NOT treated as a
|
||||
# space. The space MUST therefore be encoded as %20, never as '+'.
|
||||
assert "+" not in out
|
||||
assert "options=-c%20search_path%3Ddeerflow" in out
|
||||
parts = urlsplit(out)
|
||||
query = parse_qs(parts.query)
|
||||
assert query["options"] == ["-c search_path=deerflow"]
|
||||
|
||||
def test_merges_with_existing_query(self):
|
||||
dsn = "postgresql://u:p@h:5432/db?sslmode=require"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
query = parse_qs(urlsplit(out).query)
|
||||
assert query["sslmode"] == ["require"]
|
||||
assert query["options"] == ["-c search_path=deerflow"]
|
||||
|
||||
def test_replaces_existing_options_query(self):
|
||||
dsn = "postgresql://u:p@h:5432/db?options=-c%20search_path%3Dpublic"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
query = parse_qs(urlsplit(out).query)
|
||||
assert query["options"] == ["-c search_path=deerflow"]
|
||||
|
||||
def test_preserves_existing_options_query(self):
|
||||
dsn = "postgresql://u:p@h:5432/db?options=-c%20statement_timeout%3D5000"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
query = parse_qs(urlsplit(out).query)
|
||||
assert query["options"] == ["-c statement_timeout=5000 -c search_path=deerflow"]
|
||||
|
||||
def test_replaces_only_existing_search_path_option(self):
|
||||
dsn = "postgresql://u:p@h:5432/db?options=-c%20statement_timeout%3D5000%20-c%20search_path%3Dpublic"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
query = parse_qs(urlsplit(out).query)
|
||||
assert query["options"] == ["-c statement_timeout=5000 -c search_path=deerflow"]
|
||||
|
||||
def test_supports_keyword_dsn(self):
|
||||
pytest.importorskip("psycopg")
|
||||
from psycopg.conninfo import conninfo_to_dict
|
||||
|
||||
dsn = "host=localhost dbname=deerflow user=postgres"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
assert conninfo_to_dict(out) == {
|
||||
"host": "localhost",
|
||||
"dbname": "deerflow",
|
||||
"user": "postgres",
|
||||
"options": "-c search_path=deerflow",
|
||||
}
|
||||
|
||||
def test_preserves_keyword_dsn_options(self):
|
||||
pytest.importorskip("psycopg")
|
||||
from psycopg.conninfo import conninfo_to_dict
|
||||
|
||||
dsn = "host=localhost dbname=deerflow options='-c statement_timeout=5000'"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
assert conninfo_to_dict(out)["options"] == "-c statement_timeout=5000 -c search_path=deerflow"
|
||||
|
||||
def test_normalizes_sqlalchemy_driver_scheme(self):
|
||||
# DatabaseConfig.postgres_url may carry a +asyncpg suffix; the libpq DSN
|
||||
# produced for psycopg must drop the driver and still inject search_path.
|
||||
dsn = "postgresql+asyncpg://u:p@h:5432/db"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
parts = urlsplit(out)
|
||||
assert parts.scheme == "postgresql"
|
||||
query = parse_qs(parts.query)
|
||||
assert query["options"] == ["-c search_path=deerflow"]
|
||||
|
||||
def test_rejects_non_postgres_url_scheme(self):
|
||||
try:
|
||||
dsn_with_search_path("mysql://localhost/db", "deerflow")
|
||||
except ValueError as exc:
|
||||
assert "Unsupported PostgreSQL DSN scheme" in str(exc)
|
||||
else:
|
||||
raise AssertionError("Expected ValueError")
|
||||
|
||||
def test_roundtrip_preserves_host_and_db(self):
|
||||
dsn = "postgresql://u:p@h:5432/db"
|
||||
out = dsn_with_search_path(dsn, "deerflow")
|
||||
parts = urlsplit(out)
|
||||
assert parts.hostname == "h"
|
||||
assert parts.port == 5432
|
||||
assert parts.path == "/db"
|
||||
|
||||
def test_preserves_option_value_containing_space(self):
|
||||
# libpq's options parameter separates args on spaces unless they are
|
||||
# backslash-escaped. shlex.join would emit single-quotes, which libpq
|
||||
# treats as literal characters and would corrupt the option. A token
|
||||
# carrying a space must round-trip as a single backslash-escaped token.
|
||||
from deerflow.persistence.postgres_schema import _merge_search_path_option
|
||||
|
||||
merged = _merge_search_path_option(r"-c application_name=My\ App", "deerflow")
|
||||
assert "'" not in merged
|
||||
assert r"application_name=My\ App" in merged
|
||||
assert merged.endswith("-c search_path=deerflow")
|
||||
|
||||
def test_preserves_option_value_containing_tab(self):
|
||||
# Non-space whitespace (TAB/CR/LF) inside an existing escaped token must
|
||||
# also be re-escaped on re-join, otherwise libpq re-tokenizes on the bare
|
||||
# whitespace byte and the round-trip is lossy.
|
||||
from deerflow.persistence.postgres_schema import (
|
||||
_merge_search_path_option,
|
||||
_split_libpq_options,
|
||||
)
|
||||
|
||||
merged = _merge_search_path_option("-c application_name=My\\\tApp", "deerflow")
|
||||
assert "'" not in merged
|
||||
# The tab-bearing value must round-trip back to a single token.
|
||||
tokens = _split_libpq_options(merged)
|
||||
assert "application_name=My\tApp" in tokens
|
||||
assert merged.endswith("-c search_path=deerflow")
|
||||
|
||||
|
||||
class TestNormalizeLibpqDsn:
|
||||
def test_strips_asyncpg_driver_suffix(self):
|
||||
assert normalize_libpq_dsn("postgresql+asyncpg://u:p@h:5432/db") == "postgresql://u:p@h:5432/db"
|
||||
|
||||
def test_leaves_bare_postgres_scheme_unchanged(self):
|
||||
dsn = "postgresql://u:p@h:5432/db"
|
||||
assert normalize_libpq_dsn(dsn) == dsn
|
||||
|
||||
def test_leaves_keyword_dsn_unchanged(self):
|
||||
dsn = "host=localhost dbname=deerflow"
|
||||
assert normalize_libpq_dsn(dsn) == dsn
|
||||
|
||||
def test_rejects_non_postgres_scheme(self):
|
||||
with pytest.raises(ValueError, match="Unsupported PostgreSQL DSN scheme"):
|
||||
normalize_libpq_dsn("mysql://localhost/db")
|
||||
@ -1815,6 +1815,9 @@ skill_evolution:
|
||||
# checkpointer:
|
||||
# type: postgres
|
||||
# connection_string: postgresql://user:password@localhost:5432/deerflow
|
||||
# # Optional: place LangGraph checkpointer/store tables in this schema.
|
||||
# # Leave empty to use the server default search_path, usually public.
|
||||
# postgres_schema: deerflow
|
||||
|
||||
# ============================================================================
|
||||
# Database
|
||||
@ -1866,6 +1869,51 @@ skill_evolution:
|
||||
# database:
|
||||
# backend: postgres
|
||||
# postgres_url: $DATABASE_URL
|
||||
# # Optional: place DeerFlow ORM tables and LangGraph checkpointer/store
|
||||
# # tables in this schema. Leave empty to use the server default
|
||||
# # search_path, usually public.
|
||||
# postgres_schema: deerflow
|
||||
#
|
||||
# postgres_schema only takes effect for new tables. On startup DeerFlow runs
|
||||
# CREATE SCHEMA IF NOT EXISTS and pins the connection search_path; it never
|
||||
# moves existing tables, so third-party tables in `public` are left untouched.
|
||||
#
|
||||
# Migrating an EXISTING deployment from `public` to a dedicated schema:
|
||||
# 1. Stop the DeerFlow services.
|
||||
# 2. Move EVERY DeerFlow-owned table AND the Alembic state into the new schema
|
||||
# (run as the DB owner). Missing any table strands its rows: after restart
|
||||
# DeerFlow recreates an empty counterpart in the target schema and the old
|
||||
# rows stay invisible in `public`. Moving the ORM tables but leaving
|
||||
# `alembic_version` behind is especially dangerous -- bootstrap then treats
|
||||
# the target schema as unversioned, re-baselines it, and replays migrations.
|
||||
# CREATE SCHEMA IF NOT EXISTS deerflow;
|
||||
# -- Application ORM tables:
|
||||
# ALTER TABLE public.runs SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.run_events SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.threads_meta SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.feedback SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.users SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.agents SET SCHEMA deerflow;
|
||||
# -- IM channel tables:
|
||||
# ALTER TABLE public.channel_connections SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.channel_credentials SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.channel_oauth_states SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.channel_conversations SET SCHEMA deerflow;
|
||||
# -- Scheduled-task tables:
|
||||
# ALTER TABLE public.scheduled_tasks SET SCHEMA deerflow;
|
||||
# ALTER TABLE public.scheduled_task_runs SET SCHEMA deerflow;
|
||||
# -- Alembic migration state (REQUIRED -- see note above):
|
||||
# ALTER TABLE public.alembic_version SET SCHEMA deerflow;
|
||||
# The DeerFlow-owned set grows over time; discover any tables the list
|
||||
# above misses (as the DeerFlow DB owner/role) with:
|
||||
# SELECT table_name FROM information_schema.tables WHERE table_schema='public';
|
||||
# LangGraph checkpointer/store table names also vary by version -- list them
|
||||
# the same way, then ALTER ... SET SCHEMA deerflow for each (e.g. checkpoints,
|
||||
# checkpoint_blobs, checkpoint_writes, checkpoint_migrations, store,
|
||||
# store_migrations).
|
||||
# 3. Set postgres_schema: deerflow here and restart.
|
||||
# 4. Verify with `SHOW search_path;` and a smoke test (create a run, read
|
||||
# its history). The `public` schema should gain no new DeerFlow tables.
|
||||
database:
|
||||
backend: sqlite
|
||||
sqlite_dir: .deer-flow/data
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user