Revert "fix SQLAlchemy compatibility issues on macOS" (#8638)

* Revert "fix SQLAlchemy compatibility issues on macOS (#7724)"

This reverts commit 2d78626840.

* fix

* chore: add busy timeout pragma

astrbot/core/db/__init__.py

@@ -5,10 +5,7 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass
 
 from deprecated import deprecated
-from sqlalchemy import event
-from sqlalchemy.engine import make_url
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
-from sqlalchemy.pool import NullPool
 
 from astrbot.core.db.po import (
     ApiKey,
@@ -32,19 +29,6 @@ from astrbot.core.db.po import (
 )
 
 
-def _configure_sqlite_connection(dbapi_connection, connection_record) -> None:
-    cursor = dbapi_connection.cursor()
-    try:
-        cursor.execute("PRAGMA journal_mode=WAL")
-        cursor.execute("PRAGMA synchronous=NORMAL")
-        cursor.execute("PRAGMA cache_size=20000")
-        cursor.execute("PRAGMA temp_store=MEMORY")
-        cursor.execute("PRAGMA mmap_size=134217728")
-        cursor.execute("PRAGMA optimize")
-    finally:
-        cursor.close()
-
-
 @dataclass
 class BaseDatabase(abc.ABC):
     """数据库基类"""
@@ -57,29 +41,14 @@ class BaseDatabase(abc.ABC):
         # second write is attempted.  Setting timeout=30 tells SQLite to
         # wait up to 30 s for the lock, which is enough to ride out brief
         # write bursts from concurrent agent/metrics/session operations.
-        db_url = make_url(self.DATABASE_URL)
-        is_sqlite = db_url.get_backend_name() == "sqlite"
+        is_sqlite = "sqlite" in self.DATABASE_URL
         connect_args = {"timeout": 30} if is_sqlite else {}
-        engine_kwargs = {
-            "echo": False,
-            "future": True,
-            "connect_args": connect_args,
-        }
-        if is_sqlite:
-            # Keep SQLite async engines off SQLAlchemy's default async queue
-            # pool so packaged runtimes don't depend on dialect-specific pool
-            # event support.
-            engine_kwargs["poolclass"] = NullPool
         self.engine = create_async_engine(
             self.DATABASE_URL,
-            **engine_kwargs,
+            echo=False,
+            future=True,
+            connect_args=connect_args,
         )
-        if is_sqlite:
-            event.listen(
-                self.engine.sync_engine,
-                "connect",
-                _configure_sqlite_connection,
-            )
         self.AsyncSessionLocal = async_sessionmaker(
             self.engine,
             class_=AsyncSession,

astrbot/core/db/sqlite.py

@@ -53,6 +53,7 @@ class SQLiteDatabase(BaseDatabase):
         async with self.engine.begin() as conn:
             await conn.run_sync(SQLModel.metadata.create_all)
             await conn.execute(text("PRAGMA journal_mode=WAL"))
+            await conn.execute(text("PRAGMA busy_timeout=30000"))
             await conn.execute(text("PRAGMA synchronous=NORMAL"))
             await conn.execute(text("PRAGMA cache_size=20000"))
             await conn.execute(text("PRAGMA temp_store=MEMORY"))

astrbot/core/db/vec_db/faiss_impl/document_storage.py

@@ -5,11 +5,8 @@ from datetime import datetime
 from pathlib import Path
 
 from sqlalchemy import Column, Text, bindparam
-from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
-from sqlalchemy.pool import NullPool
-from sqlalchemy.schema import CreateTable
 from sqlmodel import Field, MetaData, SQLModel, col, func, select, text
 
 from astrbot.core import logger
@@ -63,7 +60,8 @@ class DocumentStorage:
         """Initialize the SQLite database and create the documents table if it doesn't exist."""
         await self.connect()
         async with self.engine.begin() as conn:  # type: ignore
-            await self._ensure_documents_table(conn)
+            # Create tables using SQLModel
+            await conn.run_sync(BaseDocModel.metadata.create_all)
 
             try:
                 await conn.execute(
@@ -93,59 +91,15 @@ class DocumentStorage:
             except BaseException:
                 pass
 
+            await conn.execute(
+                text(
+                    "CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_doc_id_unique ON documents(doc_id)",
+                ),
+            )
+
             await self._initialize_fts5(conn)
             await conn.commit()
 
-    async def _ensure_documents_table(self, executor) -> None:
-        """Create the document table from the SQLModel definition."""
-        result = await executor.execute(
-            text(
-                """
-                SELECT 1
-                FROM sqlite_master
-                WHERE type='table' AND name=:table_name
-                LIMIT 1
-                """,
-            ),
-            {"table_name": Document.__tablename__},
-        )
-        if result.scalar_one_or_none() is not None:
-            await self._ensure_doc_id_unique_index(executor)
-            return
-
-        create_table = CreateTable(Document.__table__, if_not_exists=True)  # type: ignore[attr-defined]
-
-        await executor.execute(
-            text(str(create_table.compile(dialect=sqlite.dialect())))
-        )
-        await self._ensure_doc_id_unique_index(executor)
-
-    async def _ensure_doc_id_unique_index(self, executor) -> None:
-        duplicate_result = await executor.execute(
-            text(
-                """
-                SELECT doc_id
-                FROM documents
-                GROUP BY doc_id
-                HAVING COUNT(*) > 1
-                LIMIT 1
-                """,
-            ),
-        )
-        if duplicate_result.scalar_one_or_none() is not None:
-            logger.warning(
-                "Skipping documents.doc_id unique index migration because duplicate "
-                f"doc_id values already exist in {self.db_path}.",
-            )
-            return
-
-        await executor.execute(
-            text(
-                "CREATE UNIQUE INDEX IF NOT EXISTS "
-                "idx_documents_doc_id_unique ON documents(doc_id)",
-            ),
-        )
-
     async def _initialize_fts5(self, executor) -> None:
         try:
             await self._create_fts5_table(executor, if_not_exists=True)
@@ -249,7 +203,6 @@ class DocumentStorage:
                 self.DATABASE_URL,
                 echo=False,
                 future=True,
-                poolclass=NullPool,
             )
             self.async_session_maker = sessionmaker(
                 self.engine,  # type: ignore

astrbot/core/knowledge_base/kb_db_sqlite.py

@@ -2,9 +2,8 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from sqlalchemy import delete, event, func, select, text, update
+from sqlalchemy import delete, func, select, text, update
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
-from sqlalchemy.pool import NullPool
 from sqlmodel import col, desc
 
 from astrbot.core import logger
@@ -20,19 +19,6 @@ if TYPE_CHECKING:
     from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
 
 
-def _configure_sqlite_connection(dbapi_connection, connection_record) -> None:
-    cursor = dbapi_connection.cursor()
-    try:
-        cursor.execute("PRAGMA journal_mode=WAL")
-        cursor.execute("PRAGMA synchronous=NORMAL")
-        cursor.execute("PRAGMA cache_size=20000")
-        cursor.execute("PRAGMA temp_store=MEMORY")
-        cursor.execute("PRAGMA mmap_size=134217728")
-        cursor.execute("PRAGMA optimize")
-    finally:
-        cursor.close()
-
-
 class KBSQLiteDatabase:
     def __init__(self, db_path: str | None = None) -> None:
         """初始化知识库数据库
@@ -54,12 +40,8 @@ class KBSQLiteDatabase:
         self.engine = create_async_engine(
             self.DATABASE_URL,
             echo=False,
-            poolclass=NullPool,
-        )
-        event.listen(
-            self.engine.sync_engine,
-            "connect",
-            _configure_sqlite_connection,
+            pool_pre_ping=True,
+            pool_recycle=3600,
         )
 
         # 创建会话工厂

main.py

@@ -54,9 +54,6 @@ def check_env() -> None:
 
     site_packages_path = get_astrbot_site_packages_path()
     if not is_packaged_desktop_runtime() and site_packages_path not in sys.path:
-        # Packaged desktop runtime keeps shared plugin dependencies out of the
-        # global import path so bundled core libraries don't mix with user-
-        # installed wheels from ~/.astrbot/data/site-packages.
         sys.path.append(site_packages_path)
 
     os.makedirs(get_astrbot_config_path(), exist_ok=True)

tests/unit/test_document_storage_fts.py

@@ -1,7 +1,6 @@
 import sqlite3
 
 import pytest
-from sqlalchemy.exc import IntegrityError
 
 from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage
 
@@ -102,38 +101,3 @@ async def test_document_storage_fts_recovers_from_legacy_non_fts_table(tmp_path)
     assert [result["doc_id"] for result in results] == ["legacy-fix"]
 
     await storage.close()
-
-
-@pytest.mark.asyncio
-async def test_document_storage_adds_unique_doc_id_index_to_existing_table(tmp_path):
-    db_path = tmp_path / "doc.db"
-    conn = sqlite3.connect(db_path)
-    conn.execute(
-        """
-        CREATE TABLE documents (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            doc_id VARCHAR NOT NULL,
-            text VARCHAR NOT NULL,
-            metadata TEXT,
-            created_at DATETIME,
-            updated_at DATETIME
-        )
-        """,
-    )
-    conn.execute(
-        "INSERT INTO documents (doc_id, text) VALUES ('legacy-chunk', 'legacy text')"
-    )
-    conn.commit()
-    conn.close()
-
-    storage = DocumentStorage(str(db_path))
-    await storage.initialize()
-
-    with pytest.raises(IntegrityError):
-        await storage.insert_document(
-            doc_id="legacy-chunk",
-            text="duplicate text",
-            metadata={},
-        )
-
-    await storage.close()