221 lines
7.1 KiB
Python
221 lines
7.1 KiB
Python
"""数据库会话与引擎管理。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import traceback
|
||
from collections.abc import AsyncGenerator
|
||
|
||
from fastapi import Request
|
||
from sqlalchemy import event
|
||
from sqlalchemy.ext.asyncio import (
|
||
AsyncEngine,
|
||
AsyncSession,
|
||
async_sessionmaker,
|
||
create_async_engine,
|
||
)
|
||
from sqlalchemy.orm import Session
|
||
|
||
from .core.config import settings
|
||
from .core.logging import get_logger
|
||
from .models import OperationLog
|
||
from .utils.audited_write import get_audit_context
|
||
|
||
DATABASE_URL = settings.database_url
|
||
logger = get_logger(__name__)
|
||
|
||
AUDIT_ENFORCE_KEY = "_pmflow_audit_enforce"
|
||
AUDIT_WRITTEN_KEY = "_pmflow_audit_written"
|
||
AUDIT_BUSINESS_WRITE_KEY = "_pmflow_business_write"
|
||
AUDIT_METHOD_KEY = "_pmflow_request_method"
|
||
AUDIT_PATH_KEY = "_pmflow_interface_path"
|
||
AUDIT_TENANT_KEY = "_pmflow_tenant_id"
|
||
|
||
DDL_PREFIXES = (
|
||
"create ",
|
||
"alter ",
|
||
"drop ",
|
||
"truncate ",
|
||
"rename ",
|
||
"comment ",
|
||
"grant ",
|
||
"revoke ",
|
||
)
|
||
|
||
|
||
def _is_postgresql() -> bool:
|
||
normalized = str(DATABASE_URL or "").strip().lower()
|
||
return normalized.startswith("postgresql") or normalized.startswith("postgres")
|
||
|
||
|
||
if not _is_postgresql():
|
||
raise RuntimeError(
|
||
"仅支持 PostgreSQL 运行时连接，请设置 DATABASE_URL=postgresql+asyncpg://..."
|
||
)
|
||
|
||
|
||
engine_kwargs = {
|
||
"echo": settings.debug,
|
||
"future": True,
|
||
"pool_pre_ping": True,
|
||
"pool_size": max(5, int(settings.db_pool_size)),
|
||
"max_overflow": max(0, int(settings.db_max_overflow)),
|
||
"pool_timeout": max(5, int(settings.db_pool_timeout)),
|
||
"pool_recycle": max(60, int(settings.db_pool_recycle)),
|
||
# 强制数据库会话时区为北京时间，确保 now()/CURRENT_TIMESTAMP 为 Asia/Shanghai。
|
||
"connect_args": {"server_settings": {"timezone": "Asia/Shanghai"}},
|
||
}
|
||
|
||
engine: AsyncEngine = create_async_engine(
|
||
DATABASE_URL,
|
||
**engine_kwargs,
|
||
)
|
||
|
||
|
||
def _strip_sql_leading_comments(statement: str) -> str:
|
||
text = str(statement or "")
|
||
while True:
|
||
stripped = text.lstrip()
|
||
if stripped.startswith("/*") and "*/" in stripped:
|
||
text = stripped.split("*/", 1)[1]
|
||
continue
|
||
if stripped.startswith("--") and "\n" in stripped:
|
||
text = stripped.split("\n", 1)[1]
|
||
continue
|
||
return stripped
|
||
|
||
|
||
def _is_ddl_statement(statement: str) -> bool:
|
||
normalized = _strip_sql_leading_comments(statement).lower()
|
||
return any(normalized.startswith(prefix) for prefix in DDL_PREFIXES)
|
||
|
||
|
||
def _is_migration_stack() -> bool:
|
||
for frame in traceback.extract_stack(limit=50):
|
||
filename = str(frame.filename or "").replace("\\", "/").lower()
|
||
if "/alembic/" in filename:
|
||
return True
|
||
return False
|
||
|
||
|
||
@event.listens_for(engine.sync_engine, "before_cursor_execute", retval=True)
|
||
def _block_runtime_ddl(_conn, _cursor, statement, parameters, _context, _executemany): # pragma: no cover
|
||
if os.getenv("PMFLOW_ALLOW_RUNTIME_DDL", "0").lower() in {"1", "true", "yes"}:
|
||
return statement, parameters
|
||
if not _is_ddl_statement(str(statement or "")):
|
||
return statement, parameters
|
||
if _is_migration_stack():
|
||
return statement, parameters
|
||
context = get_audit_context()
|
||
logger.error(
|
||
"runtime_ddl_blocked",
|
||
trace_id=context.trace_id,
|
||
interface_path=context.interface_path,
|
||
statement_preview=str(statement or "")[:200],
|
||
)
|
||
raise RuntimeError("运行期禁止 DDL，请使用 Alembic migration。")
|
||
|
||
|
||
def _is_mutation_method(method: str) -> bool:
|
||
return str(method or "").upper() in {"POST", "PUT", "PATCH", "DELETE"}
|
||
|
||
|
||
def _is_business_change(session: Session, instance: object) -> bool:
|
||
if isinstance(instance, OperationLog):
|
||
session.info[AUDIT_WRITTEN_KEY] = True
|
||
return False
|
||
if instance in session.dirty and not session.is_modified(instance, include_collections=False):
|
||
return False
|
||
return True
|
||
|
||
|
||
def _collect_business_write_flag(session: Session) -> None:
|
||
if not session.info.get(AUDIT_ENFORCE_KEY):
|
||
return
|
||
if session.info.get(AUDIT_BUSINESS_WRITE_KEY):
|
||
return
|
||
changed = set(session.new).union(session.dirty).union(session.deleted)
|
||
if any(_is_business_change(session, row) for row in changed):
|
||
session.info[AUDIT_BUSINESS_WRITE_KEY] = True
|
||
|
||
|
||
def _build_auto_audit_log(session: Session) -> OperationLog:
|
||
context = get_audit_context()
|
||
tenant_id = int(context.tenant_id or session.info.get(AUDIT_TENANT_KEY) or 1)
|
||
user_id = int(context.user_id or os.getenv("PMFLOW_AUDIT_SYSTEM_USER_ID", "1"))
|
||
path = str(context.interface_path or session.info.get(AUDIT_PATH_KEY) or "unknown")[:255]
|
||
method = str(context.method or session.info.get(AUDIT_METHOD_KEY) or "UNKNOWN").upper()[:16]
|
||
trace_id = str(context.trace_id or "").strip()[:64] or None
|
||
caller = str(context.caller or "system").strip().lower()[:16] or "system"
|
||
entity_id = trace_id or f"{path}:{method.lower()}"
|
||
return OperationLog(
|
||
tenant_id=tenant_id,
|
||
user_id=user_id,
|
||
project_id=None,
|
||
entity_type="interface_write",
|
||
entity_id=entity_id[:128],
|
||
action=f"auto:{method.lower()}",
|
||
before_data=None,
|
||
after_data={
|
||
"method": method,
|
||
"path": path,
|
||
"trace_id": trace_id,
|
||
},
|
||
remark="auto generated audit log for write transaction",
|
||
trace_id=trace_id,
|
||
interface_path=path,
|
||
caller=caller,
|
||
)
|
||
|
||
|
||
@event.listens_for(Session, "before_flush")
|
||
def _mark_session_business_write(session: Session, _flush_context, _instances) -> None: # pragma: no cover
|
||
_collect_business_write_flag(session)
|
||
|
||
|
||
@event.listens_for(Session, "before_commit")
|
||
def _ensure_request_audit_log(session: Session) -> None: # pragma: no cover
|
||
if not session.info.get(AUDIT_ENFORCE_KEY):
|
||
return
|
||
if not session.info.get(AUDIT_BUSINESS_WRITE_KEY):
|
||
return
|
||
if session.info.get(AUDIT_WRITTEN_KEY):
|
||
return
|
||
session.add(_build_auto_audit_log(session))
|
||
session.info[AUDIT_WRITTEN_KEY] = True
|
||
|
||
|
||
@event.listens_for(Session, "after_commit")
|
||
@event.listens_for(Session, "after_rollback")
|
||
def _clear_audit_flags(session: Session) -> None: # pragma: no cover
|
||
session.info.pop(AUDIT_BUSINESS_WRITE_KEY, None)
|
||
session.info.pop(AUDIT_WRITTEN_KEY, None)
|
||
|
||
AsyncSessionLocal = async_sessionmaker(
|
||
bind=engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False,
|
||
autoflush=False,
|
||
)
|
||
|
||
|
||
def _resolve_tenant_from_request(request: Request) -> int:
|
||
raw = str(request.headers.get("X-Tenant-ID") or "").strip()
|
||
if not raw:
|
||
return 1
|
||
try:
|
||
return int(raw)
|
||
except ValueError:
|
||
return 1
|
||
|
||
|
||
async def get_db(request: Request) -> AsyncGenerator[AsyncSession, None]:
|
||
"""FastAPI 依赖：提供数据库会话。"""
|
||
async with AsyncSessionLocal() as session:
|
||
method = str(request.method or "").upper()
|
||
session.info[AUDIT_ENFORCE_KEY] = _is_mutation_method(method)
|
||
session.info[AUDIT_METHOD_KEY] = method[:16]
|
||
session.info[AUDIT_PATH_KEY] = str(request.url.path or "")[:255] or "/"
|
||
session.info[AUDIT_TENANT_KEY] = _resolve_tenant_from_request(request)
|
||
yield session
|