32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
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() |