Your Name
938d26927f
- Agent per-server API Key认证 (agent.py) - JWT updated_at校验与会话超时 (auth_jwt.py) - 服务器凭据Fernet加密存储+密码脱敏 (servers.py) - WebSSH服务器级授权检查 (webssh.py) - Config同步去重+回滚机制 (sync_engine_v2.py) - DB层分页offset/limit (server_repo.py) - session.py logger修复 + @property死代码清理 - 心跳刷入rollback误用修复 (heartbeat_flush.py) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
116 lines
3.9 KiB
Python
116 lines
3.9 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
|
|
|
|
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:
|
|
_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,
|
|
)
|
|
|
|
|
|
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():
|
|
"""Dependency injection: yield async DB session"""
|
|
_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:
|
|
# Index already exists or table missing — both non-fatal
|
|
err = str(e).lower()
|
|
if "duplicate" in err or "already exists" in err:
|
|
pass # Index already exists — expected
|
|
else:
|
|
logger.warning(f"Migration: index {idx_name} skipped: {e}")
|