80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
from sqlalchemy import Column, Integer, String, Boolean, DateTime, Text
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.orm import sessionmaker
|
|
from datetime import datetime
|
|
|
|
from exceptions import ServiceNotConfiguredError
|
|
from services.config_service import ConfigService
|
|
|
|
Base = declarative_base()
|
|
|
|
|
|
class ApiKey(Base):
|
|
__tablename__ = "api_keys"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
provider = Column(String(50), unique=True, index=True, nullable=False)
|
|
api_key_encrypted = Column(Text, nullable=False) # Encrypted API key
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
updated_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
|
|
class ModelConfig(Base):
|
|
__tablename__ = "model_configs"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
provider = Column(String(50), nullable=False)
|
|
model_name = Column(String(100), nullable=False)
|
|
display_name = Column(String(100)) # Optional display name
|
|
is_active = Column(Boolean, default=True)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
updated_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
__table_args__ = {'mysql_charset': 'utf8mb4'}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Lazy engine / session factory — only created when first requested
|
|
# ---------------------------------------------------------------------------
|
|
_engine = None
|
|
_SessionLocal = None
|
|
|
|
|
|
def _get_engine():
|
|
from sqlalchemy import create_engine
|
|
global _engine
|
|
if _engine is None:
|
|
db_url = ConfigService.get_database_url()
|
|
if not db_url:
|
|
raise ServiceNotConfiguredError("数据库未配置")
|
|
_engine = create_engine(db_url)
|
|
return _engine
|
|
|
|
|
|
def _get_session_factory():
|
|
global _SessionLocal
|
|
if _SessionLocal is None:
|
|
_SessionLocal = sessionmaker(
|
|
autocommit=False, autoflush=False, bind=_get_engine()
|
|
)
|
|
return _SessionLocal
|
|
|
|
|
|
def reload_db_connection():
|
|
"""Dispose current engine and reset so next call rebuilds from config."""
|
|
global _engine, _SessionLocal
|
|
if _engine is not None:
|
|
_engine.dispose()
|
|
_engine = None
|
|
_SessionLocal = None
|
|
|
|
|
|
def get_db():
|
|
"""FastAPI dependency that yields a DB session."""
|
|
session_factory = _get_session_factory()
|
|
db = session_factory()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|