ForcePilot/backend/package/yuxi/storage/postgres/manager.py
Kris fb754d82f1 feat(channel): 添加全渠道网关核心基础架构
新增渠道扩展框架的核心基础层,为所有渠道扩展提供统一的基础设施。

主要变更:
- 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: 日志配置更新
2026-05-21 12:08:03 +08:00

495 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""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 字符的 IDkb_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()