memuron 0.1.1__py3-none-any.whl

memuron/persistence/__init__.py

@@ -0,0 +1 @@
+"""Database, identity-store, and projection persistence helpers."""

memuron/persistence/db_pool.py

@@ -0,0 +1,57 @@
+"""PostgreSQL connection pooling for memuron hot paths."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+from psycopg.rows import dict_row
+from psycopg_pool import ConnectionPool
+
+logger = logging.getLogger(__name__)
+
+_POOL_ATTR = "_memuron_connection_pool"
+_ORIGINAL_CONNECT_ATTR = "_memuron_original_connect"
+
+
+def install_postgres_pool(store: object, *, min_size: int = 1, max_size: int = 8) -> bool:
+    """Replace per-query connects with a shared pool on Postgres stores."""
+    database_url = getattr(store, "database_url", None)
+    original_connect = getattr(store, "_connect", None)
+    if not database_url or original_connect is None:
+        return False
+    if getattr(store, _POOL_ATTR, None) is not None:
+        return True
+
+    pool = ConnectionPool(
+        conninfo=database_url,
+        min_size=min_size,
+        max_size=max_size,
+        kwargs={"row_factory": dict_row},
+        open=True,
+    )
+    setattr(store, _ORIGINAL_CONNECT_ATTR, original_connect)
+
+    @contextmanager
+    def pooled_connect() -> Iterator[Any]:
+        with pool.connection() as conn:
+            yield conn
+
+    store._connect = pooled_connect  # type: ignore[method-assign]
+    setattr(store, _POOL_ATTR, pool)
+    logger.info("Postgres connection pool ready (min=%s max=%s)", min_size, max_size)
+    return True
+
+
+def close_postgres_pool(store: object) -> None:
+    pool = getattr(store, _POOL_ATTR, None)
+    if pool is None:
+        return
+    pool.close()
+    original_connect = getattr(store, _ORIGINAL_CONNECT_ATTR, None)
+    if original_connect is not None:
+        store._connect = original_connect  # type: ignore[method-assign]
+    delattr(store, _POOL_ATTR)
+    if hasattr(store, _ORIGINAL_CONNECT_ATTR):
+        delattr(store, _ORIGINAL_CONNECT_ATTR)