yuxuanhui
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
"""Bounded PostgreSQL pool ownership for selection adapters."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
from typing import Any, Protocol, cast
|
|
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
|
|
class SelectionConnectionPool(Protocol):
|
|
"""Small pool surface shared by psycopg and unit-test fakes."""
|
|
|
|
def open(self, *, wait: bool = True) -> None: ...
|
|
|
|
def close(self) -> None: ...
|
|
|
|
def connection(self) -> Any: ...
|
|
|
|
|
|
class SelectionPostgresPool:
|
|
"""Own one bounded PostgreSQL pool for the selection read/write adapters.
|
|
|
|
The owner opens lazily on the first borrowed connection, which keeps app
|
|
construction cheap while still making the pool lifetime process-scoped in
|
|
the HTTP composition layer. A fake pool can be injected for unit tests.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
database_url: str,
|
|
*,
|
|
max_connections: int,
|
|
pool: SelectionConnectionPool | None = None,
|
|
) -> None:
|
|
"""Create a bounded pool owner with an optional injected pool."""
|
|
|
|
if max_connections < 1:
|
|
raise ValueError("max_connections must be at least 1")
|
|
self.database_url = database_url
|
|
self.max_connections = max_connections
|
|
self.pool: SelectionConnectionPool = pool or cast(
|
|
SelectionConnectionPool,
|
|
ConnectionPool(
|
|
conninfo=database_url,
|
|
min_size=1,
|
|
max_size=max_connections,
|
|
open=False,
|
|
),
|
|
)
|
|
self._pool_open = False
|
|
self._pool_state_lock = threading.Lock()
|
|
|
|
def open(self) -> None:
|
|
"""Open the underlying pool once and wait for its minimum connection."""
|
|
|
|
with self._pool_state_lock:
|
|
if self._pool_open:
|
|
return
|
|
if bool(getattr(self.pool, "_opened", False)):
|
|
self._pool_open = True
|
|
return
|
|
self.pool.open(wait=True)
|
|
self._pool_open = True
|
|
|
|
def close(self) -> None:
|
|
"""Close the pool after borrowed connections have been returned."""
|
|
|
|
with self._pool_state_lock:
|
|
if self._pool_open or bool(getattr(self.pool, "_opened", False)):
|
|
self.pool.close()
|
|
self._pool_open = False
|
|
|
|
@contextmanager
|
|
def connection(self) -> Generator[Any, None, None]:
|
|
"""Borrow one connection and return it to the bounded pool."""
|
|
|
|
self.open()
|
|
with self.pool.connection() as connection:
|
|
yield connection
|
|
|
|
def __enter__(self) -> SelectionPostgresPool:
|
|
"""Open and return this resource owner."""
|
|
|
|
self.open()
|
|
return self
|
|
|
|
def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
|
|
"""Release the owned pool at the end of a context."""
|
|
|
|
self.close()
|
|
|
|
|
|
__all__ = ["SelectionConnectionPool", "SelectionPostgresPool"]
|