提交内容:为数据库生命周期增加按白名单清单重建 Workflow Studio 表的受限流程,限定仅 MySQL 且检查当前连接数据库名称一致,避免误删非目标 schema。 启动流程中新增 DEER_FLOW_REBUILD_WORKFLOW_SCHEMA 环境变量入口与持久层重建 API,并在 lifespan 里按需执行;导出到 persistence 包,便于后续调用与测试复用。 迁移与模型修复:将 workflow 相关大文本字段切到 MySQL LONGTEXT,不再给 TEXT/LONGTEXT 设置 server_default;新增 draft_canvas_schema_json 的迁移回填与非空约束变更,规避 MySQL 5.7 对默认值限制。 测试覆盖:补充 workflow 定义 ddl 长文本断言与 schema 重建保护性测试(仅允许清单表、数据库名不匹配直接拒绝),并把新测试文件加入版本追踪。
726 lines
32 KiB
Python
726 lines
32 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 asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from importlib.util import find_spec
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.engine import make_url
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
|
|
def _json_serializer(obj: object) -> str:
|
|
"""JSON serializer with ensure_ascii=False for Chinese character support."""
|
|
return json.dumps(obj, ensure_ascii=False)
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_engine: AsyncEngine | None = None
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
# The event loop on which the async engine was created. asyncmy/asyncpg connection
|
|
# pools bind their connections to the loop that opened them; awaiting a write on a
|
|
# *different* loop reuses those connections cross-loop and raises
|
|
# "got Future attached to a different loop". We capture the engine's home loop here so
|
|
# best-effort writes coming from worker threads / other loops can be routed back to it
|
|
# (see ``fire_and_forget_db_write``).
|
|
_engine_loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
# These are the complete workflow-studio persistence tables owned by DeerFlow.
|
|
# Keep this explicit allow-list instead of treating ``workflow_%`` as a SQL
|
|
# wildcard: another application might share a schema and have its own table
|
|
# with that prefix.
|
|
_WORKFLOW_SCHEMA_TABLE_NAMES = (
|
|
"workflow_definitions",
|
|
"workflow_versions",
|
|
"workflow_runs",
|
|
"workflow_node_runs",
|
|
"workflow_run_events",
|
|
"workflow_run_artifacts",
|
|
"workflow_data_sources",
|
|
)
|
|
|
|
|
|
def _add_runtime_vendor_path() -> None:
|
|
"""Allow offline deployments to place optional DB drivers in DEER_FLOW_HOME.
|
|
|
|
The offline Docker image keeps the Python environment inside the image, while
|
|
``/app/backend/.deer-flow`` is a persistent bind mount. The MySQL switch
|
|
script installs ``asyncmy`` into ``{DEER_FLOW_HOME}/python-vendor`` so the
|
|
unchanged startup script can still import it after container recreation.
|
|
"""
|
|
try:
|
|
from deerflow.config.runtime_paths import runtime_home
|
|
|
|
vendor_dir = runtime_home() / "python-vendor"
|
|
except Exception:
|
|
vendor_dir = Path(".deer-flow/python-vendor").resolve()
|
|
|
|
vendor = str(vendor_dir)
|
|
if vendor_dir.exists() and vendor not in sys.path:
|
|
sys.path.insert(0, vendor)
|
|
|
|
|
|
def _mysql_driver_name(url: str) -> str:
|
|
if url.startswith("mysql+aiomysql://"):
|
|
return "aiomysql"
|
|
return "asyncmy"
|
|
|
|
|
|
async def _run_alembic_upgrade(url: str) -> None:
|
|
"""Apply pending Alembic migrations against ``url`` (``alembic upgrade head``).
|
|
|
|
Runs after ``Base.metadata.create_all`` so brand-new deployments still get
|
|
a complete schema even without Alembic, while existing deployments pick up
|
|
in-place schema fixes (column widenings, new indexes, etc.) at the next
|
|
restart — no manual ``alembic upgrade`` step required.
|
|
|
|
Notes:
|
|
|
|
* Each historical migration in ``migrations/versions/`` already guards
|
|
itself with ``_has_table`` / ``_has_column``, so re-running them against
|
|
a database that was originally built via ``create_all`` is a sequence of
|
|
no-ops up to the first migration that actually has work to do.
|
|
* The Alembic ``env.py`` calls ``asyncio.run`` internally, which cannot run
|
|
inside an already-running loop — so we hop into a worker thread.
|
|
* Failures here are logged but do not abort startup. The app remains
|
|
operable on the ``create_all`` schema; the operator just won't get
|
|
in-place fixes until they investigate.
|
|
"""
|
|
if os.getenv("DEER_FLOW_SKIP_ALEMBIC", "").lower() in ("1", "true", "yes"):
|
|
logger.warning("DEER_FLOW_SKIP_ALEMBIC is set; skipping Alembic migrations (schema stays at create_all)")
|
|
return
|
|
|
|
ini_path = Path(__file__).parent / "migrations" / "alembic.ini"
|
|
if not ini_path.exists():
|
|
logger.warning("Alembic config not found at %s; skipping auto-migrate", ini_path)
|
|
return
|
|
|
|
# Alembic's ``Config`` is backed by ConfigParser, which performs ``%``
|
|
# interpolation on every option value. URL-encoded credentials such as
|
|
# ``kland%4088124932`` (the ``@`` in ``kland@88124932``) trip that up with
|
|
# ``ValueError: invalid interpolation syntax ... at position N``. Doubling
|
|
# the percent signs neutralises the interpolation; ConfigParser un-doubles
|
|
# them back to the original value when read.
|
|
escaped_url = url.replace("%", "%%")
|
|
|
|
def _do_upgrade() -> None:
|
|
# Import inside the worker so a missing Alembic install can't break
|
|
# the engine import chain.
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
|
|
cfg = Config(str(ini_path))
|
|
cfg.set_main_option("sqlalchemy.url", escaped_url)
|
|
# The ``script_location`` in alembic.ini is ``%(here)s`` which points
|
|
# at the directory of alembic.ini itself — already correct for us.
|
|
command.upgrade(cfg, "head")
|
|
|
|
# Bound the upgrade so a stuck DDL (e.g. MySQL metadata-lock wait) cannot
|
|
# hang startup forever. Tunable via DEER_FLOW_ALEMBIC_TIMEOUT (seconds).
|
|
# NOTE: asyncio.wait_for cancels the *await* on timeout; the worker thread
|
|
# itself keeps running (Python can't kill threads). That's acceptable —
|
|
# the orphan thread will exit when its MySQL connection eventually errors
|
|
# or when the process restarts. The point here is to unblock startup.
|
|
try:
|
|
timeout_s = float(os.getenv("DEER_FLOW_ALEMBIC_TIMEOUT", "60"))
|
|
except ValueError:
|
|
timeout_s = 60.0
|
|
|
|
try:
|
|
await asyncio.wait_for(asyncio.to_thread(_do_upgrade), timeout=timeout_s)
|
|
logger.info("Alembic migrations applied (or already up-to-date)")
|
|
except TimeoutError:
|
|
logger.error(
|
|
"Alembic upgrade did not finish within %.0fs; continuing startup on create_all schema. "
|
|
"Likely cause: MySQL metadata-lock wait — check SHOW FULL PROCESSLIST for a blocking transaction.",
|
|
timeout_s,
|
|
)
|
|
except Exception:
|
|
logger.exception("Auto-applying Alembic migrations failed; continuing with create_all schema")
|
|
|
|
|
|
async def _auto_create_postgres_db(url: str) -> None:
|
|
"""Connect to the ``postgres`` maintenance DB and CREATE DATABASE.
|
|
|
|
The target database name is extracted from *url*. The connection is
|
|
made to the default ``postgres`` database on the same server using
|
|
``AUTOCOMMIT`` isolation (CREATE DATABASE cannot run inside a
|
|
transaction).
|
|
"""
|
|
from sqlalchemy import text
|
|
from sqlalchemy.engine.url import make_url
|
|
|
|
parsed = make_url(url)
|
|
db_name = parsed.database
|
|
if not db_name:
|
|
raise ValueError("Cannot auto-create database: no database name in URL")
|
|
|
|
# Connect to the default 'postgres' database to issue CREATE DATABASE
|
|
maint_url = parsed.set(database="postgres")
|
|
maint_engine = create_async_engine(maint_url, isolation_level="AUTOCOMMIT")
|
|
try:
|
|
async with maint_engine.connect() as conn:
|
|
await conn.execute(text(f'CREATE DATABASE "{db_name}"'))
|
|
logger.info("Auto-created PostgreSQL database: %s", db_name)
|
|
finally:
|
|
await maint_engine.dispose()
|
|
|
|
|
|
async def _auto_create_mysql_db(url: str) -> None:
|
|
"""Create the target MySQL database when it is missing.
|
|
|
|
``create_all`` can create tables, but MySQL requires the database/schema to
|
|
exist before the application URL can connect. We connect to the built-in
|
|
``mysql`` maintenance database on the same server and issue a conservative
|
|
``CREATE DATABASE IF NOT EXISTS``.
|
|
"""
|
|
|
|
parsed = make_url(url)
|
|
db_name = parsed.database
|
|
if not db_name:
|
|
raise ValueError("Cannot auto-create MySQL database: no database name in URL")
|
|
|
|
maint_url = parsed.set(database="mysql")
|
|
maint_engine = create_async_engine(
|
|
maint_url.render_as_string(hide_password=False),
|
|
isolation_level="AUTOCOMMIT",
|
|
pool_pre_ping=True,
|
|
)
|
|
quoted = "`" + db_name.replace("`", "``") + "`"
|
|
try:
|
|
async with maint_engine.connect() as conn:
|
|
await conn.execute(
|
|
text(
|
|
f"CREATE DATABASE IF NOT EXISTS {quoted} "
|
|
"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
|
|
)
|
|
)
|
|
logger.info("Auto-created MySQL database if missing: %s", db_name)
|
|
finally:
|
|
await maint_engine.dispose()
|
|
|
|
|
|
def _ensure_orm_columns_sync(sync_conn) -> None:
|
|
"""Add ORM-declared columns that are missing from already-existing tables.
|
|
|
|
``Base.metadata.create_all`` creates brand-new tables but never patches
|
|
existing ones, and Alembic auto-upgrade is disabled (see the commented-out
|
|
``_run_alembic_upgrade`` call) — so on an existing database a new
|
|
``mapped_column`` on an ORM model would otherwise error at first query
|
|
(e.g. ``no such column: roundtable_drafts.task_id`` on SQLite).
|
|
|
|
This sweep compares each existing table against the ORM metadata and issues
|
|
``ALTER TABLE … ADD COLUMN`` for every missing column that can be added
|
|
safely (nullable, or carrying a server_default). Indexes covering a newly
|
|
added column are created as well. Table-level ``UniqueConstraint`` objects
|
|
that are declared on the model but missing from the table are also created
|
|
(as a ``CREATE UNIQUE INDEX`` of the same name, which is what the migrations
|
|
emit too, and works on SQLite/MySQL/PostgreSQL alike) — without this a newly
|
|
declared uniqueness rule (e.g. ``active_dedupe_key``) would silently be absent
|
|
on an existing database, weakening dedupe.
|
|
Works on sqlite/mysql/postgres — the column DDL is compiled through the
|
|
active dialect, so portable types (``PortableJSON``/``PortableLongText``/
|
|
``BeijingDateTime``) resolve to the right backend type. Everything is
|
|
per-column best-effort: a failure is logged and startup continues.
|
|
"""
|
|
from sqlalchemy import UniqueConstraint
|
|
from sqlalchemy import inspect as sa_inspect
|
|
from sqlalchemy.schema import CreateColumn
|
|
|
|
from deerflow.persistence.base import Base
|
|
|
|
try:
|
|
import deerflow.persistence.models # noqa: F401 — register all ORM models on Base.metadata
|
|
except ImportError:
|
|
pass
|
|
|
|
inspector = sa_inspect(sync_conn)
|
|
existing_tables = set(inspector.get_table_names())
|
|
preparer = sync_conn.dialect.identifier_preparer
|
|
|
|
for table in Base.metadata.sorted_tables:
|
|
if table.name not in existing_tables:
|
|
continue
|
|
existing_cols = {c["name"] for c in inspector.get_columns(table.name)}
|
|
missing = [c for c in table.columns if c.name not in existing_cols]
|
|
if not missing:
|
|
continue
|
|
|
|
added: set[str] = set()
|
|
for col in missing:
|
|
if not col.nullable and col.server_default is None:
|
|
logger.warning(
|
|
"Cannot auto-add NOT NULL column without server_default: %s.%s — add it manually",
|
|
table.name,
|
|
col.name,
|
|
)
|
|
continue
|
|
try:
|
|
spec = CreateColumn(col).compile(dialect=sync_conn.dialect).string
|
|
sync_conn.exec_driver_sql(f"ALTER TABLE {preparer.quote(table.name)} ADD COLUMN {spec}")
|
|
added.add(col.name)
|
|
logger.info("Auto-added missing column %s.%s", table.name, col.name)
|
|
except Exception:
|
|
logger.exception("Failed to auto-add column %s.%s", table.name, col.name)
|
|
|
|
if not added:
|
|
continue
|
|
try:
|
|
existing_idx = {ix["name"] for ix in inspector.get_indexes(table.name)}
|
|
except Exception:
|
|
existing_idx = set()
|
|
for idx in table.indexes:
|
|
if idx.name in existing_idx or not any(c.name in added for c in idx.columns):
|
|
continue
|
|
try:
|
|
idx.create(sync_conn)
|
|
logger.info("Auto-created index %s on %s", idx.name, table.name)
|
|
except Exception:
|
|
logger.exception("Failed to auto-create index %s on %s", idx.name, table.name)
|
|
|
|
# ── 表级 UNIQUE 约束补齐 ──────────────────────────────────────────────────
|
|
# create_all 只为全新表建 UniqueConstraint;既有表上新增声明的唯一约束(如
|
|
# roundtable_jobs.active_dedupe_key)会静默缺失,去重/幂等的正确性就靠它,故
|
|
# 这里显式补建。迁移可能用 CREATE UNIQUE INDEX(同名)实现唯一性,故同时查
|
|
# get_unique_constraints 与 get_indexes 的名字集合,避免重复建。
|
|
for table in Base.metadata.sorted_tables:
|
|
if table.name not in existing_tables:
|
|
continue
|
|
named_uqs = [
|
|
c for c in table.constraints
|
|
if isinstance(c, UniqueConstraint) and c.name is not None
|
|
]
|
|
if not named_uqs:
|
|
continue
|
|
try:
|
|
existing_uq_names = {
|
|
uq["name"] for uq in inspector.get_unique_constraints(table.name) if uq.get("name")
|
|
}
|
|
except Exception:
|
|
existing_uq_names = set()
|
|
try:
|
|
existing_idx_names = {
|
|
ix["name"] for ix in inspector.get_indexes(table.name) if ix.get("name")
|
|
}
|
|
except Exception:
|
|
existing_idx_names = set()
|
|
for uq in named_uqs:
|
|
if uq.name in existing_uq_names or uq.name in existing_idx_names:
|
|
continue
|
|
# 用 CREATE UNIQUE INDEX(而非 ALTER TABLE ADD CONSTRAINT):迁移本身也用
|
|
# CREATE UNIQUE INDEX,且它在 SQLite/MySQL/PostgreSQL 都受支持(SQLite 不
|
|
# 支持 ALTER TABLE ADD CONSTRAINT),功能等价地强制唯一性。
|
|
cols = ", ".join(preparer.quote(c.name) for c in uq.columns)
|
|
try:
|
|
sync_conn.exec_driver_sql(
|
|
f"CREATE UNIQUE INDEX {preparer.quote(uq.name)} "
|
|
f"ON {preparer.quote(table.name)} ({cols})"
|
|
)
|
|
logger.info("Auto-added unique index %s on %s", uq.name, table.name)
|
|
except Exception:
|
|
logger.exception("Failed to add unique index %s on %s", uq.name, table.name)
|
|
|
|
|
|
def _is_mysql_unknown_database_error(exc: Exception) -> bool:
|
|
message = str(exc).lower()
|
|
return "unknown database" in message or "1049" in message
|
|
|
|
|
|
async def init_engine(
|
|
backend: str,
|
|
*,
|
|
url: str = "",
|
|
echo: bool = False,
|
|
pool_size: int = 5,
|
|
max_overflow: int = 10,
|
|
pool_timeout: int = 30,
|
|
pool_recycle: int = 1800,
|
|
sqlite_dir: str = "",
|
|
) -> None:
|
|
"""Create the async engine and session factory, then auto-create tables.
|
|
|
|
Args:
|
|
backend: "memory", "sqlite", "postgres", or "mysql".
|
|
url: SQLAlchemy async URL (for sqlite/postgres/mysql).
|
|
echo: Echo SQL to log.
|
|
pool_size: Base connection pool size (postgres/mysql).
|
|
max_overflow: Extra connections allowed beyond pool_size under load (postgres/mysql).
|
|
pool_timeout: Seconds to wait for a free pooled connection before erroring (postgres/mysql).
|
|
pool_recycle: Recycle connections after N seconds (postgres/mysql). Prevents stale
|
|
connections dropped by firewalls / MySQL wait_timeout / pgbouncer.
|
|
sqlite_dir: Directory to create for SQLite (ensured to exist).
|
|
"""
|
|
global _engine, _session_factory, _engine_loop
|
|
|
|
if backend == "memory":
|
|
logger.info("Persistence backend=memory -- ORM engine not initialized")
|
|
return
|
|
|
|
# Remember the loop that owns the engine's connection pool, so off-loop
|
|
# best-effort writes (metrics from worker threads, etc.) can be routed here
|
|
# instead of reusing pooled connections cross-loop.
|
|
try:
|
|
_engine_loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
_engine_loop = None
|
|
|
|
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 == "mysql":
|
|
_add_runtime_vendor_path()
|
|
driver = _mysql_driver_name(url)
|
|
if find_spec(driver) is None:
|
|
raise ImportError(
|
|
"database.backend is set to 'mysql' but the MySQL async driver is not installed.\n"
|
|
"Install asyncmy, or run switch_database_to_mysql.sh so it can install asyncmy into "
|
|
f"{Path('.deer-flow/python-vendor')} from the offline wheelhouse."
|
|
)
|
|
|
|
if backend == "sqlite":
|
|
import os
|
|
|
|
from sqlalchemy import event
|
|
|
|
os.makedirs(sqlite_dir or ".", exist_ok=True)
|
|
_engine = create_async_engine(url, echo=echo, json_serializer=_json_serializer)
|
|
|
|
# Enable WAL on every new connection. SQLite PRAGMA settings are
|
|
# per-connection, so we wire the listener instead of running PRAGMA
|
|
# once at startup. WAL gives concurrent reads + writers without
|
|
# blocking and is the standard recommendation for any production
|
|
# SQLite deployment (TC-UPG-06 in AUTH_TEST_PLAN.md). The companion
|
|
# ``synchronous=NORMAL`` is the safe-and-fast pairing — fsync only
|
|
# at WAL checkpoint boundaries instead of every commit.
|
|
# Note: we do not set PRAGMA busy_timeout here — Python's sqlite3
|
|
# driver already defaults to a 5-second busy timeout (see the
|
|
# ``timeout`` kwarg of ``sqlite3.connect``), and aiosqlite /
|
|
# SQLAlchemy's aiosqlite dialect inherit that default. Setting
|
|
# it again would be a no-op.
|
|
@event.listens_for(_engine.sync_engine, "connect")
|
|
def _enable_sqlite_wal(dbapi_conn, _record): # noqa: ARG001 — SQLAlchemy contract
|
|
cursor = dbapi_conn.cursor()
|
|
try:
|
|
cursor.execute("PRAGMA journal_mode=WAL;")
|
|
cursor.execute("PRAGMA synchronous=NORMAL;")
|
|
cursor.execute("PRAGMA foreign_keys=ON;")
|
|
finally:
|
|
cursor.close()
|
|
elif backend == "postgres":
|
|
_engine = create_async_engine(
|
|
url,
|
|
echo=echo,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_timeout=pool_timeout,
|
|
pool_pre_ping=True,
|
|
pool_recycle=pool_recycle,
|
|
json_serializer=_json_serializer,
|
|
)
|
|
elif backend == "mysql":
|
|
_engine = create_async_engine(
|
|
url,
|
|
echo=echo,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_timeout=pool_timeout,
|
|
pool_pre_ping=True,
|
|
pool_recycle=pool_recycle,
|
|
json_serializer=_json_serializer,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown persistence backend: {backend!r}")
|
|
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
|
|
# Auto-create tables for fresh deployments, then let Alembic apply any
|
|
# in-place schema fixes on top (see ``_run_alembic_upgrade`` below). Both
|
|
# are idempotent so repeating them across restarts is harmless.
|
|
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:
|
|
# Models package not yet available — tables won't be auto-created.
|
|
# This is expected during initial scaffolding or minimal installs.
|
|
logger.debug("deerflow.persistence.models not found; skipping auto-create tables")
|
|
|
|
try:
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
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
|
|
await _engine.dispose()
|
|
_engine = create_async_engine(url, echo=echo, pool_size=pool_size, max_overflow=max_overflow, pool_timeout=pool_timeout, pool_pre_ping=True, pool_recycle=pool_recycle, json_serializer=_json_serializer)
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
elif backend == "mysql" and _is_mysql_unknown_database_error(exc):
|
|
# Database not yet created. Create it, rebuild the app engine
|
|
# against the target DB, then let SQLAlchemy create all tables.
|
|
await _auto_create_mysql_db(url)
|
|
await _engine.dispose()
|
|
_engine = create_async_engine(
|
|
url,
|
|
echo=echo,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_timeout=pool_timeout,
|
|
pool_pre_ping=True,
|
|
pool_recycle=pool_recycle,
|
|
json_serializer=_json_serializer,
|
|
)
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
else:
|
|
raise
|
|
|
|
# Patch existing tables with ORM-declared columns that ``create_all``
|
|
# cannot add (it only creates brand-new tables). Best-effort: a failure is
|
|
# logged and startup continues on the current schema.
|
|
try:
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(_ensure_orm_columns_sync)
|
|
except Exception:
|
|
logger.exception("Auto-adding missing ORM columns failed; continuing with current schema")
|
|
|
|
# Apply in-place schema fixes (column widenings, etc.) that ``create_all``
|
|
# cannot do on its own. Failures don't abort startup — see the helper.
|
|
# await _run_alembic_upgrade(url)
|
|
|
|
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,
|
|
max_overflow=getattr(config, "max_overflow", 10),
|
|
pool_timeout=getattr(config, "pool_timeout", 30),
|
|
pool_recycle=getattr(config, "pool_recycle", 1800),
|
|
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
|
|
|
|
|
|
def get_engine_loop() -> asyncio.AbstractEventLoop | None:
|
|
"""Return the event loop that owns the engine's connection pool (or None)."""
|
|
return _engine_loop
|
|
|
|
|
|
async def rebuild_workflow_schema_for_mysql() -> list[str]:
|
|
"""Drop and recreate DeerFlow's workflow tables in the current MySQL schema.
|
|
|
|
This is a deliberately narrow recovery operation for an unused Workflow
|
|
Studio installation. It never uses a prefix wildcard, never changes the
|
|
selected database, and refuses to run unless the engine is connected to
|
|
MySQL and ``SELECT DATABASE()`` confirms the URL's exact target schema.
|
|
|
|
Returns the allow-listed workflow tables that existed before the rebuild.
|
|
"""
|
|
engine = get_engine()
|
|
if engine is None:
|
|
raise RuntimeError("Cannot rebuild workflow schema: persistence engine is not initialized")
|
|
if engine.dialect.name != "mysql":
|
|
raise RuntimeError(
|
|
"DEER_FLOW_REBUILD_WORKFLOW_SCHEMA is supported only for a MySQL persistence backend"
|
|
)
|
|
|
|
expected_database = engine.url.database
|
|
if not expected_database:
|
|
raise RuntimeError("Cannot rebuild workflow schema: MySQL URL has no database name")
|
|
|
|
# Make the operation safe when this helper is invoked independently from
|
|
# normal app startup as well.
|
|
import deerflow.persistence.models # noqa: F401 — register ORM models
|
|
from deerflow.persistence.base import Base
|
|
|
|
missing_models = [name for name in _WORKFLOW_SCHEMA_TABLE_NAMES if name not in Base.metadata.tables]
|
|
if missing_models:
|
|
raise RuntimeError(
|
|
"Cannot rebuild workflow schema: ORM metadata is missing " + ", ".join(missing_models)
|
|
)
|
|
workflow_tables = [Base.metadata.tables[name] for name in _WORKFLOW_SCHEMA_TABLE_NAMES]
|
|
|
|
def _rebuild(sync_conn) -> list[str]:
|
|
from sqlalchemy import inspect as sa_inspect
|
|
|
|
current_database = sync_conn.exec_driver_sql("SELECT DATABASE()").scalar_one_or_none()
|
|
if current_database != expected_database:
|
|
raise RuntimeError(
|
|
"Refusing workflow schema rebuild: connected database "
|
|
f"{current_database!r} does not match configured database {expected_database!r}"
|
|
)
|
|
|
|
existing = set(sa_inspect(sync_conn).get_table_names())
|
|
present = [name for name in _WORKFLOW_SCHEMA_TABLE_NAMES if name in existing]
|
|
# SQLAlchemy orders parent/child tables around the foreign key from
|
|
# workflow_versions to workflow_definitions. Only tables in the exact
|
|
# allow-list above are ever passed to drop_all/create_all.
|
|
Base.metadata.drop_all(sync_conn, tables=workflow_tables, checkfirst=True)
|
|
Base.metadata.create_all(sync_conn, tables=workflow_tables, checkfirst=True)
|
|
return present
|
|
|
|
async with engine.begin() as conn:
|
|
rebuilt = await conn.run_sync(_rebuild)
|
|
|
|
logger.warning(
|
|
"Rebuilt workflow schema in current MySQL database %s; tables reset: %s",
|
|
expected_database,
|
|
", ".join(rebuilt) if rebuilt else "none (created fresh)",
|
|
)
|
|
return rebuilt
|
|
|
|
|
|
def fire_and_forget_db_write(coro) -> None:
|
|
"""Schedule a best-effort DB write on the engine's home loop, from any loop/thread.
|
|
|
|
The async engine's connection pool is bound to the loop that created it
|
|
(``init_engine``, normally the Gateway's main loop). Running a write on a
|
|
*different* loop — e.g. a model/tool call that LangGraph executed in a worker
|
|
thread, which then does ``asyncio.run(store.record(...))`` — reuses those pooled
|
|
connections cross-loop and raises ``got Future attached to a different loop``.
|
|
Under concurrency this happens constantly (many calls land in worker threads at
|
|
once), spamming the log and churning the pool; single-threaded it almost never does.
|
|
|
|
This routes every such write back onto the engine's home loop:
|
|
- already on that loop -> ``create_task`` (in-loop fire-and-forget);
|
|
- on another loop/thread -> ``run_coroutine_threadsafe`` to the home loop.
|
|
|
|
The write itself stays best-effort: if no engine loop is captured (memory backend
|
|
/ not initialized) the coroutine is closed (avoids a "never awaited" warning) and
|
|
the call is a no-op. Never raises.
|
|
"""
|
|
loop = _engine_loop
|
|
if loop is None or loop.is_closed():
|
|
close = getattr(coro, "close", None)
|
|
if callable(close):
|
|
close()
|
|
return
|
|
try:
|
|
running = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
running = None
|
|
try:
|
|
if running is loop:
|
|
loop.create_task(coro)
|
|
else:
|
|
asyncio.run_coroutine_threadsafe(coro, loop)
|
|
except Exception: # noqa: BLE001 — best-effort; scheduling must never break the caller
|
|
close = getattr(coro, "close", None)
|
|
if callable(close):
|
|
close()
|
|
logger.debug("fire_and_forget_db_write: failed to schedule DB write", exc_info=True)
|
|
|
|
|
|
def _run_coro_on_throwaway_loop(coro):
|
|
"""Last-resort: run *coro* to completion on a fresh event loop.
|
|
|
|
Only reached when there is no usable engine home loop, or when the caller is
|
|
*synchronously* inside the home loop (routing back would deadlock). This
|
|
touches the shared connection pool from a throwaway loop, so it can still leak
|
|
a dead-loop-bound connection — but the common worker-thread path no longer
|
|
gets here (see ``run_db_blocking``).
|
|
"""
|
|
try:
|
|
asyncio.get_running_loop()
|
|
in_loop = True
|
|
except RuntimeError:
|
|
in_loop = False
|
|
|
|
if in_loop:
|
|
import concurrent.futures
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
|
return pool.submit(asyncio.run, coro).result()
|
|
return asyncio.run(coro)
|
|
|
|
|
|
def run_db_blocking(coro, *, timeout: float | None = 30.0):
|
|
"""Run an awaited DB coroutine from sync code, safe w.r.t. the engine's home loop.
|
|
|
|
The blocking-read sibling of :func:`fire_and_forget_db_write`: where that one
|
|
drops the result, this returns it — used by the sync→async bridges that must
|
|
read a row before proceeding (thread-path owner resolution, always-on skills,
|
|
skill visibility).
|
|
|
|
The async engine's connection pool is bound to the loop that created it
|
|
(``init_engine``). Running a DB coroutine on a *throwaway* loop
|
|
(``asyncio.run``) checks out pooled connections cross-loop and **leaks** them:
|
|
the connection ends up bound to a loop that is then closed, and when the pool
|
|
later recycles/disposes it, asyncmy's ``_terminate_graceful_close`` fires on the
|
|
wrong loop and raises ``got Future <...> attached to a different loop``. Under
|
|
concurrency (scheduler / subagent / sandbox-tool worker threads each spinning a
|
|
fresh loop) this floods the logs.
|
|
|
|
So when we are **off** the home loop — the common case, a worker thread with no
|
|
running loop — route the coroutine onto the home loop and block the *caller*
|
|
thread for its result. The pool is only ever touched on its own loop, so nothing
|
|
leaks. Only when there is no home loop, or we are synchronously inside it (where
|
|
blocking would deadlock), do we fall back to a throwaway loop.
|
|
|
|
Never swallows the coroutine's own exception — the sync bridges wrap this call in
|
|
their own try/except and degrade gracefully.
|
|
"""
|
|
loop = _engine_loop
|
|
try:
|
|
running = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
running = None
|
|
if loop is not None and not loop.is_closed() and running is not loop:
|
|
# Off the home loop (worker thread, or a different loop): execute the coro
|
|
# on the engine's loop and block this thread for the result. asyncmy's
|
|
# pooled connections stay on their own loop, so none leak cross-loop.
|
|
return asyncio.run_coroutine_threadsafe(coro, loop).result(timeout)
|
|
# No home loop, or we're already synchronously inside it: last resort.
|
|
return _run_coro_on_throwaway_loop(coro)
|
|
|
|
|
|
async def close_engine() -> None:
|
|
"""Dispose the engine, release all connections."""
|
|
global _engine, _session_factory, _engine_loop
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
logger.info("Persistence engine closed")
|
|
_engine = None
|
|
_session_factory = None
|
|
_engine_loop = None
|