mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-20 15:11:09 +00:00
rayhpeng
1ff6b5f7ab
Introduce a unified database configuration (DatabaseConfig) that controls both the LangGraph checkpointer and the DeerFlow application persistence layer from a single `database:` config section. New modules: - deerflow.config.database_config — Pydantic config with memory/sqlite/postgres backends - deerflow.persistence — async engine lifecycle, DeclarativeBase with to_dict mixin, Alembic skeleton - deerflow.runtime.runs.store — RunStore ABC + MemoryRunStore implementation Gateway integration initializes/tears down the persistence engine in the existing langgraph_runtime() context manager. Legacy checkpointer config is preserved for backward compatibility. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
"""Async SQLAlchemy engine lifecycle management.
|
|
|
|
Initializes at Gateway startup, provides session factory for
|
|
repositories, disposes at shutdown.
|
|
|
|
When database.backend="memory", init_engine is a no-op and
|
|
get_session_factory() returns None. Repositories must check for
|
|
None and fall back to in-memory implementations.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_engine: AsyncEngine | None = None
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
|
|
async def init_engine(
|
|
backend: str,
|
|
*,
|
|
url: str = "",
|
|
echo: bool = False,
|
|
pool_size: int = 5,
|
|
sqlite_dir: str = "",
|
|
) -> None:
|
|
"""Create the async engine and session factory, then auto-create tables.
|
|
|
|
Args:
|
|
backend: "memory", "sqlite", or "postgres".
|
|
url: SQLAlchemy async URL (for sqlite/postgres).
|
|
echo: Echo SQL to log.
|
|
pool_size: Postgres connection pool size.
|
|
sqlite_dir: Directory to create for SQLite (ensured to exist).
|
|
"""
|
|
global _engine, _session_factory
|
|
|
|
if backend == "memory":
|
|
logger.info("Persistence backend=memory -- ORM engine not initialized")
|
|
return
|
|
|
|
if backend == "postgres":
|
|
try:
|
|
import asyncpg # noqa: F401
|
|
except ImportError:
|
|
raise ImportError("database.backend is set to 'postgres' but asyncpg is not installed.\nInstall it with:\n uv sync --extra postgres\nOr switch to backend: sqlite in config.yaml for single-node deployment.") from None
|
|
|
|
if backend == "sqlite":
|
|
import os
|
|
|
|
os.makedirs(sqlite_dir or ".", exist_ok=True)
|
|
_engine = create_async_engine(url, echo=echo)
|
|
elif backend == "postgres":
|
|
_engine = create_async_engine(
|
|
url,
|
|
echo=echo,
|
|
pool_size=pool_size,
|
|
pool_pre_ping=True,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown persistence backend: {backend!r}")
|
|
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
|
|
# Auto-create tables (dev convenience). Production should use Alembic.
|
|
from deerflow.persistence.base import Base
|
|
|
|
# Import all models so Base.metadata discovers them.
|
|
# When no models exist yet (scaffolding phase), this is a no-op.
|
|
try:
|
|
import deerflow.persistence.models # noqa: F401
|
|
except ImportError:
|
|
pass
|
|
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
logger.info("Persistence engine initialized: backend=%s", backend)
|
|
|
|
|
|
async def init_engine_from_config(config) -> None:
|
|
"""Convenience: init engine from a DatabaseConfig object."""
|
|
if config.backend == "memory":
|
|
await init_engine("memory")
|
|
return
|
|
await init_engine(
|
|
backend=config.backend,
|
|
url=config.app_sqlalchemy_url,
|
|
echo=config.echo_sql,
|
|
pool_size=config.pool_size,
|
|
sqlite_dir=config.sqlite_dir if config.backend == "sqlite" else "",
|
|
)
|
|
|
|
|
|
def get_session_factory() -> async_sessionmaker[AsyncSession] | None:
|
|
"""Return the async session factory, or None if backend=memory."""
|
|
return _session_factory
|
|
|
|
|
|
def get_engine() -> AsyncEngine | None:
|
|
"""Return the async engine, or None if not initialized."""
|
|
return _engine
|
|
|
|
|
|
async def close_engine() -> None:
|
|
"""Dispose the engine, release all connections."""
|
|
global _engine, _session_factory
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
logger.info("Persistence engine closed")
|
|
_engine = None
|
|
_session_factory = None
|