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