1081 lines
48 KiB
Python
1081 lines
48 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 hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
import time
|
|
import weakref
|
|
from contextvars import ContextVar, Token
|
|
from importlib.util import find_spec
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from sqlalchemy import event, 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
|
|
_performance_request_context: ContextVar[tuple[str, Any] | None] = ContextVar("deerflow_performance_request", default=None)
|
|
|
|
|
|
def bind_performance_request(request_id: str) -> Token:
|
|
"""Bind a request ID to this task so slow SQL logs can be correlated safely."""
|
|
|
|
try:
|
|
task = asyncio.current_task()
|
|
except RuntimeError:
|
|
task = None
|
|
return _performance_request_context.set((request_id, weakref.ref(task) if task is not None else None))
|
|
|
|
|
|
def reset_performance_request(token: Token) -> None:
|
|
_performance_request_context.reset(token)
|
|
|
|
|
|
def _current_performance_request_id() -> str:
|
|
context = _performance_request_context.get()
|
|
if context is None:
|
|
return "-"
|
|
try:
|
|
task = asyncio.current_task()
|
|
except RuntimeError:
|
|
task = None
|
|
return context[0] if task is not None and context[1] is not None and context[1]() is task else "-"
|
|
|
|
|
|
def _slow_sql_threshold_ms() -> float:
|
|
try:
|
|
return max(0.0, min(float(os.getenv("DEERFLOW_SLOW_SQL_LOG_MS", "500")), 60_000.0))
|
|
except (TypeError, ValueError):
|
|
return 500.0
|
|
|
|
|
|
def _sql_operation(statement: str) -> str:
|
|
parts = statement.lstrip().split(None, 1)
|
|
candidate = parts[0].upper() if parts else "EMPTY"
|
|
return candidate if re.fullmatch(r"[A-Z_]+", candidate) else "OTHER"
|
|
|
|
|
|
def _sql_table(statement: str) -> str:
|
|
"""Return the target/source table only, without query text or values."""
|
|
match = re.search(
|
|
r"\b(?:INSERT\s+INTO|UPDATE|DELETE\s+FROM)\s+([`\"\w.]+)",
|
|
statement,
|
|
re.IGNORECASE,
|
|
)
|
|
if match is None:
|
|
match = re.search(r"\bFROM\s+([`\"\w.]+)", statement, re.IGNORECASE)
|
|
if match is None:
|
|
return "unknown"
|
|
return match.group(1).strip("`\".")[-128:] or "unknown"
|
|
|
|
|
|
def _sql_fingerprint(statement: str) -> str:
|
|
return hashlib.sha256(statement.encode("utf-8", "replace")).hexdigest()[:12]
|
|
|
|
|
|
def _database_error_metadata(exception: BaseException | None) -> tuple[str, str, str]:
|
|
"""Extract safe driver diagnostics without logging SQL or bound values."""
|
|
if exception is None:
|
|
return "unknown", "-", "-"
|
|
|
|
code = getattr(exception, "errno", None) or getattr(exception, "code", None)
|
|
args = getattr(exception, "args", ())
|
|
if code is None and args and isinstance(args[0], (int, str)):
|
|
candidate = str(args[0])
|
|
if candidate.isdigit():
|
|
code = candidate
|
|
|
|
sqlstate = getattr(exception, "sqlstate", None) or getattr(exception, "pgcode", None)
|
|
diagnostic = getattr(exception, "diag", None)
|
|
constraint = getattr(diagnostic, "constraint_name", None)
|
|
|
|
def _safe_identifier(value: Any) -> str:
|
|
text_value = str(value or "")
|
|
return text_value[:128] if re.fullmatch(r"[A-Za-z0-9_.-]{1,128}", text_value) else "-"
|
|
|
|
return _safe_identifier(code), _safe_identifier(sqlstate), _safe_identifier(constraint)
|
|
|
|
|
|
def _install_slow_sql_logging(engine: AsyncEngine, backend: str) -> None:
|
|
"""Log slow SQL timings without logging query text or bound values."""
|
|
|
|
sync_engine = engine.sync_engine
|
|
|
|
def before_cursor_execute(_conn, _cursor, statement, _parameters, context, _executemany):
|
|
if context is not None:
|
|
context._deerflow_perf_started = time.perf_counter()
|
|
context._deerflow_perf_operation = _sql_operation(statement)
|
|
context._deerflow_perf_statement = statement
|
|
|
|
def after_cursor_execute(_conn, _cursor, _statement, _parameters, context, executemany):
|
|
if context is None:
|
|
return
|
|
started = getattr(context, "_deerflow_perf_started", None)
|
|
if started is None:
|
|
return
|
|
elapsed_ms = round((time.perf_counter() - started) * 1000, 2)
|
|
if elapsed_ms >= _slow_sql_threshold_ms():
|
|
fingerprint = _sql_fingerprint(getattr(context, "_deerflow_perf_statement", ""))
|
|
logger.warning(
|
|
"PERF db_query request_id=%s backend=%s operation=%s table=%s duration_ms=%.2f executemany=%s fingerprint=%s outcome=ok",
|
|
_current_performance_request_id(),
|
|
backend,
|
|
getattr(context, "_deerflow_perf_operation", "OTHER"),
|
|
_sql_table(getattr(context, "_deerflow_perf_statement", "")),
|
|
elapsed_ms,
|
|
bool(executemany),
|
|
fingerprint,
|
|
)
|
|
|
|
def handle_error(exception_context):
|
|
context = exception_context.execution_context
|
|
started = getattr(context, "_deerflow_perf_started", None) if context is not None else None
|
|
elapsed_ms = round((time.perf_counter() - started) * 1000, 2) if started is not None else None
|
|
fingerprint = _sql_fingerprint(getattr(context, "_deerflow_perf_statement", "")) if context is not None else "unknown"
|
|
original_exception = exception_context.original_exception
|
|
error_code, sqlstate, constraint = _database_error_metadata(original_exception)
|
|
table = _sql_table(getattr(context, "_deerflow_perf_statement", "")) if context is not None else "unknown"
|
|
duplicate_key = error_code == "1062" or sqlstate == "23505" or (
|
|
table == "scheduled_task_runs"
|
|
and "UNIQUE constraint failed" in str(original_exception or "")
|
|
)
|
|
log = logger.info if table == "scheduled_task_runs" and duplicate_key else logger.error
|
|
outcome = "duplicate_ignored" if table == "scheduled_task_runs" and duplicate_key else "error"
|
|
log(
|
|
"PERF db_query request_id=%s backend=%s operation=%s table=%s duration_ms=%s executemany=%s fingerprint=%s outcome=%s error_type=%s error_code=%s sqlstate=%s constraint=%s",
|
|
_current_performance_request_id(),
|
|
backend,
|
|
getattr(context, "_deerflow_perf_operation", "OTHER") if context is not None else "OTHER",
|
|
table,
|
|
elapsed_ms,
|
|
bool(getattr(context, "executemany", False)) if context is not None else False,
|
|
fingerprint,
|
|
outcome,
|
|
type(original_exception).__name__ if original_exception is not None else "unknown",
|
|
error_code,
|
|
sqlstate,
|
|
constraint,
|
|
)
|
|
|
|
event.listen(sync_engine, "before_cursor_execute", before_cursor_execute)
|
|
event.listen(sync_engine, "after_cursor_execute", after_cursor_execute)
|
|
event.listen(sync_engine, "handle_error", handle_error)
|
|
|
|
# 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"
|
|
|
|
|
|
def _pooled_engine_kwargs(*, backend, url, echo, pool_size, max_overflow, pool_timeout, pool_recycle, command_timeout):
|
|
"""Use identical pool/deadline settings on initial connect and DB-create retry."""
|
|
kwargs = dict(
|
|
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,
|
|
)
|
|
# asyncpg command deadlines also bound the SELECT used by pool_pre_ping.
|
|
# These are not MySQL/psycopg connection arguments.
|
|
if backend in {"postgres", "gauss"} and make_url(url).drivername.endswith("+asyncpg"):
|
|
kwargs["connect_args"] = {"command_timeout": command_timeout}
|
|
return kwargs
|
|
|
|
|
|
def _pooled_engine_kwargs(*, backend, url, echo, pool_size, max_overflow, pool_timeout, pool_recycle, command_timeout):
|
|
"""Use identical pool/deadline settings on initial connect and DB-create retry."""
|
|
kwargs = dict(
|
|
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,
|
|
)
|
|
# asyncpg command deadlines also bound the SELECT used by pool_pre_ping.
|
|
# These are not MySQL/psycopg connection arguments.
|
|
if backend in {"postgres", "gauss"} and make_url(url).drivername.endswith("+asyncpg"):
|
|
kwargs["connect_args"] = {"command_timeout": command_timeout}
|
|
return kwargs
|
|
|
|
|
|
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-compatible 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). Missing ORM indexes 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/openGauss 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/gauss — 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]
|
|
|
|
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}")
|
|
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)
|
|
|
|
try:
|
|
existing_idx = {ix["name"] for ix in inspector.get_indexes(table.name)}
|
|
except Exception:
|
|
existing_idx = set()
|
|
for idx in table.indexes:
|
|
if not idx.name or idx.name in existing_idx:
|
|
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 _backfill_thread_meta_task_projections_sync(sync_conn) -> None:
|
|
"""Populate task lookup columns added by :func:`_ensure_orm_columns_sync`.
|
|
|
|
Alembic auto-upgrade is intentionally disabled in deployed Gateway
|
|
processes, while the ORM column synchronizer is always run. The latter
|
|
can add the task lookup columns and indexes but cannot derive values from
|
|
JSON metadata. Backfill task rows once at startup so an upgrade never
|
|
temporarily hides every historical task conversation from ``/threads/search``.
|
|
"""
|
|
from sqlalchemy import inspect as sa_inspect
|
|
|
|
table_name = "threads_meta"
|
|
required_columns = {"metadata_json", "task_id", "agent_id", "action_id"}
|
|
try:
|
|
inspector = sa_inspect(sync_conn)
|
|
if table_name not in inspector.get_table_names():
|
|
return
|
|
columns = {column["name"] for column in inspector.get_columns(table_name)}
|
|
if not required_columns <= columns:
|
|
return
|
|
|
|
dialect = sync_conn.dialect.name
|
|
if dialect == "sqlite":
|
|
task_id = "json_extract(metadata_json, '$.taskId')"
|
|
agent_id = "json_extract(metadata_json, '$.agent_id')"
|
|
action_id = "json_extract(metadata_json, '$.actionId')"
|
|
valid_json = "json_valid(metadata_json)"
|
|
elif dialect == "mysql":
|
|
task_id = "JSON_UNQUOTE(JSON_EXTRACT(metadata_json, '$.taskId'))"
|
|
agent_id = "JSON_UNQUOTE(JSON_EXTRACT(metadata_json, '$.agent_id'))"
|
|
action_id = "JSON_UNQUOTE(JSON_EXTRACT(metadata_json, '$.actionId'))"
|
|
valid_json = "JSON_VALID(metadata_json)"
|
|
elif dialect in {"postgresql", "opengauss"}:
|
|
task_id = "metadata_json ->> 'taskId'"
|
|
agent_id = "metadata_json ->> 'agent_id'"
|
|
action_id = "metadata_json ->> 'actionId'"
|
|
valid_json = "TRUE"
|
|
else:
|
|
return
|
|
|
|
result = sync_conn.execute(
|
|
text(
|
|
"UPDATE threads_meta "
|
|
f"SET task_id = {task_id}, agent_id = {agent_id}, action_id = {action_id} "
|
|
f"WHERE task_id IS NULL AND {valid_json} AND {task_id} IS NOT NULL"
|
|
)
|
|
)
|
|
if result.rowcount:
|
|
logger.info("Backfilled task lookup projections for %s historical thread(s)", result.rowcount)
|
|
except Exception:
|
|
# A malformed legacy JSON row or a backend-specific JSON function must
|
|
# not prevent startup. The generic metadata search remains available,
|
|
# and the failure is visible to the operator for remediation.
|
|
logger.exception("Failed to backfill thread task lookup projections")
|
|
|
|
|
|
def _is_mysql_unknown_database_error(exc: Exception) -> bool:
|
|
message = str(exc).lower()
|
|
return "unknown database" in message or "1049" in message
|
|
|
|
|
|
def _acquire_gauss_mysql_schema_bootstrap_lock():
|
|
"""Serialize Dolphin schema bootstrap across workers on this host."""
|
|
|
|
from deerflow.persistence.gauss_mysql_compat import is_gauss_mysql_product
|
|
|
|
if not is_gauss_mysql_product():
|
|
return None
|
|
try:
|
|
import fcntl
|
|
except ImportError:
|
|
# Production offline deployments are Linux. Keeping a no-lock fallback
|
|
# makes imports and unit tests work on Windows without changing MySQL.
|
|
logger.warning("fcntl is unavailable; openGauss schema bootstrap is not cross-process locked")
|
|
return None
|
|
|
|
try:
|
|
from deerflow.config.runtime_paths import runtime_home
|
|
|
|
home = runtime_home()
|
|
except Exception:
|
|
home = Path(os.getenv("DEER_FLOW_HOME", ".deer-flow")).resolve()
|
|
lock_path = home / "data" / "gauss-schema-bootstrap.lock"
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
handle = lock_path.open("a+", encoding="utf-8")
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
|
logger.info("Acquired openGauss schema bootstrap lock: %s", lock_path)
|
|
return handle
|
|
|
|
|
|
def _release_gauss_mysql_schema_bootstrap_lock(handle) -> None: # noqa: ANN001
|
|
if handle is None:
|
|
return
|
|
import fcntl
|
|
|
|
try:
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
|
finally:
|
|
handle.close()
|
|
|
|
|
|
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 | None = None,
|
|
command_timeout: float | None = 30,
|
|
sqlite_dir: str = "",
|
|
) -> None:
|
|
"""Create the async engine and session factory, then auto-create tables.
|
|
|
|
Args:
|
|
backend: "memory", "sqlite", "postgres", "gauss", or "mysql".
|
|
url: SQLAlchemy async URL (for sqlite/postgres/gauss/mysql).
|
|
echo: Echo SQL to log.
|
|
pool_size: Base connection pool size (postgres/gauss/mysql).
|
|
max_overflow: Extra connections allowed beyond pool_size under load (postgres/gauss/mysql).
|
|
pool_timeout: Seconds to wait for a free pooled connection before erroring (postgres/gauss/mysql).
|
|
pool_recycle: Recycle connections after N seconds (postgres/gauss/mysql). Prevents stale
|
|
connections dropped by firewalls / MySQL wait_timeout / pgbouncer.
|
|
command_timeout: asyncpg command deadline, including pool pre-ping (None disables).
|
|
sqlite_dir: Directory to create for SQLite (ensured to exist).
|
|
"""
|
|
global _engine, _session_factory, _engine_loop
|
|
|
|
if pool_recycle is None:
|
|
pool_recycle = 300 if backend == "postgres" else 1800
|
|
|
|
if pool_recycle is None:
|
|
pool_recycle = 300 if backend == "postgres" else 1800
|
|
|
|
if backend == "memory":
|
|
logger.info("Persistence backend=memory -- ORM engine not initialized")
|
|
return
|
|
|
|
from deerflow.persistence.gauss_mysql_compat import is_gauss_mysql_product
|
|
|
|
gauss_mysql = backend == "mysql" and is_gauss_mysql_product()
|
|
|
|
# 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 in {"postgres", "gauss"}:
|
|
try:
|
|
import asyncpg # noqa: F401
|
|
except ImportError:
|
|
extra = "gauss" if backend == "gauss" else "postgres"
|
|
raise ImportError(
|
|
f"database.backend is set to {backend!r} but asyncpg is not installed.\n"
|
|
f"Install it with:\n uv sync --extra {extra}\n"
|
|
"Or switch to backend: sqlite in config.yaml for single-node deployment."
|
|
) from None
|
|
if backend == "gauss" and find_spec("opengauss_sqlalchemy") is None:
|
|
raise ImportError(
|
|
"database.backend is set to 'gauss' but the openGauss SQLAlchemy dialect "
|
|
"is not installed.\nInstall it with:\n uv sync --extra gauss"
|
|
)
|
|
if backend == "gauss":
|
|
from deerflow.persistence.gauss_compat import install_gauss_dialect_compatibility
|
|
|
|
install_gauss_dialect_compatibility()
|
|
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 gauss_mysql:
|
|
from deerflow.persistence.gauss_mysql_compat import install_gauss_mysql_dialect_compatibility
|
|
|
|
install_gauss_mysql_dialect_compatibility(driver)
|
|
|
|
if backend == "sqlite":
|
|
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 in {"postgres", "gauss"}:
|
|
_engine = create_async_engine(
|
|
url,
|
|
**_pooled_engine_kwargs(
|
|
backend=backend, url=url, echo=echo, pool_size=pool_size,
|
|
max_overflow=max_overflow, pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle, command_timeout=command_timeout,
|
|
),
|
|
)
|
|
elif backend == "mysql":
|
|
_engine = create_async_engine(
|
|
url,
|
|
**_pooled_engine_kwargs(
|
|
backend=backend, url=url, echo=echo, pool_size=pool_size,
|
|
max_overflow=max_overflow, pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle, command_timeout=command_timeout,
|
|
),
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown persistence backend: {backend!r}")
|
|
|
|
_install_slow_sql_logging(_engine, backend)
|
|
_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")
|
|
|
|
async def _create_all_with_product_integrity(conn) -> None: # noqa: ANN001
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
if gauss_mysql:
|
|
from deerflow.persistence.gauss_schema import ensure_gauss_schema_sync
|
|
|
|
await conn.run_sync(ensure_gauss_schema_sync, Base.metadata)
|
|
|
|
# Schema bootstrap is deliberately best-effort. A single incompatible or
|
|
# permission-restricted table must not make the whole Gateway unavailable:
|
|
# routes backed by healthy existing tables can still work, while a route
|
|
# whose table is missing will report its own request-time error. Keep the
|
|
# auto-create-database retries, but contain failures from both the initial
|
|
# create_all and the retry path.
|
|
schema_bootstrap_lock = None
|
|
if gauss_mysql:
|
|
try:
|
|
schema_bootstrap_lock = _acquire_gauss_mysql_schema_bootstrap_lock()
|
|
except Exception:
|
|
# A lock-file permission problem is operationally important, but
|
|
# the global startup invariant still requires the Gateway to come
|
|
# up and report schema errors at the affected routes.
|
|
logger.exception("Could not acquire openGauss schema bootstrap lock; continuing without the host lock")
|
|
try:
|
|
try:
|
|
async with _engine.begin() as conn:
|
|
await _create_all_with_product_integrity(conn)
|
|
except Exception as exc:
|
|
if backend in {"postgres", "gauss"} and "does not exist" in str(exc).lower():
|
|
# 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,
|
|
**_pooled_engine_kwargs(
|
|
backend=backend, url=url, echo=echo, pool_size=pool_size,
|
|
max_overflow=max_overflow, pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle, command_timeout=command_timeout,
|
|
),
|
|
)
|
|
_install_slow_sql_logging(_engine, backend)
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
async with _engine.begin() as conn:
|
|
await _create_all_with_product_integrity(conn)
|
|
elif backend == "mysql" and not gauss_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,
|
|
**_pooled_engine_kwargs(
|
|
backend=backend, url=url, echo=echo, pool_size=pool_size,
|
|
max_overflow=max_overflow, pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle, command_timeout=command_timeout,
|
|
),
|
|
)
|
|
_install_slow_sql_logging(_engine, backend)
|
|
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
|
|
async with _engine.begin() as conn:
|
|
await _create_all_with_product_integrity(conn)
|
|
else:
|
|
raise
|
|
except Exception:
|
|
logger.exception(
|
|
"Automatic ORM schema creation failed; continuing Gateway startup "
|
|
"with the existing database schema"
|
|
)
|
|
finally:
|
|
try:
|
|
_release_gauss_mysql_schema_bootstrap_lock(schema_bootstrap_lock)
|
|
except Exception:
|
|
logger.exception("Could not release openGauss schema bootstrap lock")
|
|
|
|
# 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.
|
|
if gauss_mysql:
|
|
# Dolphin returns PostgreSQL-style DDL from SHOW CREATE TABLE. The
|
|
# stock MySQL inspector misreads it and may issue duplicate ALTERs,
|
|
# aborting the transaction with error 322. The native-catalog helper
|
|
# above is the only schema repair path for this product.
|
|
logger.info("Skipping MySQL reflection and JSON backfill for openGauss Dolphin")
|
|
else:
|
|
try:
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(_ensure_orm_columns_sync)
|
|
await conn.run_sync(_backfill_thread_meta_task_projections_sync)
|
|
except Exception:
|
|
logger.exception("Auto-updating ORM columns or task lookup projections 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 dialect=%s", backend, _engine.dialect.name)
|
|
|
|
|
|
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", None),
|
|
command_timeout=getattr(config, "command_timeout", 30),
|
|
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."""
|
|
from deerflow.persistence.loop_isolation import isolated_session_factory
|
|
|
|
return isolated_session_factory() or _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
|
|
uses an isolated, non-pooled engine for compatibility reads. The coroutine
|
|
must obtain its session factory inside its async body so the isolated
|
|
context applies; no connection is returned to the Gateway pool.
|
|
"""
|
|
from deerflow.persistence.loop_isolation import run_isolated_db_read
|
|
|
|
coro = run_isolated_db_read(coro)
|
|
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
|