92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
"""Database dialect guard helpers for app startup."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
from .core.config import settings
|
||
from .core.logging import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
def is_missing_table_error(exc: Exception) -> bool:
|
||
text = str(exc).lower()
|
||
return (
|
||
"no such table" in text
|
||
or "undefined table" in text
|
||
or "relation" in text and "does not exist" in text
|
||
)
|
||
|
||
|
||
def normalize_dialect(value: str) -> str:
|
||
raw = str(value or "").strip().lower()
|
||
if not raw:
|
||
return ""
|
||
if raw.startswith("postgres"):
|
||
return "postgresql"
|
||
if raw.startswith("mysql"):
|
||
return "mysql"
|
||
if "+" in raw:
|
||
raw = raw.split("+", 1)[0]
|
||
if ":" in raw:
|
||
raw = raw.split(":", 1)[0]
|
||
return raw
|
||
|
||
|
||
def mask_db_url(url: str) -> str:
|
||
value = str(url or "").strip()
|
||
if "@" not in value:
|
||
return value
|
||
prefix, suffix = value.split("@", 1)
|
||
if "://" not in prefix:
|
||
return f"***@{suffix}"
|
||
scheme, creds = prefix.split("://", 1)
|
||
user = creds.split(":", 1)[0] if ":" in creds else creds
|
||
return f"{scheme}://{user}:***@{suffix}"
|
||
|
||
|
||
def check_database_dialect_guard(sync_conn) -> None:
|
||
configured = normalize_dialect(settings.database_url)
|
||
actual = normalize_dialect(sync_conn.dialect.name)
|
||
expected = normalize_dialect(os.getenv("PMFLOW_EXPECTED_DB_DIALECT", "postgresql")) or "postgresql"
|
||
|
||
env_urls = [
|
||
os.getenv("DATABASE_URL", ""),
|
||
os.getenv("SQLALCHEMY_DATABASE_URL", ""),
|
||
os.getenv("SYNC_DATABASE_URL", ""),
|
||
]
|
||
configured_families = {normalize_dialect(item) for item in env_urls if str(item or "").strip()}
|
||
configured_families.discard("")
|
||
|
||
logger.info(
|
||
"database_dialect_guard",
|
||
configured=configured,
|
||
actual=actual,
|
||
expected=expected or "unset",
|
||
database_url=mask_db_url(settings.database_url),
|
||
)
|
||
print(
|
||
"[DB-Guard] "
|
||
f"configured={configured}, actual={actual}, expected={expected or 'unset'}, "
|
||
f"url={mask_db_url(settings.database_url)}"
|
||
)
|
||
if configured != "postgresql":
|
||
raise RuntimeError(f"DATABASE_URL 必须为 PostgreSQL，当前 configured={configured or 'unset'}。")
|
||
if actual != "postgresql":
|
||
raise RuntimeError(f"运行时数据库必须为 PostgreSQL，当前 actual={actual or 'unset'}。")
|
||
if configured and actual and configured != actual:
|
||
raise RuntimeError(
|
||
f"数据库方言不一致: configured={configured}, actual={actual}。"
|
||
"请确保同一服务只连接单一数据库方言。"
|
||
)
|
||
if expected and actual and expected != actual:
|
||
raise RuntimeError(
|
||
f"数据库方言与 PMFLOW_EXPECTED_DB_DIALECT 不一致: expected={expected}, actual={actual}。"
|
||
)
|
||
if len(configured_families) > 1:
|
||
raise RuntimeError(
|
||
f"检测到多种数据库方言配置: {sorted(configured_families)}。"
|
||
"禁止同一服务混用多数据库方言。"
|
||
)
|