PostgresPoolManager -- asyncpg connection pool wrapper.
Provides a reusable async context manager around asyncpg's Pool.
Shared by both PostgresGraphBackend and PostgresVectorBackend.
Optional dependency: asyncpg must be installed.
pgvector codec registration (v7.3, Item 13.3)
Every physical connection created by the pool runs the pgvector asyncpg
codec registration in its init=... callback. Once registered,
asyncpg returns vector columns as pgvector.Vector instances
(iterable of floats) instead of the slower text form "[0.1,0.2]".
The graph backend's _decode_embedding accepts BOTH shapes — text
form AND iterable — so this PR can land without breaking the 7.2.1
hotfix. When pgvector is not importable (degraded mode) the init
callback logs once and the text-form path keeps working.
PostgresPoolManager
PostgresPoolManager(dsn: str, min_size: int = 2, max_size: int = 10)
Async connection pool manager wrapping asyncpg.Pool.
Usage::
pool = PostgresPoolManager(dsn="postgresql://user:pw@host/db")
await pool.open()
async with pool as conn:
await conn.execute("SELECT 1")
await pool.close()
Or as a context manager for the pool itself::
async with PostgresPoolManager(dsn=...) as pool:
conn = await pool.acquire()
try:
await conn.execute("SELECT 1")
finally:
await pool.release(conn)
Source code in src/symfonic/memory/backends/pool.py
| def __init__(
self,
dsn: str,
min_size: int = 2,
max_size: int = 10,
) -> None:
if not _ASYNCPG_AVAILABLE:
raise ImportError(
"PostgresPoolManager requires asyncpg. "
"Install it with: pip install symfonic-core[postgres]"
)
self._dsn = dsn
self._min_size = min_size
self._max_size = max_size
self._pool: Any = None # asyncpg.Pool at runtime
self._guards: tuple[Any, ...] = () # CORE-171 erase anchor: postgres_erase_incarnation
|
dsn
property
Return the configured connection string (read-only).
Used by the v7.1.1 PostgresCheckpointerFactory to build a
separate psycopg pool against the same database. The two pools
intentionally do not share connections (asyncpg vs. psycopg
protocols are incompatible inside langgraph-checkpoint-postgres).
raw_pool
property
Expose the underlying asyncpg.Pool for direct use (advanced).
acquire
async
acquire(*, timeout: float | None = None) -> Any
Acquire a connection from the pool.
Caller is responsible for releasing via :meth:release.
Inside :meth:transaction this hands back that transaction's own
connection, so every backend on this pool joins it without knowing it
exists (:mod:.pool_transaction).
Source code in src/symfonic/memory/backends/pool.py
| async def acquire(self, *, timeout: float | None = None) -> Any:
"""Acquire a connection from the pool.
Caller is responsible for releasing via :meth:`release`.
Inside :meth:`transaction` this hands back that transaction's own
connection, so every backend on this pool joins it without knowing it
exists (:mod:`.pool_transaction`).
"""
if foreign_bound_transaction(self):
raise RuntimeError("inherited PostgreSQL transaction belongs to another task")
bound = bound_connection(self)
if bound is not None:
for guard in self._guards:
await guard(bound)
return bound
if self._pool is None:
raise RuntimeError(
"Pool is not open. Call await pool.open() first "
"or use 'async with pool' context manager."
)
if timeout is None:
conn = await self._pool.acquire()
else:
conn = await self._pool.acquire(timeout=timeout)
if self._guards:
try:
for guard in self._guards: # every checkout, before any caller
await guard(conn)
except BaseException:
await self._pool.release(conn)
raise
return conn
|
anchor_graph
anchor_graph(eraser: Any) -> Any
CORE-171: guard this pool with eraser's anchor proof, or refuse native reads.
Source code in src/symfonic/memory/backends/pool.py
| def anchor_graph(self, eraser: Any) -> Any:
"""CORE-171: guard this pool with ``eraser``'s anchor proof, or refuse native reads."""
from symfonic.memory.backends.postgres_erase_incarnation import anchored_graph
return anchored_graph(self, eraser)
|
close
async
Gracefully close all connections in the pool.
Source code in src/symfonic/memory/backends/pool.py
| async def close(self) -> None:
"""Gracefully close all connections in the pool."""
if self._pool is not None:
await self._pool.close()
self._pool = None
|
guard_connections
guard_connections(guard: Any) -> None
Also run guard on every connection acquire hands out, before the caller.
Guards only add: a connection is handed out once every guard admitted it.
A placeholder guard (refuse: nothing proves this pool) yields to,
and is never added beside, a guard that proves every checkout.
Source code in src/symfonic/memory/backends/pool.py
| def guard_connections(self, guard: Any) -> None:
"""Also run ``guard`` on every connection ``acquire`` hands out, before the caller.
Guards only add: a connection is handed out once every guard admitted it.
A ``placeholder`` guard (refuse: nothing proves this pool) yields to,
and is never added beside, a guard that proves every checkout.
"""
domain = getattr(guard, "guard_domain", "native")
same_domain = tuple(
g for g in self._guards if getattr(g, "guard_domain", "native") == domain
)
proving = tuple(g for g in same_domain if not getattr(g, "placeholder", False))
if getattr(guard, "placeholder", False):
if not proving and guard not in same_domain:
self._guards += (guard,)
elif guard not in proving:
# One authority cannot substitute for an unrelated authority.
self._guards = tuple(
g
for g in self._guards
if getattr(g, "guard_domain", "native") != domain
or not getattr(g, "placeholder", False)
) + (guard,)
|
open
async
Create the underlying asyncpg connection pool.
Passes init=_register_pgvector_on_connection so every new
physical connection registers the pgvector binary codec before
the pool hands it out. See module docstring for the rationale.
Source code in src/symfonic/memory/backends/pool.py
| async def open(self) -> None:
"""Create the underlying asyncpg connection pool.
Passes ``init=_register_pgvector_on_connection`` so every new
physical connection registers the pgvector binary codec before
the pool hands it out. See module docstring for the rationale.
"""
if self._pool is not None:
return
self._pool = await asyncpg.create_pool(
self._dsn,
min_size=self._min_size,
max_size=self._max_size,
init=_register_pgvector_on_connection,
)
|
release
async
release(conn: Any, *, timeout: float | None = None) -> None
Return a connection to the pool.
A no-op for a connection :meth:transaction holds: releasing it
mid-way would hand a half-written batch to somebody else.
Source code in src/symfonic/memory/backends/pool.py
| async def release(self, conn: Any, *, timeout: float | None = None) -> None:
"""Return a connection to the pool.
A no-op for a connection :meth:`transaction` holds: releasing it
mid-way would hand a half-written batch to somebody else.
"""
if is_bound(conn):
return
if self._pool is not None:
if timeout is None:
await self._pool.release(conn)
return
task = asyncio.ensure_future(self._pool.release(conn, timeout=max(0, timeout)))
try:
done, _ = await asyncio.wait({task}, timeout=max(0, timeout))
if not done:
raise TimeoutError("PostgreSQL release was not acknowledged")
task.result()
except BaseException:
# A detached proxy may already have been released; never mask cancellation.
with suppress(Exception):
conn.terminate() # Discard is not evidence of transaction rollback.
task.cancel()
task.add_done_callback(
lambda finished: None if finished.cancelled() else finished.exception()
)
raise
|
transaction
Run a block so every operation on this pool commits or fails together.
Ambient rather than passed: :mod:.pool_transaction has the why.
Source code in src/symfonic/memory/backends/pool.py
| def transaction(self) -> Any:
"""Run a block so every operation on this pool commits or fails together.
Ambient rather than passed: :mod:`.pool_transaction` has the why.
"""
return run_in_transaction(self)
|
connection_has_pgvector_codec
connection_has_pgvector_codec(conn: Any) -> bool
Return True if conn had pgvector's codec successfully registered.
Set by :func:_register_pgvector_on_connection per physical
connection. Defaults to False when the connection has not been
seen by the init hook — e.g. on mock connections in unit tests, or
on a connection that was not created by a pool we control.
Source code in src/symfonic/memory/backends/pool_pgvector.py
| def connection_has_pgvector_codec(conn: Any) -> bool:
"""Return True if ``conn`` had pgvector's codec successfully registered.
Set by :func:`_register_pgvector_on_connection` per physical
connection. Defaults to ``False`` when the connection has not been
seen by the init hook — e.g. on mock connections in unit tests, or
on a connection that was not created by a pool we control.
"""
raw = _raw_connection(conn)
try:
return bool(_CODEC_REGISTRY.get(raw, False))
except TypeError: # raw is not hashable (mocks, etc.) — degrade safely
return False
|