#221: implement connection pool (#408)

Kraysent · web-flow · commit defdd1715aa0 · 2026-03-08T18:19:43.000+03:00
diff --git a/app/data/repositories/layer0/tables.py b/app/data/repositories/layer0/tables.py
@@ -104,7 +104,7 @@ def insert_raw_data(self, data: model.Layer0RawData) -> None:
             field_identifiers,
             placeholders,
         )
-        query_str = query.as_string(self._storage.get_connection())
+        query_str = self._storage.query_str(query)
 
         rows: list[list[Any]] = []
         for row in data.data.to_dict(orient="records"):
@@ -420,12 +420,11 @@ def update_column_metadata(self, table_name: str, column_description: model.Colu
 
         modification_query = "UPDATE layer0.tables SET modification_dt = now() WHERE id = %s"
 
-        with self.with_tx():
-            self._storage.exec(
-                "SELECT meta.setparams(%s, %s, %s, %s::json)",
-                params=[RAWDATA_SCHEMA, table_name, column_description.name, json.dumps(column_params)],
-            )
-            self._storage.exec(modification_query, params=[table_id])
+        self._storage.exec(
+            "SELECT meta.setparams(%s, %s, %s, %s::json)",
+            params=[RAWDATA_SCHEMA, table_name, column_description.name, json.dumps(column_params)],
+        )
+        self._storage.exec(modification_query, params=[table_id])
 
     def search_tables(
         self,
diff --git a/app/lib/storage/postgres/postgres_storage.py b/app/lib/storage/postgres/postgres_storage.py
@@ -1,3 +1,5 @@
+import threading
+import time
 from collections.abc import Sequence
 from typing import Any
 
@@ -6,6 +8,7 @@
 import structlog
 from psycopg import rows, sql
 from psycopg.types import enum, numeric
+from psycopg_pool import ConnectionPool
 
 from app.lib.storage import enums
 from app.lib.storage.postgres import config
@@ -33,7 +36,7 @@ def dump(self, obj: Any) -> bytes | bytearray | memoryview:
     (np.int64, NumpyIntDumper),
 ]
 
-DEFAULT_ENUMS = [
+DEFAULT_ENUMS: list[tuple[type[enum.Enum], str]] = [
     (enums.DataType, "common.datatype"),
     (enums.RecordCrossmatchStatus, "layer0.crossmatch_status"),
     (enums.RecordTriageStatus, "layer0.triage_status"),
@@ -43,89 +46,112 @@ def dump(self, obj: Any) -> bytes | bytearray | memoryview:
 class PgStorage:
     def __init__(self, cfg: config.PgStorageConfig, logger: structlog.stdlib.BoundLogger) -> None:
         self._config = cfg
-        self._connection: psycopg.Connection | None = None
+        self._pool: ConnectionPool | None = None
         self._logger = logger
+        self._local = threading.local()
+        self._extra_enums: list[tuple[type[enum.Enum], str]] = []
 
-    def connect(self) -> None:
-        self._connection = psycopg.connect(self._config.get_dsn(), row_factory=rows.dict_row, autocommit=True)
-        if self._connection is None:
-            raise InternalError("unable to create database connection")
+    def _configure_connection(self, conn: psycopg.Connection) -> None:
+        for python_type, dumper in DEFAULT_DUMPERS:
+            conn.adapters.register_dumper(python_type, dumper)
+        for enum_type, pg_type in DEFAULT_ENUMS + self._extra_enums:
+            type_info = enum.EnumInfo.fetch(conn, pg_type)
+            if type_info is None:
+                raise RuntimeError(f"Unable to find enum {pg_type} in DB")
+            enum.register_enum(
+                type_info,
+                conn,
+                enum_type,
+                mapping={m: m.value for m in enum_type},
+            )
 
+    def connect(self) -> None:
         self._logger.debug("connecting to Postgres", endpoint=self._config.endpoint, port=self._config.port)
+        self._pool = ConnectionPool(
+            self._config.get_dsn(),
+            min_size=2,
+            max_size=10,
+            kwargs={"row_factory": rows.dict_row, "autocommit": True},
+            configure=self._configure_connection,
+        )
 
-        for python_type, dumper in DEFAULT_DUMPERS:
-            self._connection.adapters.register_dumper(python_type, dumper)
+    def register_type(self, enum_type: type[enum.Enum], pg_type: str) -> None:
+        self._extra_enums.append((enum_type, pg_type))
 
-        for python_type, pg_type in DEFAULT_ENUMS:
-            self.register_type(python_type, pg_type)
+    def get_thread_conn(self) -> psycopg.Connection | None:
+        return getattr(self._local, "conn", None)
 
-    def register_type(self, enum_type: type[enum.Enum], pg_type: str) -> None:
-        if self._connection is None:
-            raise RuntimeError("did not connect to database")
-
-        type_info = enum.EnumInfo.fetch(self._connection, pg_type)
-        if type_info is None:
-            raise RuntimeError(f"Unable to find enum {pg_type} in DB")
-
-        enum.register_enum(
-            type_info,
-            self._connection,
-            enum_type,
-            mapping={m: m.value for m in enum_type},
-        )
+    def set_thread_conn(self, conn: psycopg.Connection | None) -> None:
+        self._local.conn = conn
 
-    def get_connection(self) -> psycopg.Connection:
-        if self._connection is None:
-            raise InternalError("unable to create database connection")
+    def get_pool(self) -> ConnectionPool:
+        if self._pool is None:
+            raise InternalError("connection pool is not initialized")
+        return self._pool
 
-        return self._connection
+    def get_connection(self) -> psycopg.Connection:
+        conn = self.get_thread_conn()
+        if conn is not None:
+            return conn
+        raise InternalError("no active transaction connection on this thread")
 
     def disconnect(self) -> None:
-        if self._connection is not None:
+        if self._pool is not None:
             self._logger.debug("disconnecting from Postgres", endpoint=self._config.endpoint, port=self._config.port)
+            self._pool.close()
 
-            self._connection.close()
-
-    def _query_str(self, query: str | sql.SQL | sql.Composed) -> str:
+    def query_str(self, query: str | sql.SQL | sql.Composed) -> str:
         if isinstance(query, str):
             return query
-        return query.as_string(self._connection)
+        conn = self.get_thread_conn()
+        if conn is not None:
+            return query.as_string(conn)
+        with self.get_pool().connection() as c:
+            return query.as_string(c)
 
     def exec(self, query: str | sql.SQL | sql.Composed, *, params: list[Any] | None = None) -> None:
         if params is None:
             params = []
-        if self._connection is None:
-            raise RuntimeError("Unable to execute query: connection to Postgres was not established")
-
-        log.debug("SQL query", query=self._query_str(query).replace("\n", " "), args=params)
 
-        cursor = self._connection.cursor()
-        cursor.execute(query, params)
+        log.debug("SQL query", query=self.query_str(query).replace("\n", " "), args=params)
 
-    def execute_batch(self, query: str, rows: Sequence[Sequence[Any]]) -> None:
-        if self._connection is None:
-            raise RuntimeError("Unable to execute query: connection to Postgres was not established")
+        conn = self.get_thread_conn()
+        if conn is not None:
+            conn.cursor().execute(query, params)
+        else:
+            with self.get_pool().connection() as c:
+                c.cursor().execute(query, params)
 
-        log.debug("SQL execute batch", query=query.replace("\n", " "), num_rows=len(rows))
+    def execute_batch(self, query: str, rows_data: Sequence[Sequence[Any]]) -> None:
+        log.debug("SQL execute batch", query=query.replace("\n", " "), num_rows=len(rows_data))
 
-        cursor = self._connection.cursor()
-        cursor.executemany(query, rows)
+        conn = self.get_thread_conn()
+        if conn is not None:
+            conn.cursor().executemany(query, rows_data)
+        else:
+            with self.get_pool().connection() as c:
+                c.cursor().executemany(query, rows_data)
 
     def query(self, query: str | sql.SQL | sql.Composed, *, params: list[Any] | None = None) -> list[rows.DictRow]:
         if params is None:
             params = []
-        if self._connection is None:
-            raise RuntimeError("Unable to execute query: connection to Postgres was not established")
-
-        log.debug("SQL query", query=self._query_str(query).replace("\n", " "), args=params)
-
-        cursor = self._connection.cursor()
-        cursor.execute(query, params)
-
-        result = cursor.fetchall()
-        log.debug("SQL result", num_rows=len(result))
 
-        return result
+        log.debug("SQL query", query=self.query_str(query).replace("\n", " "), args=params)
+
+        def _run(conn: psycopg.Connection) -> list[rows.DictRow]:
+            cursor = conn.cursor()
+            start = time.monotonic()
+            cursor.execute(query, params)
+            result = cursor.fetchall()
+            elapsed = time.monotonic() - start
+            log.debug("SQL result", num_rows=len(result), elapsed_seconds=round(elapsed, 4))
+            return result
+
+        conn = self.get_thread_conn()
+        if conn is not None:
+            return _run(conn)
+        with self.get_pool().connection() as c:
+            return _run(c)
 
     def query_one(self, query: str | sql.SQL | sql.Composed, *, params: list[Any] | None = None) -> rows.DictRow:
         result = self.query(query, params=params)
diff --git a/app/lib/storage/postgres/transactional.py b/app/lib/storage/postgres/transactional.py
@@ -9,9 +9,10 @@ def __init__(self, storage: postgres.PgStorage) -> None:
 
     @contextmanager
     def with_tx(self):
-        conn = self._storage.get_connection()
-        try:
-            with conn.transaction():
-                yield
-        except Exception as e:
-            raise e
+        with self._storage.get_pool().connection() as conn:
+            self._storage.set_thread_conn(conn)
+            try:
+                with conn.transaction():
+                    yield
+            finally:
+                self._storage.set_thread_conn(None)
diff --git a/tests/integration/transactional_storage_test.py b/tests/integration/transactional_storage_test.py
@@ -18,13 +18,12 @@ def test_several_queries(self):
 
         self.assertEqual(result["id"], 42)
 
-    def test_nested_transactions(self):
+    def test_multiple_statements_in_transaction(self):
         repo = transactional.TransactionalPGRepository(self.pg_storage.get_storage())
         with repo.with_tx():
             self.pg_storage.get_storage().exec("CREATE TABLE test_table2 (id INTEGER)")
-            with repo.with_tx():
-                self.pg_storage.get_storage().exec("INSERT INTO test_table2 VALUES (42)")
-                result = self.pg_storage.get_storage().query_one("SELECT id FROM test_table2 LIMIT 1")
+            self.pg_storage.get_storage().exec("INSERT INTO test_table2 VALUES (42)")
+            result = self.pg_storage.get_storage().query_one("SELECT id FROM test_table2 LIMIT 1")
 
         self.assertEqual(result["id"], 42)
 
