182 lines
8.4 KiB
Python
182 lines
8.4 KiB
Python
from sqlalchemy import create_engine
|
||
from sqlalchemy.orm import sessionmaker, declarative_base
|
||
from app.core.config import settings
|
||
|
||
# 数据库连接字符串，SQLite 相对路径已在配置层解析为 Server 目录下的绝对路径
|
||
SQLALCHEMY_DATABASE_URL = getattr(
|
||
settings, "DATABASE_URL", "sqlite:///./designercep.db"
|
||
)
|
||
|
||
# 创建数据库引擎
|
||
engine = create_engine(
|
||
SQLALCHEMY_DATABASE_URL,
|
||
pool_pre_ping=True, # 每次连接前 ping 一下，防止连接断开
|
||
pool_recycle=3600, # 1小时回收连接
|
||
pool_size=10, # 连接池大小
|
||
max_overflow=20, # 最大溢出连接数
|
||
connect_args={
|
||
"connect_timeout": 60, # 增加连接超时
|
||
"read_timeout": 60, # 增加读取超时
|
||
"write_timeout": 60 # 增加写入超时
|
||
} if not SQLALCHEMY_DATABASE_URL.startswith("sqlite") else {"check_same_thread": False}
|
||
)
|
||
|
||
# 会话工厂与 ORM 基类
|
||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||
Base = declarative_base()
|
||
|
||
def get_db():
|
||
# FastAPI 依赖注入使用的数据库会话
|
||
db = SessionLocal()
|
||
try:
|
||
yield db
|
||
finally:
|
||
db.close()
|
||
|
||
def init_db():
|
||
# 初始化数据库（创建所有 ORM 映射的表）
|
||
from app.models.user import User
|
||
from app.models.group import PluginGroup
|
||
from app.models.session import UserSession
|
||
from app.models.business import FeatureConfig, VipConfig, CheckInConfig, CheckInRecord, PointsHistory
|
||
from app.models.logs import PltProcessRecord, UserActionLog, PltPiece
|
||
from app.models.chat import ChatSession, ChatMessage
|
||
Base.metadata.create_all(bind=engine)
|
||
ensure_migrations()
|
||
seed_data()
|
||
|
||
def seed_data():
|
||
"""Ensure default data exists"""
|
||
from app.models.group import PluginGroup
|
||
from app.models.business import FeatureConfig, VipConfig, CheckInConfig
|
||
db = SessionLocal()
|
||
try:
|
||
# Seed Groups
|
||
default_group = db.query(PluginGroup).filter(PluginGroup.name == "default").first()
|
||
if not default_group:
|
||
print("Creating 'default' group...")
|
||
new_group = PluginGroup(name="default", comment="Default User Group")
|
||
db.add(new_group)
|
||
|
||
# Seed VIP Config
|
||
if db.query(VipConfig).count() == 0:
|
||
print("Seeding VIP Config...")
|
||
db.add(VipConfig(vip_type="vip", name="VIP会员", price=30.0, daily_quota=20, points_multiplier=1.5))
|
||
db.add(VipConfig(vip_type="svip", name="SVIP会员", price=88.0, daily_quota=-1, points_multiplier=2.0))
|
||
|
||
# Seed Checkin Config
|
||
if db.query(CheckInConfig).count() == 0:
|
||
print("Seeding Checkin Config...")
|
||
db.add(CheckInConfig(consecutive_days=1, base_points=10, bonus_points=0, total_points=10))
|
||
db.add(CheckInConfig(consecutive_days=3, base_points=10, bonus_points=5, total_points=15))
|
||
db.add(CheckInConfig(consecutive_days=7, base_points=10, bonus_points=20, total_points=30))
|
||
|
||
# Seed Features
|
||
if db.query(FeatureConfig).count() == 0:
|
||
print("Seeding Features...")
|
||
db.add(FeatureConfig(feature_key="ai_remove_bg", feature_name="智能抠图", points_cost=10, category="ai"))
|
||
|
||
db.commit()
|
||
except Exception as e:
|
||
print(f"Error seeding data: {e}")
|
||
finally:
|
||
db.close()
|
||
|
||
def ensure_migrations():
|
||
# 轻量级迁移：仅作为本地 SQLite 的兜底补列逻辑。
|
||
# 正式环境请优先使用 Alembic（见 Server/migrations/README.md）。
|
||
if not SQLALCHEMY_DATABASE_URL.startswith("sqlite"):
|
||
return
|
||
with engine.connect() as conn:
|
||
def has_column(table: str, col: str) -> bool:
|
||
try:
|
||
rows = conn.exec_driver_sql(f"PRAGMA table_info('{table}')").fetchall()
|
||
names = {r[1] for r in rows} if rows else set()
|
||
return col in names
|
||
except Exception:
|
||
return False
|
||
|
||
def add_col(table: str, col: str, type_sql: str):
|
||
try:
|
||
conn.exec_driver_sql(f"ALTER TABLE {table} ADD COLUMN {col} {type_sql}")
|
||
except Exception as e:
|
||
print(f"Migration error for {table}.{col}: {e}")
|
||
|
||
# user_sessions 需要的列
|
||
if not has_column("user_sessions", "login_at"):
|
||
add_col("user_sessions", "login_at", "TIMESTAMP NULL")
|
||
if not has_column("user_sessions", "logout_at"):
|
||
add_col("user_sessions", "logout_at", "TIMESTAMP NULL")
|
||
if not has_column("user_sessions", "duration_seconds"):
|
||
add_col("user_sessions", "duration_seconds", "INTEGER NULL")
|
||
if not has_column("user_sessions", "last_seen_at"):
|
||
add_col("user_sessions", "last_seen_at", "TIMESTAMP NULL")
|
||
|
||
# users 需要的列
|
||
if not has_column("users", "group_id"):
|
||
add_col("users", "group_id", "INTEGER NULL REFERENCES plugin_groups(id)")
|
||
if not has_column("users", "permissions"):
|
||
add_col("users", "permissions", "TEXT NULL")
|
||
if not has_column("users", "expire_date"):
|
||
add_col("users", "expire_date", "TIMESTAMP NULL")
|
||
|
||
# users email & verification columns
|
||
if not has_column("users", "email"):
|
||
add_col("users", "email", "VARCHAR(255) NULL")
|
||
if not has_column("users", "is_verified"):
|
||
add_col("users", "is_verified", "BOOLEAN DEFAULT 0")
|
||
if not has_column("users", "verification_code"):
|
||
add_col("users", "verification_code", "VARCHAR(6) NULL")
|
||
if not has_column("users", "reset_token"):
|
||
add_col("users", "reset_token", "VARCHAR(128) NULL")
|
||
if not has_column("users", "reset_token_expire"):
|
||
add_col("users", "reset_token_expire", "TIMESTAMP NULL")
|
||
|
||
# users profile & vip columns
|
||
if not has_column("users", "nickname"):
|
||
add_col("users", "nickname", "VARCHAR(50) NULL")
|
||
if not has_column("users", "avatar"):
|
||
add_col("users", "avatar", "VARCHAR(500) NULL")
|
||
if not has_column("users", "points"):
|
||
add_col("users", "points", "INTEGER DEFAULT 0")
|
||
if not has_column("users", "level"):
|
||
add_col("users", "level", "INTEGER DEFAULT 1")
|
||
if not has_column("users", "vip_type"):
|
||
add_col("users", "vip_type", "VARCHAR(20) DEFAULT 'none'")
|
||
if not has_column("users", "vip_expire"):
|
||
add_col("users", "vip_expire", "TIMESTAMP NULL")
|
||
if not has_column("users", "vip_daily_quota"):
|
||
add_col("users", "vip_daily_quota", "INTEGER DEFAULT 0")
|
||
if not has_column("users", "vip_quota_reset_date"):
|
||
add_col("users", "vip_quota_reset_date", "DATE NULL")
|
||
if not has_column("users", "total_check_in_days"):
|
||
add_col("users", "total_check_in_days", "INTEGER DEFAULT 0")
|
||
if not has_column("users", "consecutive_check_in"):
|
||
add_col("users", "consecutive_check_in", "INTEGER DEFAULT 0")
|
||
if not has_column("users", "last_check_in_date"):
|
||
add_col("users", "last_check_in_date", "DATE NULL")
|
||
if not has_column("users", "ai_provider"):
|
||
add_col("users", "ai_provider", "VARCHAR(32) NULL")
|
||
if not has_column("users", "ai_api_key"):
|
||
add_col("users", "ai_api_key", "TEXT NULL")
|
||
if not has_column("users", "ai_base_url"):
|
||
add_col("users", "ai_base_url", "VARCHAR(255) NULL")
|
||
if not has_column("users", "ai_chat_base_url"):
|
||
add_col("users", "ai_chat_base_url", "VARCHAR(255) NULL")
|
||
if not has_column("users", "ai_vision_base_url"):
|
||
add_col("users", "ai_vision_base_url", "VARCHAR(255) NULL")
|
||
if not has_column("users", "ai_image_base_url"):
|
||
add_col("users", "ai_image_base_url", "VARCHAR(255) NULL")
|
||
if not has_column("users", "ai_model"):
|
||
add_col("users", "ai_model", "VARCHAR(120) NULL")
|
||
if not has_column("users", "ai_vision_model"):
|
||
add_col("users", "ai_vision_model", "VARCHAR(120) NULL")
|
||
if not has_column("users", "ai_image_model"):
|
||
add_col("users", "ai_image_model", "VARCHAR(120) NULL")
|
||
|
||
# chat_sessions 会话工作流字段
|
||
if not has_column("chat_sessions", "active_skill_id"):
|
||
add_col("chat_sessions", "active_skill_id", "VARCHAR(80) NULL")
|
||
if not has_column("chat_sessions", "workflow_state"):
|
||
add_col("chat_sessions", "workflow_state", "TEXT NULL")
|