wiki_crawler/backend/core/database.py

from sqlalchemy import create_engine, MetaData, Table
from pgvector.sqlalchemy import Vector
from .config import settings

class Database:
    """
    数据库单例类
    负责初始化连接池并反射加载现有的表结构
    """
    def __init__(self):
        # 1. 创建引擎
        # pool_pre_ping=True 用于解决数据库连接长时间空闲后断开的问题
        self.engine = create_engine(settings.DATABASE_URL, pool_pre_ping=True)
        
        # 2. 注册 pgvector 类型
        # 这是为了让 SQLAlchemy 反射机制能识别数据库中的 'vector' 类型
        self.engine.dialect.ischema_names['vector'] = Vector
        
        self.metadata = MetaData()
        self.tasks = None
        self.queue = None
        self.chunks = None
        
        self._reflect_tables()

    def _reflect_tables(self):
        """自动从数据库加载表定义"""
        try:
            # autoload_with 会查询数据库元数据，自动填充 Column 信息
            self.tasks = Table('crawl_tasks', self.metadata, autoload_with=self.engine)
            self.queue = Table('crawl_queue', self.metadata, autoload_with=self.engine)
            self.chunks = Table('knowledge_chunks', self.metadata, autoload_with=self.engine)
            print("[INFO] Database tables reflected successfully.")
        except Exception as e:
            print(f"[ERROR] Failed to reflect tables: {e}")

# 全局数据库实例
db = Database()