rssKeeper/backend/database.py

"""数据库连接与初始化"""
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker, declarative_base
from config import DATABASE_URL

# SQLite 连接
engine = create_engine(
    f"sqlite:///{DATABASE_URL}",
    connect_args={"check_same_thread": False},
    echo=False,
)

# 启用 SQLite 外键约束
@event.listens_for(engine, "connect")
def set_sqlite_pragma(dbapi_conn, connection_record):
    cursor = dbapi_conn.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()

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():
    """初始化数据库表"""
    from models import Feed, Article, FetchLog  # noqa

    Base.metadata.create_all(bind=engine)
    _migrate(engine)
    init_fts5()


def _migrate(engine):
    """处理数据库增量迁移（添加新列）"""
    import logging
    logger = logging.getLogger(__name__)
    conn = engine.raw_connection()
    cursor = conn.cursor()

    # 获取 feeds 表现有列
    cursor.execute("PRAGMA table_info(feeds)")
    existing = {row[1] for row in cursor.fetchall()}

    migrations = [
        ("feeds", "error_type", "VARCHAR(32) DEFAULT ''"),
    ]
    for table, column, col_type in migrations:
        if column not in existing:
            logger.info(f"迁移: ALTER TABLE {table} ADD COLUMN {column} {col_type}")
            cursor.execute(f"ALTER TABLE {table} ADD COLUMN {column} {col_type}")

    conn.commit()

    # 对已有错误数据分类
    from rss_fetcher import classify_error
    cursor.execute("SELECT id, last_error FROM feeds WHERE last_error != '' AND (error_type IS NULL OR error_type = '')")
    rows = cursor.fetchall()
    for row in rows:
        feed_id, error = row
        etype = classify_error(error)
        if etype:
            cursor.execute("UPDATE feeds SET error_type = ? WHERE id = ?", (etype, feed_id))
    if rows:
        conn.commit()
        logger.info(f"迁移: 已分类 {len(rows)} 条历史错误")

    cursor.close()
    conn.close()


def init_fts5():
    """初始化 FTS5 全文搜索虚拟表"""
    conn = engine.raw_connection()
    cursor = conn.cursor()

    import logging
    logger = logging.getLogger(__name__)

    # 检查 FTS5 扩展是否可用
    try:
        cursor.execute("SELECT sqlite_compileoption_used('ENABLE_FTS5')")
        has_fts5 = cursor.fetchone()[0]
        if not has_fts5:
            logger.warning("SQLite 未启用 FTS5 扩展，全文搜索将不可用")
            return
    except Exception as e:
        logger.error(f"FTS5 检测失败: {e}")
        return

    # 创建 FTS5 虚拟表
    cursor.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS articles_fts USING fts5(
            title, content,
            content='articles',
            content_rowid='id'
        )
    """)

    # 创建触发器，自动同步 articles 表到 FTS5
    cursor.execute("""
        CREATE TRIGGER IF NOT EXISTS articles_fts_insert AFTER INSERT ON articles BEGIN
            INSERT INTO articles_fts(rowid, title, content)
            VALUES (new.id, new.title, new.content);
        END
    """)
    cursor.execute("""
        CREATE TRIGGER IF NOT EXISTS articles_fts_delete AFTER DELETE ON articles BEGIN
            INSERT INTO articles_fts(articles_fts, rowid, title, content)
            VALUES ('delete', old.id, old.title, old.content);
        END
    """)
    cursor.execute("""
        CREATE TRIGGER IF NOT EXISTS articles_fts_update AFTER UPDATE ON articles BEGIN
            INSERT INTO articles_fts(articles_fts, rowid, title, content)
            VALUES ('delete', old.id, old.title, old.content);
            INSERT INTO articles_fts(rowid, title, content)
            VALUES (new.id, new.title, new.content);
        END
    """)

    conn.commit()
    cursor.close()
    conn.close()