Kris
c61d5f0163
本次提交完成了一系列核心功能迭代与优化： 1. 新增并完善了多个领域模型与端口定义，补充了`__all__`导出规范 2. 优化了会话、绑定、出箱等模块的数据模型，修复了时间字段类型不一致问题 3. 新增了代理ID解析、缓存发布等接口，扩展了系统能力 4. 重构了去重中间件逻辑，优化了空内容校验规则 5. 新增了认证中间件的匿名访问支持，完善了鉴权流程 6. 优化了SSE连接管理，增加了单会话连接上限限制 7. 重构了消息日志与仓储相关代码，将数据类迁移至对应模型目录 8. 新增了重复绑定校验、绑定更新接口，完善了绑定服务逻辑 9. 优化了健康检查逻辑，新增了环境变量控制启动时间线展示 10. 重构了出箱重试工作线程，使用缓存端口替代直接redis操作，新增了消息处理标记逻辑 11. 完善了飞书、Web、钩子等通道的翻译器逻辑，补充了账户ID传递 12. 新增了多种自定义异常类型，优化了异常映射与错误处理流程 13. 完善了配置热重载逻辑，同步认证凭证与校验器配置 14. 重构了Redis缓存实现，增加了异常捕获与包装
435 lines
23 KiB
Python
435 lines
23 KiB
Python
"""PostgreSQL 数据库管理器 - 支持知识库和业务数据"""
|
||
|
||
import json
|
||
import os
|
||
from contextlib import asynccontextmanager
|
||
|
||
from psycopg_pool import AsyncConnectionPool
|
||
from sqlalchemy import text
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||
from sqlalchemy.orm import declarative_base
|
||
from yuxi.storage.postgres.models_business import Base as BusinessBase
|
||
from yuxi.storage.postgres.models_knowledge import Base as KnowledgeBase
|
||
from yuxi.utils import logger
|
||
|
||
from yuxi.utils.singleton import SingletonMeta
|
||
|
||
# 合并两个 Base
|
||
CombinedBase = declarative_base()
|
||
|
||
# 继承所有表
|
||
for module in [KnowledgeBase, BusinessBase]:
|
||
for table_name in dir(module):
|
||
table = getattr(module, table_name)
|
||
if isinstance(table, type) and hasattr(table, "__tablename__"):
|
||
setattr(CombinedBase, table_name, table)
|
||
|
||
|
||
class PostgresManager(metaclass=SingletonMeta):
|
||
"""PostgreSQL 数据库管理器 - 支持知识库和业务数据"""
|
||
|
||
# 知识库 PostgreSQL URL 环境变量名
|
||
KB_DATABASE_URL_ENV = "POSTGRES_URL"
|
||
|
||
def __init__(self):
|
||
self.async_engine = None
|
||
self.AsyncSession = None
|
||
self.langgraph_pool = None
|
||
self._initialized = False
|
||
|
||
def initialize(self):
|
||
"""初始化数据库连接"""
|
||
if self._initialized:
|
||
return
|
||
|
||
db_url = os.getenv(self.KB_DATABASE_URL_ENV)
|
||
if not db_url:
|
||
logger.error(
|
||
f"环境变量 {self.KB_DATABASE_URL_ENV} 未设置，"
|
||
"请在 docker-compose.yml 或 .env 中配置 PostgreSQL 连接字符串"
|
||
)
|
||
return
|
||
|
||
try:
|
||
# 创建异步 SQLAlchemy 引擎
|
||
self.async_engine = create_async_engine(
|
||
db_url,
|
||
json_serializer=lambda obj: json.dumps(obj, ensure_ascii=False),
|
||
json_deserializer=json.loads,
|
||
pool_pre_ping=True,
|
||
pool_recycle=1800,
|
||
pool_size=10,
|
||
max_overflow=20,
|
||
)
|
||
|
||
# 创建异步会话工厂
|
||
self.AsyncSession = async_sessionmaker(
|
||
bind=self.async_engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False,
|
||
)
|
||
|
||
# ==========================================
|
||
# 2. 为 LangGraph 专门初始化一个原生 psycopg_pool
|
||
# ==========================================
|
||
# ⚠️ 注意：psycopg 不认识 "+asyncpg" 这样的 SQLAlchemy 方言标识。
|
||
# 如果你的 db_url 是 "postgresql+asyncpg://user:pwd@host/db"，
|
||
# 需要把它清洗成标准的 "postgresql://user:pwd@host/db"
|
||
langgraph_db_url = db_url.replace("+asyncpg", "").replace("+psycopg", "")
|
||
|
||
# 创建 LangGraph 专属连接池
|
||
self.langgraph_pool = AsyncConnectionPool(
|
||
conninfo=langgraph_db_url,
|
||
max_size=10, # 根据你的 Agent 并发情况设置，通常 5-10 足够了
|
||
kwargs={"autocommit": True}, # LangGraph Checkpoint 强依赖 autocommit
|
||
)
|
||
|
||
self._initialized = True
|
||
logger.info(f"PostgreSQL manager initialized for knowledge base: {db_url.split('@')[0]}://***")
|
||
except Exception as e:
|
||
logger.error(f"Failed to initialize PostgreSQL manager: {e}")
|
||
# 不抛出异常，允许应用启动，但在使用时会报错
|
||
|
||
def _check_initialized(self):
|
||
"""检查是否已初始化"""
|
||
if not self._initialized:
|
||
raise RuntimeError("PostgreSQL manager not initialized. Please check configuration.")
|
||
|
||
async def create_tables(self):
|
||
"""创建所有表（知识库和业务表）"""
|
||
self._check_initialized()
|
||
async with self.async_engine.begin() as conn:
|
||
await conn.run_sync(KnowledgeBase.metadata.create_all)
|
||
await conn.run_sync(BusinessBase.metadata.create_all)
|
||
logger.info("PostgreSQL tables created/checked (knowledge + business)")
|
||
|
||
async def create_business_tables(self):
|
||
"""创建所有业务数据表"""
|
||
self._check_initialized()
|
||
async with self.async_engine.begin() as conn:
|
||
await conn.run_sync(BusinessBase.metadata.create_all)
|
||
logger.info("PostgreSQL business tables created/checked")
|
||
|
||
async def drop_tables(self):
|
||
"""删除所有表（慎用！）"""
|
||
self._check_initialized()
|
||
async with self.async_engine.begin() as conn:
|
||
await conn.run_sync(BusinessBase.metadata.drop_all)
|
||
await conn.run_sync(KnowledgeBase.metadata.drop_all)
|
||
logger.info("PostgreSQL tables dropped")
|
||
|
||
async def ensure_knowledge_schema(self):
|
||
"""确保知识库 schema 包含所有必要字段"""
|
||
self._check_initialized()
|
||
stmts = [
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS embed_info JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS llm_info JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS query_params JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS additional_params JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS share_config JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS mindmap JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS sample_questions JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_bases ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS parent_id VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS original_filename VARCHAR(512)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS file_type VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS path VARCHAR(1024)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS minio_url VARCHAR(1024)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS markdown_file VARCHAR(1024)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS status VARCHAR(32)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS content_hash VARCHAR(128)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS file_size BIGINT",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS content_type VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS processing_params JSONB",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS is_folder BOOLEAN",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS error_message TEXT",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS created_by VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS updated_by VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ",
|
||
"ALTER TABLE IF EXISTS evaluation_benchmarks ADD COLUMN IF NOT EXISTS data_file_path VARCHAR(1024)",
|
||
"ALTER TABLE IF EXISTS evaluation_benchmarks ADD COLUMN IF NOT EXISTS created_by VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS evaluation_benchmarks ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS metrics JSONB",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS overall_score DOUBLE PRECISION",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS total_questions INTEGER",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS completed_questions INTEGER",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS started_at TIMESTAMPTZ",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS completed_at TIMESTAMPTZ",
|
||
"ALTER TABLE IF EXISTS evaluation_results ADD COLUMN IF NOT EXISTS created_by VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS evaluation_result_details ADD COLUMN IF NOT EXISTS gold_chunk_ids JSONB",
|
||
"ALTER TABLE IF EXISTS evaluation_result_details ADD COLUMN IF NOT EXISTS gold_answer TEXT",
|
||
"ALTER TABLE IF EXISTS evaluation_result_details ADD COLUMN IF NOT EXISTS generated_answer TEXT",
|
||
"ALTER TABLE IF EXISTS evaluation_result_details ADD COLUMN IF NOT EXISTS retrieved_chunks JSONB",
|
||
"ALTER TABLE IF EXISTS evaluation_result_details ADD COLUMN IF NOT EXISTS metrics JSONB",
|
||
# 扩展 db_id 字段长度以支持最长 75 字符的 ID（kb_private_ + 64字符hash）
|
||
"ALTER TABLE IF EXISTS knowledge_bases ALTER COLUMN db_id TYPE VARCHAR(80)",
|
||
"ALTER TABLE IF EXISTS knowledge_files ALTER COLUMN db_id TYPE VARCHAR(80)",
|
||
"ALTER TABLE IF EXISTS evaluation_benchmarks ALTER COLUMN db_id TYPE VARCHAR(80)",
|
||
"ALTER TABLE IF EXISTS evaluation_results ALTER COLUMN db_id TYPE VARCHAR(80)",
|
||
"CREATE INDEX IF NOT EXISTS idx_kb_type ON knowledge_bases(kb_type)",
|
||
"CREATE INDEX IF NOT EXISTS idx_kb_name ON knowledge_bases(name)",
|
||
"CREATE INDEX IF NOT EXISTS idx_kf_db_id ON knowledge_files(db_id)",
|
||
"CREATE INDEX IF NOT EXISTS idx_kf_parent ON knowledge_files(parent_id)",
|
||
"CREATE INDEX IF NOT EXISTS idx_kf_status ON knowledge_files(status)",
|
||
"CREATE INDEX IF NOT EXISTS idx_kf_hash ON knowledge_files(content_hash)",
|
||
"CREATE INDEX IF NOT EXISTS idx_eb_db_id ON evaluation_benchmarks(db_id)",
|
||
"CREATE INDEX IF NOT EXISTS idx_er_db_id ON evaluation_results(db_id)",
|
||
"CREATE INDEX IF NOT EXISTS idx_er_status ON evaluation_results(status)",
|
||
"CREATE INDEX IF NOT EXISTS idx_er_started ON evaluation_results(started_at DESC)",
|
||
"CREATE INDEX IF NOT EXISTS idx_erd_task ON evaluation_result_details(task_id)",
|
||
]
|
||
|
||
async with self.async_engine.begin() as conn:
|
||
for stmt in stmts:
|
||
await conn.execute(text(stmt))
|
||
|
||
async def ensure_business_schema(self):
|
||
"""确保业务 schema 包含后续新增字段（兼容已存在表）。"""
|
||
self._check_initialized()
|
||
stmts = [
|
||
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS tool_dependencies JSONB DEFAULT '[]'::jsonb",
|
||
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS mcp_dependencies JSONB DEFAULT '[]'::jsonb",
|
||
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS skill_dependencies JSONB DEFAULT '[]'::jsonb",
|
||
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS version VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS is_builtin BOOLEAN NOT NULL DEFAULT FALSE",
|
||
"ALTER TABLE IF EXISTS skills ADD COLUMN IF NOT EXISTS content_hash VARCHAR(128)",
|
||
"ALTER TABLE IF EXISTS subagents ADD COLUMN IF NOT EXISTS enabled BOOLEAN NOT NULL DEFAULT TRUE",
|
||
"ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS is_pinned BOOLEAN NOT NULL DEFAULT FALSE",
|
||
"ALTER TABLE IF EXISTS mcp_servers ADD COLUMN IF NOT EXISTS env JSONB",
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS model_providers (
|
||
id SERIAL PRIMARY KEY,
|
||
provider_id VARCHAR(100) NOT NULL UNIQUE,
|
||
display_name VARCHAR(100) NOT NULL,
|
||
provider_type VARCHAR(32) NOT NULL DEFAULT 'openai',
|
||
default_protocol VARCHAR(64),
|
||
base_url VARCHAR(500) NOT NULL,
|
||
embedding_base_url VARCHAR(500),
|
||
rerank_base_url VARCHAR(500),
|
||
models_endpoint VARCHAR(200),
|
||
embedding_models_endpoint VARCHAR(200),
|
||
rerank_models_endpoint VARCHAR(200),
|
||
api_key_env VARCHAR(128),
|
||
api_key VARCHAR(500),
|
||
capabilities JSONB NOT NULL DEFAULT '[]'::jsonb,
|
||
enabled_models JSONB NOT NULL DEFAULT '[]'::jsonb,
|
||
headers_json JSONB,
|
||
extra_json JSONB,
|
||
is_enabled BOOLEAN NOT NULL DEFAULT TRUE,
|
||
is_builtin BOOLEAN NOT NULL DEFAULT FALSE,
|
||
created_by VARCHAR(100),
|
||
updated_by VARCHAR(100),
|
||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||
updated_at TIMESTAMPTZ DEFAULT NOW()
|
||
)
|
||
""",
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS agent_runs (
|
||
id VARCHAR(64) PRIMARY KEY,
|
||
thread_id VARCHAR(64) NOT NULL,
|
||
agent_id VARCHAR(64) NOT NULL,
|
||
user_id VARCHAR(64) NOT NULL,
|
||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||
request_id VARCHAR(64) NOT NULL UNIQUE,
|
||
input_payload JSONB NOT NULL DEFAULT '{}'::jsonb,
|
||
error_type VARCHAR(64),
|
||
error_message TEXT,
|
||
started_at TIMESTAMPTZ,
|
||
finished_at TIMESTAMPTZ,
|
||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||
)
|
||
""",
|
||
"CREATE INDEX IF NOT EXISTS idx_agent_runs_user_created ON agent_runs(user_id, created_at DESC)",
|
||
"CREATE INDEX IF NOT EXISTS idx_agent_runs_thread_created ON agent_runs(thread_id, created_at DESC)",
|
||
"CREATE INDEX IF NOT EXISTS idx_agent_runs_status_updated ON agent_runs(status, updated_at)",
|
||
"CREATE INDEX IF NOT EXISTS ix_conversations_is_pinned ON conversations(is_pinned)",
|
||
"CREATE UNIQUE INDEX IF NOT EXISTS ix_model_providers_provider_id ON model_providers(provider_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_model_providers_is_enabled ON model_providers(is_enabled)",
|
||
# conversations 表扩展渠道字段
|
||
"ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS channel_type VARCHAR(20)",
|
||
"ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS channel_session_key VARCHAR(128)",
|
||
"CREATE INDEX IF NOT EXISTS ix_conversations_channel_type ON conversations(channel_type)",
|
||
"CREATE INDEX IF NOT EXISTS ix_conversations_channel_session_key ON conversations(channel_session_key)",
|
||
"""
|
||
CREATE UNIQUE INDEX IF NOT EXISTS uq_conversation_channel_session
|
||
ON conversations (channel_type, channel_session_key, agent_id)
|
||
WHERE channel_type IS NOT NULL
|
||
""",
|
||
# channel_bindings 表
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_bindings (
|
||
id SERIAL PRIMARY KEY,
|
||
channel_type VARCHAR(20) NOT NULL,
|
||
account_id VARCHAR(64) NOT NULL DEFAULT '',
|
||
group_id VARCHAR(128) NOT NULL DEFAULT '',
|
||
agent_config_id INTEGER NOT NULL REFERENCES agent_configs(id),
|
||
is_enabled BOOLEAN NOT NULL DEFAULT TRUE,
|
||
created_by VARCHAR(64),
|
||
updated_by VARCHAR(64),
|
||
created_at TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
||
updated_at TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
||
is_deleted INTEGER NOT NULL DEFAULT 0,
|
||
deleted_at TIMESTAMP,
|
||
CONSTRAINT uq_channel_binding UNIQUE (channel_type, account_id, group_id)
|
||
)
|
||
""",
|
||
# channel_bindings 补齐旧表缺失列
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS is_enabled BOOLEAN NOT NULL DEFAULT TRUE",
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS created_by VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS updated_by VARCHAR(64)",
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS updated_at "
|
||
"TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC')",
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS is_deleted INTEGER NOT NULL DEFAULT 0",
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMP",
|
||
"ALTER TABLE IF EXISTS channel_bindings ADD COLUMN IF NOT EXISTS session_key_strategy VARCHAR(20) NOT NULL DEFAULT 'auto'",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_bindings_channel_type ON channel_bindings(channel_type)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_bindings_is_enabled ON channel_bindings(is_enabled)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_bindings_is_deleted ON channel_bindings(is_deleted)",
|
||
# channel_outbox 表
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_outbox (
|
||
id SERIAL PRIMARY KEY,
|
||
message_id VARCHAR(64) NOT NULL,
|
||
session_id VARCHAR(64) NOT NULL,
|
||
channel_type VARCHAR(20) NOT NULL,
|
||
content TEXT NOT NULL,
|
||
status VARCHAR(20) NOT NULL DEFAULT 'pending',
|
||
retry_count INTEGER NOT NULL DEFAULT 0,
|
||
max_retries INTEGER NOT NULL DEFAULT 5,
|
||
next_retry_at TIMESTAMP,
|
||
last_error TEXT,
|
||
trace_id VARCHAR(64),
|
||
extra_metadata JSONB,
|
||
created_by VARCHAR(64),
|
||
updated_by VARCHAR(64),
|
||
created_at TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
||
updated_at TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
||
is_deleted INTEGER NOT NULL DEFAULT 0,
|
||
deleted_at TIMESTAMP
|
||
)
|
||
""",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_message_id ON channel_outbox(message_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_session_id ON channel_outbox(session_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_channel_type ON channel_outbox(channel_type)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_status ON channel_outbox(status)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_next_retry_at ON channel_outbox(next_retry_at)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_trace_id ON channel_outbox(trace_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_outbox_is_deleted ON channel_outbox(is_deleted)",
|
||
# channel_message_logs 表
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS channel_message_logs (
|
||
id SERIAL PRIMARY KEY,
|
||
trace_id VARCHAR(64) NOT NULL,
|
||
message_id VARCHAR(64) NOT NULL,
|
||
channel_type VARCHAR(20) NOT NULL,
|
||
conversation_id INTEGER,
|
||
session_id VARCHAR(64),
|
||
direction VARCHAR(10) NOT NULL,
|
||
sender_id VARCHAR(64),
|
||
content_summary VARCHAR(500),
|
||
agent_config_id INTEGER,
|
||
status VARCHAR(20) NOT NULL DEFAULT 'received',
|
||
pipeline_result VARCHAR(20),
|
||
abort_reason VARCHAR(64),
|
||
worker_result VARCHAR(20),
|
||
error_message TEXT,
|
||
processing_time_ms INTEGER,
|
||
extra_metadata JSONB,
|
||
created_by VARCHAR(64),
|
||
updated_by VARCHAR(64),
|
||
created_at TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
||
updated_at TIMESTAMP NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
|
||
is_deleted INTEGER NOT NULL DEFAULT 0,
|
||
deleted_at TIMESTAMP
|
||
)
|
||
""",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_trace_id ON channel_message_logs(trace_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_message_id ON channel_message_logs(message_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_channel_type ON channel_message_logs(channel_type)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_conversation_id "
|
||
"ON channel_message_logs(conversation_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_session_id ON channel_message_logs(session_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_direction ON channel_message_logs(direction)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_agent_config_id "
|
||
"ON channel_message_logs(agent_config_id)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_status ON channel_message_logs(status)",
|
||
"CREATE INDEX IF NOT EXISTS ix_channel_message_logs_is_deleted ON channel_message_logs(is_deleted)",
|
||
]
|
||
async with self.async_engine.begin() as conn:
|
||
for stmt in stmts:
|
||
await conn.execute(text(stmt))
|
||
|
||
@property
|
||
def is_initialized(self) -> bool:
|
||
return self._initialized
|
||
|
||
@property
|
||
def is_postgresql(self) -> bool:
|
||
"""检查是否是 PostgreSQL 数据库"""
|
||
if not self._initialized:
|
||
return False
|
||
return self.async_engine.dialect.name == "postgresql"
|
||
|
||
async def get_async_session(self) -> AsyncSession:
|
||
"""获取异步数据库会话"""
|
||
self.initialize() # 确保已初始化
|
||
return self.AsyncSession()
|
||
|
||
@asynccontextmanager
|
||
async def get_async_session_context(self):
|
||
"""获取异步数据库会话的上下文管理器"""
|
||
self.initialize() # 确保已初始化
|
||
session = self.AsyncSession()
|
||
try:
|
||
yield session
|
||
await session.commit()
|
||
except Exception as e:
|
||
await session.rollback()
|
||
logger.error(f"PostgreSQL async operation failed: {e}")
|
||
raise
|
||
finally:
|
||
await session.close()
|
||
|
||
async def close(self):
|
||
"""关闭引擎"""
|
||
if self.async_engine:
|
||
await self.async_engine.dispose()
|
||
|
||
if self.langgraph_pool:
|
||
await self.langgraph_pool.close()
|
||
|
||
async def async_check_first_run(self):
|
||
"""检查是否首次运行（异步版本）- 检查用户表是否有数据"""
|
||
from sqlalchemy import func, select
|
||
|
||
self._check_initialized()
|
||
async with self.get_async_session_context() as session:
|
||
from yuxi.storage.postgres.models_business import User
|
||
|
||
result = await session.execute(select(func.count(User.id)))
|
||
count = result.scalar()
|
||
return count == 0
|
||
|
||
async def execute(self, statement):
|
||
"""直接执行 SQL 语句（用于迁移脚本）"""
|
||
self._check_initialized()
|
||
async with self.get_async_session_context() as session:
|
||
return await session.execute(statement)
|
||
|
||
async def add(self, instance):
|
||
"""添加实例到会话（用于迁移脚本）"""
|
||
self._check_initialized()
|
||
async with self.get_async_session_context() as session:
|
||
session.add(instance)
|
||
|
||
async def commit(self):
|
||
"""提交当前会话"""
|
||
self._check_initialized()
|
||
async with self.get_async_session_context():
|
||
pass # commit is automatic in context manager
|
||
|
||
|
||
# 创建全局 PostgreSQL 管理器实例
|
||
pg_manager = PostgresManager()
|