fix: Lazily access db engine for API key validation

_db_engine is not available when register_starlette_routes() is
called (it's set later during lifespan startup). Store the app
reference instead and access app.state._db_engine lazily.
This commit is contained in:
YueGuobin 2026-06-11 23:18:05 +08:00
parent d75746e029
commit aed8830052
No known key found for this signature in database

View File

@ -119,8 +119,10 @@ from .drawings import (
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
# Database engine reference — set during register_starlette_routes # FastAPI app reference — used to lazily access app.state._db_engine for API key validation.
_db_engine = None # The db engine is initialized during the lifespan startup, which runs AFTER
# register_starlette_routes() is called, so we cannot capture it at registration time.
_app = None
# ── Server ready state ──────────────────────────────────────────────── # ── Server ready state ────────────────────────────────────────────────
@ -203,22 +205,24 @@ async def _resolve_token(token: str) -> str | None:
pass pass
# Try API key # Try API key
if token.startswith("gns3_") and _db_engine is not None: if token.startswith("gns3_") and _app is not None:
try: db_engine = getattr(_app.state, "_db_engine", None)
async with AsyncSession(_db_engine, expire_on_commit=False) as db_session: if db_engine is not None:
repo = ApiKeysRepository(db_session) try:
query = select(models.ApiKey).where(models.ApiKey.revoked == False) async with AsyncSession(db_engine, expire_on_commit=False) as db_session:
result = await db_session.execute(query) repo = ApiKeysRepository(db_session)
for db_key in result.scalars().all(): query = select(models.ApiKey).where(models.ApiKey.revoked == False)
if bcrypt.checkpw(token.encode(), db_key.key_hash.encode()): result = await db_session.execute(query)
await repo.update_last_used(db_key.api_key_id) for db_key in result.scalars().all():
user_repo = UsersRepository(db_session) if bcrypt.checkpw(token.encode(), db_key.key_hash.encode()):
user = await user_repo.get_user(db_key.user_id) await repo.update_last_used(db_key.api_key_id)
if user: user_repo = UsersRepository(db_session)
svc = AuthService() user = await user_repo.get_user(db_key.user_id)
return svc.create_access_token(user.username, expires_in=5) if user:
except Exception: svc = AuthService()
pass return svc.create_access_token(user.username, expires_in=5)
except Exception:
pass
return None return None
@ -1378,8 +1382,8 @@ async def mcp_root():
def register_starlette_routes(app): def register_starlette_routes(app):
"""Mount MCP transports on the FastAPI app.""" """Mount MCP transports on the FastAPI app."""
global _db_engine global _app
_db_engine = getattr(app.state, "_db_engine", None) _app = app
sse_app = _make_auth_wrapper(mcp.sse_app(mount_path="")) sse_app = _make_auth_wrapper(mcp.sse_app(mount_path=""))
app.mount("/v3/mcp/transport", sse_app, name="mcp-sse") app.mount("/v3/mcp/transport", sse_app, name="mcp-sse")
log.info("MCP SSE server mounted at /v3/mcp/transport") log.info("MCP SSE server mounted at /v3/mcp/transport")