Your Name
126 lines
4.5 KiB
Python
126 lines
4.5 KiB
Python
"""Nexus — Database Session Management (Async SQLAlchemy)
|
|
Engine is created lazily to allow testing without MySQL.
|
|
"""
|
|
|
|
import logging
|
|
|
|
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
|
from server.domain.models import Base
|
|
from server.config import settings
|
|
from server.infrastructure.database.engine_compat import patch_aiomysql_do_ping
|
|
|
|
logger = logging.getLogger("nexus.db")
|
|
|
|
_engine = None
|
|
_session_factory = None
|
|
|
|
|
|
def _ensure_engine():
|
|
"""Lazily create engine on first access (not at import time)"""
|
|
global _engine, _session_factory
|
|
if _engine is None:
|
|
patch_aiomysql_do_ping()
|
|
_engine = create_async_engine(
|
|
settings.DATABASE_URL.replace("mysql+pymysql", "mysql+aiomysql"),
|
|
pool_size=max(5, int(settings.DB_POOL_SIZE or 160)),
|
|
max_overflow=max(5, int(settings.DB_MAX_OVERFLOW or 120)),
|
|
pool_recycle=300,
|
|
pool_pre_ping=True,
|
|
)
|
|
_session_factory = async_sessionmaker(
|
|
autocommit=False,
|
|
autoflush=False,
|
|
bind=_engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False, # Fix MissingGreenlet: keep in-memory values after commit
|
|
)
|
|
|
|
|
|
def get_engine():
|
|
_ensure_engine()
|
|
return _engine
|
|
|
|
|
|
class _SessionLocalProxy:
|
|
"""Proxy that defers session factory creation until first call"""
|
|
def __call__(self):
|
|
_ensure_engine()
|
|
return _session_factory()
|
|
|
|
def __getattr__(self, name):
|
|
_ensure_engine()
|
|
return getattr(_session_factory, name)
|
|
|
|
|
|
AsyncSessionLocal = _SessionLocalProxy()
|
|
|
|
|
|
async def get_async_session():
|
|
"""[DEPRECATED] Yield async DB session for FastAPI Depends.
|
|
|
|
Use DbSessionMiddleware + request.state.db instead (see D7).
|
|
This function is kept for backward compatibility but new code
|
|
should NOT use it — it creates unmanaged sessions that can leak.
|
|
Prefer _get_request_session() from server.api.dependencies.
|
|
"""
|
|
_ensure_engine()
|
|
session = _session_factory()
|
|
try:
|
|
yield session
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize database tables and apply schema migrations"""
|
|
_ensure_engine()
|
|
async with _engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
# Apply incremental migrations for existing tables
|
|
await _apply_migrations(conn)
|
|
|
|
|
|
async def _apply_migrations(conn):
|
|
"""Apply incremental schema migrations (idempotent — safe to run on every startup)"""
|
|
from sqlalchemy import text
|
|
|
|
# Migration 1: Add updated_at column to admins table (if missing)
|
|
try:
|
|
result = await conn.execute(text(
|
|
"SELECT COUNT(*) FROM information_schema.COLUMNS "
|
|
"WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='admins' AND COLUMN_NAME='updated_at'"
|
|
))
|
|
count = result.scalar()
|
|
if count == 0:
|
|
await conn.execute(text(
|
|
"ALTER TABLE admins ADD COLUMN updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP"
|
|
))
|
|
logger.info("Migration applied: admins.updated_at column added")
|
|
except Exception as e:
|
|
logger.warning(f"Migration check failed (admins.updated_at): {e}")
|
|
|
|
# Migration 2: Add missing indexes (idempotent — CREATE INDEX IF NOT EXISTS)
|
|
_indexes = [
|
|
("idx_login_attempts_user_time", "login_attempts", "username, attempted_at"),
|
|
("idx_audit_logs_action_time", "audit_logs", "action, created_at"),
|
|
("idx_audit_logs_created_at", "audit_logs", "created_at"),
|
|
("idx_push_schedules_enabled", "push_schedules", "enabled"),
|
|
("idx_push_retry_pending", "push_retry_jobs", "status, next_retry_at"),
|
|
("idx_nodes_parent_id", "nodes", "parent_id"),
|
|
("idx_script_exec_script_id", "script_executions", "script_id"),
|
|
]
|
|
for idx_name, table, columns in _indexes:
|
|
try:
|
|
await conn.execute(text(
|
|
f"CREATE INDEX `{idx_name}` ON `{table}` ({columns})"
|
|
))
|
|
logger.info(f"Migration: index {idx_name} created on {table}({columns})")
|
|
except Exception as e:
|
|
err = str(e).lower()
|
|
if "duplicate" in err or "already exists" in err:
|
|
pass # Index already exists — expected on non-first-run
|
|
else:
|
|
# Genuine error (table missing, column missing, etc.) — must surface
|
|
logger.error(f"Migration FAILED: index {idx_name} on {table}({columns}): {e}")
|
|
raise
|