99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
"""数据库引擎与会话管理"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from sqlalchemy import create_engine, event
|
|
from sqlalchemy.orm import sessionmaker, declarative_base
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
DB_PATH = PROJECT_ROOT / "data" / "web.db"
|
|
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
DATABASE_URL = os.getenv("DATABASE_URL", f"sqlite:///{DB_PATH}")
|
|
|
|
connect_args = (
|
|
{"check_same_thread": False, "timeout": 30}
|
|
if DATABASE_URL.startswith("sqlite")
|
|
else {}
|
|
)
|
|
|
|
engine = create_engine(
|
|
DATABASE_URL,
|
|
connect_args=connect_args,
|
|
echo=False,
|
|
)
|
|
|
|
|
|
if DATABASE_URL.startswith("sqlite"):
|
|
@event.listens_for(engine, "connect")
|
|
def _set_sqlite_pragmas(dbapi_connection, connection_record):
|
|
"""提升 SQLite 并发写入稳定性。"""
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.execute("PRAGMA busy_timeout=30000")
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.close()
|
|
|
|
SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)
|
|
Base = declarative_base()
|
|
|
|
|
|
def get_db():
|
|
"""FastAPI 依赖：提供数据库会话，请求结束自动关闭。"""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def init_db():
|
|
"""执行数据库迁移 + 写入初始数据。"""
|
|
run_migrations()
|
|
_seed()
|
|
_encrypt_existing_sensitive_data()
|
|
|
|
|
|
def run_migrations():
|
|
"""运行 Alembic 迁移到最新版本。"""
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
|
|
config = Config(str(PROJECT_ROOT / "alembic.ini"))
|
|
config.set_main_option("script_location", str(PROJECT_ROOT / "web" / "backend" / "migrations"))
|
|
config.set_main_option("sqlalchemy.url", DATABASE_URL)
|
|
command.upgrade(config, "head")
|
|
|
|
|
|
def _seed():
|
|
"""写入默认超管账号和角色（账号密码通过环境变量配置）。"""
|
|
from .models import User
|
|
from .security import hash_password
|
|
|
|
admin_username = os.getenv("ADMIN_USERNAME", "admin")
|
|
admin_password = os.getenv("ADMIN_PASSWORD", "admin123")
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
if not db.query(User).first():
|
|
admin = User(
|
|
username=admin_username,
|
|
password_hash=hash_password(admin_password),
|
|
role="super_admin",
|
|
is_active=True,
|
|
remark="默认超级管理员",
|
|
)
|
|
db.add(admin)
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _encrypt_existing_sensitive_data():
|
|
"""启动时把历史明文敏感数据迁移为密文。"""
|
|
from loguru import logger
|
|
from .crypto_storage import encrypt_existing_sensitive_data
|
|
|
|
changed = encrypt_existing_sensitive_data(engine)
|
|
if changed:
|
|
logger.info(f"已加密历史敏感字段: {changed} 个值")
|