425f3c5bb4
Convert all Chinese comments, docstrings, logger/print output, HTTPException detail messages, and API response messages to English across the entire codebase. Functional zh/ja localized strings (e.g. prompt templates, timezone display names, date formats) are preserved as-is. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
"""
|
|
Global PostgreSQL checkpointer manager.
|
|
Uses the shared database connection pool.
|
|
"""
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from psycopg_pool import AsyncConnectionPool
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
|
|
logger = logging.getLogger('app')
|
|
|
|
|
|
class CheckpointerManager:
|
|
"""
|
|
Global checkpointer manager using the shared PostgreSQL connection pool.
|
|
|
|
Main features:
|
|
1. Uses the shared connection pool from DBPoolManager
|
|
2. AsyncPostgresSaver natively supports connection pools and acquires/releases connections automatically
|
|
3. No manual return is required, avoiding long-running requests occupying connections
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._pool: Optional[AsyncConnectionPool] = None
|
|
self._initialized = False
|
|
self._closed = False
|
|
|
|
async def initialize(self, pool: AsyncConnectionPool) -> None:
|
|
"""Initialize the checkpointer using an externally provided connection pool
|
|
|
|
Args:
|
|
pool: AsyncConnectionPool instance (from DBPoolManager)
|
|
"""
|
|
if self._initialized:
|
|
return
|
|
|
|
self._pool = pool
|
|
|
|
logger.info("Initializing PostgreSQL checkpointer (using shared connection pool)...")
|
|
|
|
try:
|
|
# Create the table schema (autocommit is required for CREATE INDEX CONCURRENTLY)
|
|
async with self._pool.connection() as conn:
|
|
await conn.set_autocommit(True)
|
|
checkpointer = AsyncPostgresSaver(conn=conn)
|
|
await checkpointer.setup()
|
|
|
|
self._initialized = True
|
|
logger.info("PostgreSQL checkpointer initialized successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize PostgreSQL checkpointer: {e}")
|
|
raise
|
|
|
|
@property
|
|
def checkpointer(self) -> AsyncPostgresSaver:
|
|
"""Get the global AsyncPostgresSaver instance."""
|
|
if self._closed:
|
|
raise RuntimeError("CheckpointerManager is closed")
|
|
|
|
if not self._initialized:
|
|
raise RuntimeError("CheckpointerManager not initialized, call initialize() first")
|
|
|
|
return AsyncPostgresSaver(conn=self._pool)
|
|
|
|
@property
|
|
def pool(self) -> AsyncConnectionPool:
|
|
"""Get the connection pool for sharing with other managers."""
|
|
if self._closed:
|
|
raise RuntimeError("CheckpointerManager is closed")
|
|
|
|
if not self._initialized:
|
|
raise RuntimeError("CheckpointerManager not initialized, call initialize() first")
|
|
|
|
return self._pool
|
|
|
|
async def close(self) -> None:
|
|
"""Close the checkpointer manager without closing the pool managed by DBPoolManager."""
|
|
if self._closed:
|
|
return
|
|
|
|
logger.info("Closing CheckpointerManager...")
|
|
self._closed = True
|
|
self._initialized = False
|
|
logger.info("CheckpointerManager closed (pool managed by DBPoolManager)")
|
|
|
|
|
|
# Global singleton
|
|
_global_manager: Optional[CheckpointerManager] = None
|
|
|
|
|
|
def get_checkpointer_manager() -> CheckpointerManager:
|
|
"""Get the global CheckpointerManager singleton."""
|
|
global _global_manager
|
|
if _global_manager is None:
|
|
_global_manager = CheckpointerManager()
|
|
return _global_manager
|
|
|
|
|
|
async def init_global_checkpointer(pool: AsyncConnectionPool) -> None:
|
|
"""Initialize the global checkpointer manager
|
|
|
|
Args:
|
|
pool: AsyncConnectionPool instance (from DBPoolManager)
|
|
"""
|
|
manager = get_checkpointer_manager()
|
|
await manager.initialize(pool)
|
|
|
|
|
|
async def close_global_checkpointer() -> None:
|
|
"""Close the global checkpointer manager."""
|
|
global _global_manager
|
|
if _global_manager is not None:
|
|
await _global_manager.close()
|