新增渠道扩展框架的核心基础层,为所有渠道扩展提供统一的基础设施。 主要变更: - channel/extensions/__init__.py: 渠道扩展包初始化与注册 - channel/extensions/base.py: 渠道插件基类,定义统一接口规范 - channel_service.py: 渠道服务核心业务逻辑 - models_channels.py: 渠道相关数据库模型 - 8 个 channel_*_repo.py: 渠道数据仓库层 - storage/postgres/manager.py: 数据库管理器,新增渠道表 - models_business.py: 业务模型更新 - config/app.py: 应用配置更新 - conversation_repository.py: 会话仓库更新 - oidc_service.py: OIDC 服务更新 - logging_config.py: 日志配置更新
495 lines
26 KiB
Python
495 lines
26 KiB
Python
"""PostgreSQL 数据库管理器 - 支持知识库和业务数据"""
|
||
|
||
import json
|
||
import os
|
||
from contextlib import asynccontextmanager
|
||
|
||
import yuxi.storage.postgres.models_channels as _ # 触发 SQLAlchemy Base.metadata 注册
|
||
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 server.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,
|
||
)
|
||
|
||
# 创建异步会话工厂
|
||
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)",
|
||
]
|
||
async with self.async_engine.begin() as conn:
|
||
for stmt in stmts:
|
||
await conn.execute(text(stmt))
|
||
|
||
async def ensure_channels_schema(self):
|
||
self._check_initialized()
|
||
stmts = [
|
||
"""CREATE TABLE IF NOT EXISTS channel_configs (
|
||
id VARCHAR(64) PRIMARY KEY,
|
||
channel_type VARCHAR(50) NOT NULL,
|
||
name VARCHAR(100) NOT NULL,
|
||
config JSONB NOT NULL DEFAULT '{}',
|
||
enabled BOOLEAN DEFAULT TRUE,
|
||
created_by VARCHAR(64),
|
||
updated_by VARCHAR(64),
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_configs_type
|
||
ON channel_configs (channel_type)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_configs_type_enabled
|
||
ON channel_configs (channel_type, enabled)""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_bindings (
|
||
id VARCHAR(64) PRIMARY KEY,
|
||
channel_type VARCHAR(50) NOT NULL,
|
||
channel_config_id VARCHAR(64) REFERENCES channel_configs(id) ON DELETE SET NULL,
|
||
account_id VARCHAR(100),
|
||
peer_kind VARCHAR(20),
|
||
peer_id VARCHAR(200),
|
||
agent_config_id INTEGER NOT NULL REFERENCES agent_configs(id) ON DELETE CASCADE,
|
||
priority INTEGER DEFAULT 0,
|
||
dm_scope VARCHAR(50) DEFAULT 'per-channel-peer',
|
||
created_by VARCHAR(64),
|
||
updated_by VARCHAR(64),
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_bindings_agent
|
||
ON channel_bindings (agent_config_id)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_bindings_lookup
|
||
ON channel_bindings (channel_type, account_id, peer_kind, peer_id)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_bindings_type
|
||
ON channel_bindings (channel_type)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_bindings_config
|
||
ON channel_bindings (channel_config_id)""",
|
||
"""ALTER TABLE channel_bindings
|
||
ADD COLUMN IF NOT EXISTS guild_id VARCHAR(100)""",
|
||
"""ALTER TABLE channel_bindings
|
||
ADD COLUMN IF NOT EXISTS team_id VARCHAR(100)""",
|
||
"""ALTER TABLE channel_bindings
|
||
ADD COLUMN IF NOT EXISTS roles JSONB DEFAULT '[]'""",
|
||
"""ALTER TABLE channel_bindings
|
||
ADD COLUMN IF NOT EXISTS session_dm_scope VARCHAR(50)""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_pairing_records (
|
||
id VARCHAR(64) PRIMARY KEY,
|
||
channel_type VARCHAR(50) NOT NULL,
|
||
account_id VARCHAR(100) NOT NULL DEFAULT 'default',
|
||
peer_id VARCHAR(200) NOT NULL,
|
||
agent_config_id INTEGER REFERENCES agent_configs(id) ON DELETE SET NULL,
|
||
pairing_code VARCHAR(8) NOT NULL,
|
||
pairing_token VARCHAR(64) NOT NULL,
|
||
status VARCHAR(20) DEFAULT 'pending',
|
||
expires_at TIMESTAMP WITH TIME ZONE NOT NULL,
|
||
paired_at TIMESTAMP WITH TIME ZONE,
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_pairing_records_lookup
|
||
ON channel_pairing_records (channel_type, account_id, peer_id, status)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_pairing_records_expiry
|
||
ON channel_pairing_records (expires_at)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_pairing_records_agent
|
||
ON channel_pairing_records (agent_config_id)""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_user_mappings (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
channel_id VARCHAR(32) NOT NULL,
|
||
channel_user_id VARCHAR(128) NOT NULL,
|
||
internal_user_id VARCHAR(64) NOT NULL,
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
CONSTRAINT uq_channel_user UNIQUE (channel_id, channel_user_id)
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_user_internal
|
||
ON channel_user_mappings (internal_user_id)""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_thread_mappings (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
channel_id VARCHAR(32) NOT NULL,
|
||
channel_chat_id VARCHAR(128) NOT NULL,
|
||
internal_user_id VARCHAR(64) NOT NULL,
|
||
thread_id VARCHAR(64) NOT NULL,
|
||
agent_id VARCHAR(64),
|
||
last_active_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
CONSTRAINT uq_channel_thread UNIQUE (channel_id, channel_chat_id, internal_user_id)
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_thread_thread
|
||
ON channel_thread_mappings (thread_id)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_thread_active
|
||
ON channel_thread_mappings (last_active_at)""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_msg_records (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
channel_id VARCHAR(32) NOT NULL,
|
||
channel_type VARCHAR(32) NOT NULL,
|
||
message_id VARCHAR(128) NOT NULL,
|
||
chat_id VARCHAR(128) NOT NULL,
|
||
chat_type VARCHAR(32) NOT NULL DEFAULT 'direct',
|
||
content_type VARCHAR(32) NOT NULL DEFAULT 'text',
|
||
sender_user_id VARCHAR(128) NOT NULL,
|
||
content_preview VARCHAR(500) NOT NULL,
|
||
reply_to_message_id VARCHAR(128),
|
||
agent_config_id INTEGER,
|
||
status VARCHAR(32) NOT NULL DEFAULT 'processing',
|
||
error_message VARCHAR(500),
|
||
response_time_ms INTEGER,
|
||
reply_message_id VARCHAR(128),
|
||
reply_content_preview VARCHAR(500),
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
replied_at TIMESTAMP WITH TIME ZONE,
|
||
extra_data JSONB
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_msg_records_channel_status
|
||
ON channel_msg_records (channel_id, status)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_msg_records_created
|
||
ON channel_msg_records (created_at)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_msg_records_chat
|
||
ON channel_msg_records (channel_id, chat_id)""",
|
||
"""ALTER TABLE IF EXISTS users
|
||
ADD COLUMN IF NOT EXISTS source VARCHAR(32) NOT NULL DEFAULT 'local'""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_device_identities (
|
||
id VARCHAR(64) PRIMARY KEY,
|
||
device_id VARCHAR(64) NOT NULL,
|
||
name VARCHAR(255) NOT NULL,
|
||
public_key_pem TEXT NOT NULL,
|
||
status VARCHAR(20) DEFAULT 'active',
|
||
created_by VARCHAR(64),
|
||
last_used_at TIMESTAMP WITH TIME ZONE,
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||
)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_device_identities_device_id
|
||
ON channel_device_identities (device_id)""",
|
||
"""CREATE INDEX IF NOT EXISTS idx_channel_device_identities_status
|
||
ON channel_device_identities (status)""",
|
||
"""CREATE TABLE IF NOT EXISTS channel_identity_links (
|
||
id VARCHAR(64) PRIMARY KEY,
|
||
identity VARCHAR(255) NOT NULL,
|
||
channel_type VARCHAR(64) NOT NULL,
|
||
peer_id VARCHAR(255) NOT NULL,
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
CONSTRAINT uq_identity_link_channel_peer UNIQUE (channel_type, peer_id)
|
||
)""",
|
||
"CREATE INDEX IF NOT EXISTS idx_identity_links_identity ON channel_identity_links(identity)",
|
||
"""CREATE TABLE IF NOT EXISTS pay_notifications (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
notification_id VARCHAR(128) NOT NULL,
|
||
mch_id VARCHAR(32) NOT NULL,
|
||
event_type VARCHAR(64) NOT NULL,
|
||
out_trade_no VARCHAR(64) DEFAULT '',
|
||
transaction_id VARCHAR(64) DEFAULT '',
|
||
resource_json JSONB DEFAULT '{}',
|
||
status VARCHAR(32) DEFAULT 'pending',
|
||
error_message TEXT,
|
||
retry_count INTEGER DEFAULT 0,
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
completed_at TIMESTAMP WITH TIME ZONE,
|
||
CONSTRAINT uk_pay_notification_id UNIQUE (notification_id)
|
||
)""",
|
||
"CREATE INDEX IF NOT EXISTS idx_pay_notifications_mch_id ON pay_notifications(mch_id)",
|
||
"CREATE INDEX IF NOT EXISTS idx_pay_notifications_out_trade_no ON pay_notifications(out_trade_no)",
|
||
"CREATE INDEX IF NOT EXISTS idx_pay_notifications_status ON pay_notifications(status)",
|
||
"CREATE INDEX IF NOT EXISTS idx_pay_notifications_created_at ON pay_notifications(created_at)",
|
||
]
|
||
|
||
async with self.async_engine.begin() as conn:
|
||
for stmt in stmts:
|
||
await conn.execute(text(stmt))
|
||
|
||
@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._check_initialized()
|
||
return self.AsyncSession()
|
||
|
||
@asynccontextmanager
|
||
async def get_async_session_context(self):
|
||
"""获取异步数据库会话的上下文管理器"""
|
||
self._check_initialized()
|
||
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()
|