From 1caf96e6948670e4b9d35d786f3b89c82500794a Mon Sep 17 00:00:00 2001 From: oliver Date: Thu, 14 May 2026 10:24:23 +0800 Subject: [PATCH] fix: harden assistant SQLite/PostgreSQL dual-backend (singleton, timeouts, warnings) Co-authored-by: Cursor --- _local/system.env.example | 2 ++ docs/ASSISTANT_PG_MIGRATION.md | 7 +++--- docs/ENVIRONMENT_VARIABLES.md | 6 +++++ svc/config/database.py | 39 ++++++++++++++++++++++++++++- svc/persistence/assistant_store.py | 40 +++++++++++++++++++++++------- svc/persistence/pg_adapter.py | 16 +++++++++++- svc/persistence/sqlite_store.py | 11 +++++++- tests/test_database_backend.py | 23 +++++++++++++++++ 8 files changed, 129 insertions(+), 15 deletions(-) diff --git a/_local/system.env.example b/_local/system.env.example index c2a659b0..04aa81a5 100644 --- a/_local/system.env.example +++ b/_local/system.env.example @@ -82,6 +82,8 @@ OPS_ASSISTANT_DB_BACKEND= AIA_ASSISTANT_DATABASE_URL= OPS_ASSISTANT_DATABASE_URL= AIA_TEST_PG_URL= +# AIA_ASSISTANT_PG_CONNECT_TIMEOUT 原始 psycopg 连接超时（秒），默认 10，范围 1–120（见 svc/persistence/pg_adapter.py）。 +AIA_ASSISTANT_PG_CONNECT_TIMEOUT= # AIA_LOG_CHAT_MESSAGE_PERSIST 设为 1/true/on 时，每次成功写入 chat_message 后打一条 WARNING 日志（含 backend、session_id、message_id、role、event_type），便于在 gateway.err.log 里 grep「chat_message_persisted」；默认关闭。 # AIA_LOG_CHAT_MESSAGE_PERSIST= diff --git a/docs/ASSISTANT_PG_MIGRATION.md b/docs/ASSISTANT_PG_MIGRATION.md index e72e77ef..df2ae72b 100644 --- a/docs/ASSISTANT_PG_MIGRATION.md +++ b/docs/ASSISTANT_PG_MIGRATION.md @@ -141,7 +141,7 @@ cd D:\path\to\oclaw 2. 设置 **`AIA_ASSISTANT_DATABASE_URL`**（与导入目标一致）。 3. 重启 **网关**、**wiki worker**、**渠道 worker** 等所有持有 DB 连接的进程。 -**注意**：`get_assistant_store()` 在进程内会按 `(backend, 连接键)` 缓存；改环境后必须重启进程，否则会连旧库。 +**注意**：`get_assistant_store()` 在进程内会按 `(backend, 连接键)` 缓存；连接键对 PostgreSQL 使用 **URL 的哈希**（避免在内存键中携带明文口令）。改环境后必须重启进程，否则会连旧库。 --- @@ -180,7 +180,8 @@ cd D:\path\to\oclaw | 网关仍像连 SQLite | 环境未生效 / 未重启 | 检查进程环境、`assistant_store` 单例 | | `invalid_credentials` / 连接失败 | URL 或网络或权限 | 用 `psql` 或 `psycopg` 单独测连 | | 部分表未拷贝 | PG 无同名表 | 先 `alembic upgrade head`；查看 migrator 打印的 `skip (not in PG public schema)` | -| Windows 上 Alembic 编码问题 | 控制台代码页 | 使用 UTF-8 终端或重定向日志到文件 | +| 仍为 SQLite 但已填 PG URL | 后端未切到 `postgresql` | 首次解析 URL 会 `UserWarning`；确认 `AIA_ASSISTANT_DB_BACKEND` 与重启 | +| 连接 PG 长时间无响应 | 网络/防火墙 | 可调 `AIA_ASSISTANT_PG_CONNECT_TIMEOUT`（默认 10s，见 `ENVIRONMENT_VARIABLES.md`） | --- @@ -191,7 +192,7 @@ cd D:\path\to\oclaw | `alembic.ini` | Alembic 配置，`script_location = assistant_migrations` | | `assistant_migrations/` | PG schema 版本链 | | `svc/persistence/ddl/postgresql_bootstrap.sql` | 初始 DDL（与首版迁移同源） | -| `svc/config/database.py` | 后端与 DSN 解析 | +| `svc/persistence/pg_adapter.py` | psycopg 连接（含 `connect_timeout`） | | `svc/persistence/assistant_store.py` | `get_assistant_store()` 工厂 | | `runtime/operations/scripts/migrate_assistant_sqlite_to_postgresql.py` | 数据导入 | | `runtime/operations/scripts/cutover_sqlite_to_postgresql.ps1` / `.sh` | 备份 + 预演 + 导入 | diff --git a/docs/ENVIRONMENT_VARIABLES.md b/docs/ENVIRONMENT_VARIABLES.md index fa76a78b..0134d2c8 100644 --- a/docs/ENVIRONMENT_VARIABLES.md +++ b/docs/ENVIRONMENT_VARIABLES.md @@ -636,6 +636,7 @@ - 默认：未设置（按 `sqlite` 处理） - 取值：`sqlite`（默认）或 `postgresql`（大小写不敏感，亦接受 `pg` / `postgres` 等别名） - 作用：主 assistant 持久化后端；设为 `postgresql` 时必须配置 `AIA_ASSISTANT_DATABASE_URL`（或 `OPS_*` / `*_PG_DSN` 别名） + - 说明：仍为 `sqlite` 时若设置了 PostgreSQL URL 环境变量，会在首次解析 SQLAlchemy URL 时发出**一次性** `UserWarning`（提示 URL 被忽略），避免误配割接后仍连 SQLite - 生效：`oclaw/svc/config/database.py`, `oclaw/svc/persistence/assistant_store.py` - `AIA_ASSISTANT_DATABASE_URL` / `OPS_ASSISTANT_DATABASE_URL`（及 `AIA_ASSISTANT_PG_DSN` 等别名） @@ -643,6 +644,11 @@ - 作用：PostgreSQL 连接串（`postgresql://…` 或 `postgresql+psycopg://…`） - 生效：`oclaw/svc/config/database.py`, Alembic `assistant_migrations/env.py` +- `AIA_ASSISTANT_PG_CONNECT_TIMEOUT` + - 默认：`10`（秒） + - 作用：原始 `psycopg` 连接（`SqliteStore._connect` 每条业务连接）的 `connect_timeout` 上限，避免错误主机/防火墙时长时间挂死；范围钳制为 `1..120` + - 生效：`oclaw/svc/persistence/pg_adapter.py` + - `AIA_TEST_PG_URL` - 默认：空 - 作用：仅测试用；设置后部分 pytest 会对真实 PostgreSQL 跑冒烟（如 `tests/test_database_backend.py`） diff --git a/svc/config/database.py b/svc/config/database.py index 43f35e29..57bfa057 100644 --- a/svc/config/database.py +++ b/svc/config/database.py @@ -3,6 +3,9 @@ from __future__ import annotations import os +import warnings + +_PG_URL_IGNORED_WARNED = False def assistant_db_backend() -> str: @@ -22,8 +25,31 @@ def assistant_db_backend() -> str: ) +def _warn_if_postgres_url_ignored_for_sqlite() -> None: + """Emit once if PostgreSQL DSN env vars are set but backend is still SQLite.""" + global _PG_URL_IGNORED_WARNED + if _PG_URL_IGNORED_WARNED or assistant_db_backend() != "sqlite": + return + for var in ( + "AIA_ASSISTANT_DATABASE_URL", + "OPS_ASSISTANT_DATABASE_URL", + "AIA_ASSISTANT_PG_DSN", + "OPS_ASSISTANT_PG_DSN", + ): + if str(os.getenv(var) or "").strip(): + _PG_URL_IGNORED_WARNED = True + warnings.warn( + f"{var} is set but assistant DB backend is sqlite; the PostgreSQL URL is ignored. " + "Set AIA_ASSISTANT_DB_BACKEND=postgresql (or pg/postgres) to use PostgreSQL.", + UserWarning, + stacklevel=2, + ) + return + + def assistant_sqlalchemy_url() -> str: """SQLAlchemy URL for the assistant store (sqlite or postgresql+psycopg).""" + _warn_if_postgres_url_ignored_for_sqlite() if assistant_db_backend() == "postgresql": raw = assistant_postgres_dsn() if raw.startswith("postgresql+") or raw.startswith("postgres+"): @@ -54,4 +80,15 @@ def assistant_postgres_dsn() -> str: return url -__all__ = ["assistant_db_backend", "assistant_postgres_dsn", "assistant_sqlalchemy_url"] +def reset_assistant_db_config_warnings_for_tests() -> None: + """Reset one-shot warnings (pytest only).""" + global _PG_URL_IGNORED_WARNED + _PG_URL_IGNORED_WARNED = False + + +__all__ = [ + "assistant_db_backend", + "assistant_postgres_dsn", + "assistant_sqlalchemy_url", + "reset_assistant_db_config_warnings_for_tests", +] diff --git a/svc/persistence/assistant_store.py b/svc/persistence/assistant_store.py index 0140441f..66d09234 100644 --- a/svc/persistence/assistant_store.py +++ b/svc/persistence/assistant_store.py @@ -2,13 +2,20 @@ from __future__ import annotations +import hashlib +import threading from pathlib import Path -from svc.config.database import assistant_db_backend, assistant_postgres_dsn +from svc.config.database import ( + assistant_db_backend, + assistant_postgres_dsn, + assistant_sqlalchemy_url, +) from svc.persistence.assistant_store_protocol import AssistantStoreProtocol _singleton: AssistantStoreProtocol | None = None _singleton_key: str | None = None +_lock = threading.Lock() def reset_assistant_store_singleton() -> None: @@ -29,6 +36,12 @@ def get_assistant_store() -> AssistantStoreProtocol: does not re-run PostgreSQL bootstrap / orphan pruning on every HTTP or WS call (which could race with in-flight writes and make messages disappear after tool rounds). + The cache key for PostgreSQL uses a **hash of the SQLAlchemy URL** (not the raw DSN string) so logs + and debug dumps are less likely to embed credentials. + + When the effective key changes (e.g. tests switch ``AIA_ASSISTANT_DB_PATH`` or DSN), the SQLAlchemy + engine cache is cleared so pooled connections are not reused for the wrong database. + Tests should keep constructing ``SqliteStore(path)`` with an explicit file path; production code should prefer this factory for ``db_path()``-backed instances. @@ -40,18 +53,27 @@ def get_assistant_store() -> AssistantStoreProtocol: from svc.persistence.sqlite_store import SqliteStore if assistant_db_backend() == "postgresql": - key = f"postgresql::{assistant_postgres_dsn()}" + key = f"postgresql::{hashlib.sha256(assistant_sqlalchemy_url().encode()).hexdigest()}" else: key = f"sqlite::{Path(db_path()).resolve()}" - if _singleton is not None and _singleton_key == key: + inst = _singleton + if inst is not None and _singleton_key == key: + return inst + + with _lock: + if _singleton is not None and _singleton_key == key: + return _singleton + if _singleton_key is not None and _singleton_key != key: + from svc.persistence.db.engine import clear_assistant_engine_cache + + clear_assistant_engine_cache() + if assistant_db_backend() == "postgresql": + _singleton = SqliteStore(None, postgres_url=assistant_postgres_dsn()) + else: + _singleton = SqliteStore(db_path()) + _singleton_key = key return _singleton - if assistant_db_backend() == "postgresql": - _singleton = SqliteStore(None, postgres_url=assistant_postgres_dsn()) - else: - _singleton = SqliteStore(db_path()) - _singleton_key = key - return _singleton __all__ = ["get_assistant_store", "reset_assistant_store_singleton"] diff --git a/svc/persistence/pg_adapter.py b/svc/persistence/pg_adapter.py index 9583d6fa..04ef5d1e 100644 --- a/svc/persistence/pg_adapter.py +++ b/svc/persistence/pg_adapter.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os from typing import Any, Iterable, Sequence import psycopg @@ -61,8 +62,21 @@ class PgConnShim: self._raw.executemany(adapted, seq_of_params) +def _pg_connect_timeout_s() -> int: + raw = str(os.getenv("AIA_ASSISTANT_PG_CONNECT_TIMEOUT") or "10").strip() + try: + return max(1, min(int(raw), 120)) + except ValueError: + return 10 + + def connect_postgres(url: str) -> psycopg.Connection: - return psycopg.connect(normalize_psycopg_conninfo(url), row_factory=dict_row) + """Open a new PostgreSQL connection (short ``connect_timeout`` to avoid hanging on bad hosts).""" + return psycopg.connect( + normalize_psycopg_conninfo(url), + row_factory=dict_row, + connect_timeout=_pg_connect_timeout_s(), + ) __all__ = ["PgConnShim", "PgCursorShim", "connect_postgres", "normalize_psycopg_conninfo"] diff --git a/svc/persistence/sqlite_store.py b/svc/persistence/sqlite_store.py index 4c208dce..4f742d7d 100644 --- a/svc/persistence/sqlite_store.py +++ b/svc/persistence/sqlite_store.py @@ -325,7 +325,16 @@ class SqliteStore: self.db_path = str(db_path) self._postgres_url = "" self._use_pg = False - self._init_db() + try: + self._init_db() + except Exception as exc: + if self._use_pg: + raise RuntimeError( + "Assistant PostgreSQL initialization failed (connection, permissions, or missing tables). " + "Apply schema with `alembic upgrade head` (or `svc/persistence/ddl/postgresql_bootstrap.sql`), " + "verify AIA_ASSISTANT_DATABASE_URL, and inspect the chained exception." + ) from exc + raise @contextmanager def _connect(self) -> Iterator[Any]: diff --git a/tests/test_database_backend.py b/tests/test_database_backend.py index 7d99b16f..29d248cb 100644 --- a/tests/test_database_backend.py +++ b/tests/test_database_backend.py @@ -77,6 +77,29 @@ def test_get_assistant_store_postgresql_smoke(monkeypatch: pytest.MonkeyPatch) - assert msgs[0].content == "pg-smoke-body" +def test_postgresql_singleton_key_formula_excludes_password(monkeypatch: pytest.MonkeyPatch) -> None: + """Regression: cache key must not embed DSN credentials (logs, crash dumps).""" + import hashlib + + monkeypatch.setenv("AIA_ASSISTANT_DB_BACKEND", "postgresql") + monkeypatch.setenv("AIA_ASSISTANT_DATABASE_URL", "postgresql://u:top_secret@127.0.0.1:5432/db") + url = db_cfg.assistant_sqlalchemy_url() + key = f"postgresql::{hashlib.sha256(url.encode()).hexdigest()}" + assert "top_secret" not in key + assert key.startswith("postgresql::") + + +def test_warn_when_postgres_url_set_but_sqlite_backend(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + db_cfg.reset_assistant_db_config_warnings_for_tests() + monkeypatch.delenv("AIA_ASSISTANT_DB_BACKEND", raising=False) + monkeypatch.setenv("OPS_WORKSPACE_ROOT", str(tmp_path)) + monkeypatch.setenv("AIA_ASSISTANT_DB_PATH", str(tmp_path / "w.sqlite")) + monkeypatch.setenv("AIA_ASSISTANT_DATABASE_URL", "postgresql://u:p@127.0.0.1:5432/x") + with pytest.warns(UserWarning, match="ignored"): + db_cfg.assistant_sqlalchemy_url() + db_cfg.reset_assistant_db_config_warnings_for_tests() + + def test_assistant_db_backend_invalid(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("AIA_ASSISTANT_DB_BACKEND", "mysql") with pytest.raises(ValueError, match="Invalid assistant DB backend"):