合并到远程

src/common/database/database.py

@@ -1,10 +1,10 @@
 from rich.traceback import install
 from pathlib import Path
 from contextlib import contextmanager
-from sqlalchemy import create_engine, event, text
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import event
 from sqlalchemy.engine import Engine
-from sqlalchemy import inspect as sqlalchemy_inspect
-from sqlalchemy.orm import Session, sessionmaker
+from sqlmodel import create_engine, Session
 from typing import TYPE_CHECKING, Generator
 
 if TYPE_CHECKING:
@@ -27,19 +27,12 @@ DATABASE_URL = f"sqlite:///{_DB_FILE}"
 def set_sqlite_pragma(dbapi_connection: "SQLite3Connection", connection_record):
     """
     为每个新的数据库连接设置 SQLite PRAGMA。
-
-    这些设置优化了并发性能和数据安全性:
-    - journal_mode=WAL: 启用预写式日志,提高并发性能
-    - cache_size: 设置缓存大小为 64MB
-    - foreign_keys: 启用外键约束
-    - synchronous=NORMAL: 平衡性能和数据安全
-    - busy_timeout: 设置1秒超时,避免锁定冲突
     """
     cursor = dbapi_connection.cursor()
     cursor.execute("PRAGMA journal_mode=WAL")
     cursor.execute("PRAGMA cache_size=-64000")  # 负值表示KB,64000KB = 64MB
     cursor.execute("PRAGMA foreign_keys=ON")
-    cursor.execute("PRAGMA synchronous=NORMAL")  # NORMAL 模式在WAL下是安全的
+    cursor.execute("PRAGMA synchronous=NORMAL")
     cursor.execute("PRAGMA busy_timeout=1000")  # 1秒超时
     cursor.close()
 
@@ -52,11 +45,12 @@ engine = create_engine(
     pool_pre_ping=True,
 )
 
-# 创建会话工厂
+# 创建会话工厂（使用 sqlmodel.Session）
 SessionLocal = sessionmaker(
     autocommit=False,
     autoflush=False,
     bind=engine,
+    class_=Session,
 )
 
 
@@ -96,7 +90,6 @@ def get_db_session(auto_commit: bool = True) -> Generator[Session, None, None]:
     session = SessionLocal()
     try:
         yield session
-        # 如果启用自动提交且没有异常,则提交事务
         if auto_commit:
             session.commit()
     except Exception:
@@ -132,59 +125,3 @@ def get_db() -> Generator[Session, None, None]:
         yield session
     finally:
         session.close()
-
-
-class _AtomicContext:
-    def __init__(self) -> None:
-        self._session: Session | None = None
-
-    def __enter__(self) -> Session:
-        self._session = SessionLocal()
-        self._session.begin()
-        return self._session
-
-    def __exit__(self, exc_type, exc, tb) -> None:
-        if self._session is None:
-            return
-        try:
-            if exc_type is None:
-                self._session.commit()
-            else:
-                self._session.rollback()
-        finally:
-            self._session.close()
-
-
-class DatabaseCompat:
-    """兼容旧 db 调用接口（Peewee 风格），底层使用 SQLAlchemy。"""
-
-    def connect(self, reuse_if_open: bool = True) -> None:
-        # SQLAlchemy 由 engine 按需管理连接，这里保留兼容入口。
-        _ = reuse_if_open
-
-    def create_tables(self, models: list[type], safe: bool = True) -> None:
-        _ = safe
-        tables = [model.__table__ for model in models if hasattr(model, "__table__")]
-        if not tables:
-            return
-        from sqlmodel import SQLModel
-
-        SQLModel.metadata.create_all(engine, tables=tables)
-
-    def atomic(self) -> _AtomicContext:
-        return _AtomicContext()
-
-    def execute_sql(self, sql: str):
-        with engine.connect() as conn:
-            result = conn.execute(text(sql))
-            conn.commit()
-            return result
-
-    def table_exists(self, model: type) -> bool:
-        if not hasattr(model, "__tablename__"):
-            return False
-        inspector = sqlalchemy_inspect(engine)
-        return inspector.has_table(model.__tablename__)
-
-
-db = DatabaseCompat()