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()