Codex
105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""数据库连接解析器"""
|
|
import logging
|
|
import time
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.database import AsyncSessionLocal
|
|
from core.database_connection.model import DatabaseConnection
|
|
from core.database_connection.types import ConnectionInfo
|
|
from core.database_manager.service import parse_database_url
|
|
from utils.secret_crypto import decrypt_secret
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_CACHE: Dict[str, Tuple[ConnectionInfo, float]] = {}
|
|
_CACHE_TTL = 60.0
|
|
|
|
|
|
class ConnectionResolver:
|
|
@staticmethod
|
|
def invalidate_cache(code: Optional[str] = None) -> None:
|
|
if code:
|
|
_CACHE.pop(code, None)
|
|
else:
|
|
_CACHE.clear()
|
|
|
|
@staticmethod
|
|
def default_connection_info() -> ConnectionInfo:
|
|
db_info = parse_database_url(settings.DATABASE_URL or "")
|
|
if not db_info:
|
|
raise ValueError("Invalid DATABASE_URL configuration")
|
|
return ConnectionInfo(
|
|
code="default",
|
|
db_type=db_info["db_type"],
|
|
host=db_info["host"],
|
|
port=db_info["port"],
|
|
user=db_info["user"],
|
|
password=db_info["password"],
|
|
database=db_info["database"],
|
|
is_system=True,
|
|
display_name=db_info["database"] or "default",
|
|
)
|
|
|
|
@classmethod
|
|
async def resolve(
|
|
cls,
|
|
code: str,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> ConnectionInfo:
|
|
if code == "default":
|
|
return cls.default_connection_info()
|
|
|
|
now = time.time()
|
|
cached = _CACHE.get(code)
|
|
if cached and now - cached[1] < _CACHE_TTL:
|
|
return cached[0]
|
|
|
|
if db is None:
|
|
async with AsyncSessionLocal() as session:
|
|
info = await cls._resolve_from_db(session, code)
|
|
else:
|
|
info = await cls._resolve_from_db(db, code)
|
|
|
|
_CACHE[code] = (info, now)
|
|
return info
|
|
|
|
@classmethod
|
|
async def _resolve_from_db(cls, db: AsyncSession, code: str) -> ConnectionInfo:
|
|
result = await db.execute(
|
|
select(DatabaseConnection).where(
|
|
DatabaseConnection.code == code,
|
|
DatabaseConnection.is_deleted == False,
|
|
)
|
|
)
|
|
row = result.scalar_one_or_none()
|
|
if not row or not row.status:
|
|
raise ValueError(f"Database connection '{code}' not found or disabled")
|
|
if row.is_system:
|
|
return cls.default_connection_info()
|
|
|
|
password = decrypt_secret(row.password_enc or "")
|
|
if row.password_enc and not password:
|
|
logger.error(
|
|
"Failed to decrypt password for connection code=%s (check DB_CONN_SECRET_KEY / JWT_SECRET_KEY)",
|
|
row.code,
|
|
)
|
|
|
|
return ConnectionInfo(
|
|
code=row.code,
|
|
db_type=row.db_type,
|
|
host=row.host,
|
|
port=row.port,
|
|
user=row.user or "",
|
|
password=password,
|
|
database=row.default_database or "",
|
|
is_system=False,
|
|
display_name=row.name,
|
|
extra_options=row.extra_options or {},
|
|
)
|