247 lines
9.1 KiB
Python
247 lines
9.1 KiB
Python
|
"""
|
|||
|
|
异步 PostgreSQL 存储实现 - 用于生产环境。
|
|||
|
|
|
|||
|
|
使用 asyncpg 实现真正的异步 PostgreSQL 文档存储，支持高并发访问。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
from typing import List, Dict, Any, Optional, Iterator, Tuple, Sequence
|
|||
|
|
|
|||
|
|
from langchain_core.documents import Document
|
|||
|
|
from langchain_core.stores import BaseStore
|
|||
|
|
|
|||
|
|
import asyncpg
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
class PostgresDocStore(BaseStore[str, Any]):
|
|||
|
|
"""
|
|||
|
|
异步 PostgreSQL 文档存储实现。
|
|||
|
|
|
|||
|
|
使用 asyncpg 作为异步 PostgreSQL 客户端，支持：
|
|||
|
|
- 真正的异步操作
|
|||
|
|
- 连接池管理
|
|||
|
|
- 自动表创建
|
|||
|
|
- 批量操作（amget/amset/amdelete）
|
|||
|
|
- JSONB 数据存储
|
|||
|
|
- 并发控制
|
|||
|
|
|
|||
|
|
适用于生产环境，提供高性能的异步数据持久化。
|
|||
|
|
|
|||
|
|
Attributes:
|
|||
|
|
dsn: PostgreSQL 连接字符串
|
|||
|
|
table_name: 存储表名，默认为 "parent_documents"
|
|||
|
|
_pool: asyncpg 连接池实例
|
|||
|
|
_semaphore: 控制并发数的信号量（可选）
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
connection_string: str,
|
|||
|
|
table_name: str = "parent_documents",
|
|||
|
|
pool_config: Optional[Dict[str, Any]] = None,
|
|||
|
|
max_concurrency: Optional[int] = None
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
初始化异步 PostgreSQL 文档存储。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
connection_string: PostgreSQL 连接 URL，格式：
|
|||
|
|
"postgresql://user:password@host:port/database?sslmode=disable"
|
|||
|
|
table_name: 存储表名，默认为 "parent_documents"
|
|||
|
|
pool_config: 连接池配置字典，包含：
|
|||
|
|
- min_size: 最小连接数（默认 2）
|
|||
|
|
- max_size: 最大连接数（默认 10）
|
|||
|
|
max_concurrency: 最大并发操作数，如果为 None 则不限制
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ImportError: 未安装 asyncpg 时抛出
|
|||
|
|
|
|||
|
|
Example:
|
|||
|
|
>>> store = PostgresDocStore(
|
|||
|
|
... "postgresql://user:pass@localhost:5432/mydb",
|
|||
|
|
... table_name="parent_docs",
|
|||
|
|
... pool_config={"min_size": 5, "max_size": 20},
|
|||
|
|
... max_concurrency=10
|
|||
|
|
... )
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
self.dsn = connection_string
|
|||
|
|
self.table_name = table_name
|
|||
|
|
self._pool: Optional["asyncpg.Pool"] = None
|
|||
|
|
self._pool_config = pool_config or {}
|
|||
|
|
|
|||
|
|
# 并发控制信号量
|
|||
|
|
self._semaphore = None
|
|||
|
|
if max_concurrency is not None and max_concurrency > 0:
|
|||
|
|
self._semaphore = asyncio.Semaphore(max_concurrency)
|
|||
|
|
|
|||
|
|
# 注意：连接池的异步初始化延迟到第一次使用时
|
|||
|
|
# 表结构创建也延迟到第一次操作时
|
|||
|
|
|
|||
|
|
async def _get_pool(self):
|
|||
|
|
"""获取或创建 asyncpg 连接池。"""
|
|||
|
|
if self._pool is None:
|
|||
|
|
import asyncpg
|
|||
|
|
min_size = self._pool_config.get("min_size", 2)
|
|||
|
|
max_size = self._pool_config.get("max_size", 10)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
self._pool = await asyncpg.create_pool(
|
|||
|
|
dsn=self.dsn,
|
|||
|
|
min_size=min_size,
|
|||
|
|
max_size=max_size
|
|||
|
|
)
|
|||
|
|
logger.info(f"PostgreSQL 异步连接池已创建: {self.table_name}")
|
|||
|
|
|
|||
|
|
# 初始化表结构
|
|||
|
|
await self._create_table()
|
|||
|
|
except Exception as e:
|
|||
|
|
raise RuntimeError(f"PostgreSQL 异步连接池创建失败: {e}")
|
|||
|
|
|
|||
|
|
return self._pool
|
|||
|
|
|
|||
|
|
async def _create_table(self):
|
|||
|
|
"""创建存储表（如果不存在）。"""
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.transaction():
|
|||
|
|
await conn.execute(f"""
|
|||
|
|
CREATE TABLE IF NOT EXISTS {self.table_name} (
|
|||
|
|
key TEXT PRIMARY KEY,
|
|||
|
|
value JSONB NOT NULL,
|
|||
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|||
|
|
)
|
|||
|
|
""")
|
|||
|
|
logger.info(f"表 {self.table_name} 已就绪")
|
|||
|
|
|
|||
|
|
async def _with_concurrency_control(self, coro):
|
|||
|
|
"""使用信号量控制并发执行。"""
|
|||
|
|
if self._semaphore is None:
|
|||
|
|
return await coro
|
|||
|
|
async with self._semaphore:
|
|||
|
|
return await coro
|
|||
|
|
|
|||
|
|
# --- 同步方法（保持兼容性，但功能有限）---
|
|||
|
|
|
|||
|
|
def mget(self, keys: Sequence[str]) -> List[Optional[Any]]:
|
|||
|
|
"""不支持同步操作，请使用异步 amget 方法。"""
|
|||
|
|
raise NotImplementedError("不支持同步操作，请使用异步 amget 方法。")
|
|||
|
|
|
|||
|
|
def mset(self, key_value_pairs: Sequence[Tuple[str, Any]]) -> None:
|
|||
|
|
"""不支持同步操作，请使用异步 amset 方法。"""
|
|||
|
|
raise NotImplementedError("不支持同步操作，请使用异步 amset 方法。")
|
|||
|
|
|
|||
|
|
def mdelete(self, keys: Sequence[str]) -> None:
|
|||
|
|
"""不支持同步操作，请使用异步 amdelete 方法。"""
|
|||
|
|
raise NotImplementedError("不支持同步操作，请使用异步 amdelete 方法。")
|
|||
|
|
|
|||
|
|
def yield_keys(self, *, prefix: str | None = None) -> Iterator[str]:
|
|||
|
|
"""不支持同步操作，请使用异步 ayield_keys 方法。"""
|
|||
|
|
raise NotImplementedError("不支持同步操作，请使用异步 ayield_keys 方法。")
|
|||
|
|
|
|||
|
|
# --- 异步方法（真正的实现）---
|
|||
|
|
|
|||
|
|
async def amget(self, keys: Sequence[str]) -> List[Optional[Any]]:
|
|||
|
|
"""异步批量获取文档。"""
|
|||
|
|
if not keys:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
async def _amget():
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
rows = await conn.fetch(
|
|||
|
|
f"SELECT key, value FROM {self.table_name} WHERE key = ANY($1)",
|
|||
|
|
keys
|
|||
|
|
)
|
|||
|
|
result_map = {}
|
|||
|
|
for row in rows:
|
|||
|
|
val = row['value']
|
|||
|
|
if isinstance(val, str):
|
|||
|
|
val = json.loads(val)
|
|||
|
|
if isinstance(val, dict) and 'page_content' in val:
|
|||
|
|
result_map[row['key']] = Document(**val)
|
|||
|
|
else:
|
|||
|
|
result_map[row['key']] = val
|
|||
|
|
return [result_map.get(key) for key in keys]
|
|||
|
|
|
|||
|
|
return await self._with_concurrency_control(_amget())
|
|||
|
|
|
|||
|
|
async def amset(self, key_value_pairs: Sequence[Tuple[str, Any]]) -> None:
|
|||
|
|
"""异步批量设置文档。"""
|
|||
|
|
if not key_value_pairs:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
async def _amset():
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.transaction():
|
|||
|
|
await conn.executemany(
|
|||
|
|
f"""
|
|||
|
|
INSERT INTO {self.table_name} (key, value)
|
|||
|
|
VALUES ($1, $2)
|
|||
|
|
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value
|
|||
|
|
""",
|
|||
|
|
[
|
|||
|
|
(k, json.dumps(v.dict() if isinstance(v, Document) else v, ensure_ascii=False))
|
|||
|
|
for k, v in key_value_pairs
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
logger.debug(f"已异步批量设置 {len(key_value_pairs)} 个文档")
|
|||
|
|
|
|||
|
|
await self._with_concurrency_control(_amset())
|
|||
|
|
|
|||
|
|
async def amdelete(self, keys: Sequence[str]) -> None:
|
|||
|
|
"""异步批量删除文档。"""
|
|||
|
|
if not keys:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
async def _amdelete():
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
async with conn.transaction():
|
|||
|
|
await conn.execute(
|
|||
|
|
f"DELETE FROM {self.table_name} WHERE key = ANY($1)",
|
|||
|
|
keys
|
|||
|
|
)
|
|||
|
|
logger.debug(f"已异步批量删除 {len(keys)} 个文档")
|
|||
|
|
|
|||
|
|
await self._with_concurrency_control(_amdelete())
|
|||
|
|
|
|||
|
|
async def ayield_keys(self, *, prefix: str | None = None) -> Iterator[str]:
|
|||
|
|
"""异步迭代所有键。
|
|||
|
|
|
|||
|
|
注意：这是一个异步生成器，需要使用 async for 迭代。
|
|||
|
|
"""
|
|||
|
|
pool = await self._get_pool()
|
|||
|
|
async with pool.acquire() as conn:
|
|||
|
|
if prefix:
|
|||
|
|
rows = await conn.fetch(
|
|||
|
|
f"SELECT key FROM {self.table_name} WHERE key LIKE $1 ORDER BY key",
|
|||
|
|
f"{prefix}%"
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
rows = await conn.fetch(
|
|||
|
|
f"SELECT key FROM {self.table_name} ORDER BY key"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
for row in rows:
|
|||
|
|
yield row['key']
|
|||
|
|
|
|||
|
|
async def aclose(self) -> None:
|
|||
|
|
"""异步关闭连接池，释放资源。"""
|
|||
|
|
if self._pool:
|
|||
|
|
await self._pool.close()
|
|||
|
|
self._pool = None
|
|||
|
|
logger.info("PostgreSQL 异步连接池已关闭")
|
|||
|
|
|
|||
|
|
def close(self) -> None:
|
|||
|
|
"""同步关闭连接池（功能有限）。
|
|||
|
|
|
|||
|
|
注意：在异步环境中，请使用 aclose 方法。
|
|||
|
|
"""
|
|||
|
|
pass
|