wiki_crawler/backend/database.py

from sqlalchemy import create_engine, MetaData, Table, event
from pgvector.sqlalchemy import Vector  # 必须导入这个
from .config import settings

class Database:
    def __init__(self):
        # 1. 创建引擎
        self.engine = create_engine(settings.DATABASE_URL, pool_pre_ping=True)
        
        # 2. 【核心修复】手动注册 vector 类型，让反射能识别它
        # 这告诉 SQLAlchemy：如果在数据库里看到名为 "vector" 的类型，请使用 pgvector 库的 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:
            # 自动从数据库加载表结构
            # 因为上面注册了 ischema_names，现在 chunks_table.c.embedding 就能被正确识别为 Vector 类型了
            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)
        except Exception as e:
            print(f"❌ 数据库表加载失败: {e}")

# 全局单例
db_instance = Database()