Spaces:

sidaier
/

fc

Sleeping

fc / src /db.py

feat(mvp): 新增铁血防守 MVP（领队结算/红色榜单/🛡️芯片），原 MVP 更名最佳球员

23938a8 verified 4 months ago

3 kB

"""数据库引擎与会话管理。
engine 由 config.DATABASE_URL 决定：
- sqlite:///... 本地文件（默认）
- postgresql+psycopg://... 云端（部署时通过环境变量切换）
"""
from __future__ import annotations

import logging
from contextlib import contextmanager
from typing import Iterator

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import Session, sessionmaker

from .config import DATABASE_URL
from .models import Base

logger = logging.getLogger(__name__)

# SQLite 在多线程（websocket）下需要关闭线程检查
_connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}

# pool_pre_ping：serverless Postgres（如 Neon）空闲会休眠，用连接前先探活，
# 避免拿到已失效的连接导致首次请求报错。
engine = create_engine(
    DATABASE_URL, future=True, connect_args=_connect_args, pool_pre_ping=True
)
SessionLocal = sessionmaker(bind=engine, expire_on_commit=False, future=True)


@contextmanager
def get_session() -> Iterator[Session]:
    """提供一个自动 commit / rollback / close 的会话上下文。"""
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def _run_light_migrations() -> None:
    """轻量迁移：给「已存在的表」补上后期新增的列（幂等）。
    create_all 只建缺失的表，不会给已有表加列；线上 Neon / 本地 SQLite 都已有
    历史数据，故这里用 inspector 检测缺列并 ALTER 补齐。区分方言以兼容布尔默认值
    （Postgres 用 false，SQLite 用 0）。新增列均带默认值，老数据零影响。
    """
    insp = inspect(engine)
    existing = set(insp.get_table_names())
    bool_false = "false" if engine.dialect.name == "postgresql" else "0"
    wanted = [
        ("users", "is_goalkeeper", f"BOOLEAN NOT NULL DEFAULT {bool_false}"),
        ("stats", "saves", "INTEGER NOT NULL DEFAULT 0"),
        ("matches", "mvp_locked", f"BOOLEAN NOT NULL DEFAULT {bool_false}"),
        ("matches", "mvp_user_id", "INTEGER"),  # 可空:None=未结算
        ("matches", "mvp2_user_id", "INTEGER"),  # 可空:双 MVP 第二人
        ("matches", "dmvp_user_id", "INTEGER"),  # 可空:铁血防守 MVP
    ]
    for table, col, ddl in wanted:
        if table not in existing:
            continue  # 新库由 create_all 直接建好（已含新列）
        if col in {c["name"] for c in insp.get_columns(table)}:
            continue
        with engine.begin() as conn:
            conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {col} {ddl}"))
        logger.info("迁移：已为 %s 表新增列 %s", table, col)


def init_db() -> None:
    """创建所有表（幂等）并执行轻量迁移补列。"""
    Base.metadata.create_all(engine)
    _run_light_migrations()
    logger.info("数据库已就绪：%s", DATABASE_URL)

	"""数据库引擎与会话管理。

	engine 由 config.DATABASE_URL 决定：
	- sqlite:///... 本地文件（默认）
	- postgresql+psycopg://... 云端（部署时通过环境变量切换）
	"""
	from __future__ import annotations

	import logging
	from contextlib import contextmanager
	from typing import Iterator

	from sqlalchemy import create_engine, inspect, text
	from sqlalchemy.orm import Session, sessionmaker

	from .config import DATABASE_URL
	from .models import Base

	logger = logging.getLogger(__name__)

	# SQLite 在多线程（websocket）下需要关闭线程检查
	_connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}

	# pool_pre_ping：serverless Postgres（如 Neon）空闲会休眠，用连接前先探活，
	# 避免拿到已失效的连接导致首次请求报错。
	engine = create_engine(
	DATABASE_URL, future=True, connect_args=_connect_args, pool_pre_ping=True
	)
	SessionLocal = sessionmaker(bind=engine, expire_on_commit=False, future=True)


	@contextmanager
	def get_session() -> Iterator[Session]:
	"""提供一个自动 commit / rollback / close 的会话上下文。"""
	session = SessionLocal()
	try:
	yield session
	session.commit()
	except Exception:
	session.rollback()
	raise
	finally:
	session.close()


	def _run_light_migrations() -> None:
	"""轻量迁移：给「已存在的表」补上后期新增的列（幂等）。

	create_all 只建缺失的表，不会给已有表加列；线上 Neon / 本地 SQLite 都已有
	历史数据，故这里用 inspector 检测缺列并 ALTER 补齐。区分方言以兼容布尔默认值
	（Postgres 用 false，SQLite 用 0）。新增列均带默认值，老数据零影响。
	"""
	insp = inspect(engine)
	existing = set(insp.get_table_names())
	bool_false = "false" if engine.dialect.name == "postgresql" else "0"
	wanted = [
	("users", "is_goalkeeper", f"BOOLEAN NOT NULL DEFAULT {bool_false}"),
	("stats", "saves", "INTEGER NOT NULL DEFAULT 0"),
	("matches", "mvp_locked", f"BOOLEAN NOT NULL DEFAULT {bool_false}"),
	("matches", "mvp_user_id", "INTEGER"), # 可空:None=未结算
	("matches", "mvp2_user_id", "INTEGER"), # 可空:双 MVP 第二人
	("matches", "dmvp_user_id", "INTEGER"), # 可空:铁血防守 MVP
	]
	for table, col, ddl in wanted:
	if table not in existing:
	continue # 新库由 create_all 直接建好（已含新列）
	if col in {c["name"] for c in insp.get_columns(table)}:
	continue
	with engine.begin() as conn:
	conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {col} {ddl}"))
	logger.info("迁移：已为 %s 表新增列 %s", table, col)


	def init_db() -> None:
	"""创建所有表（幂等）并执行轻量迁移补列。"""
	Base.metadata.create_all(engine)
	_run_light_migrations()
	logger.info("数据库已就绪：%s", DATABASE_URL)