psycopg2 ThreadedConnectionPool implementation

Question

I am currently using a ThreadedConnectionPool from psycopg2 to lease transactions to a PostgreSQL database and would like a review of my current implementation.

Code

from __future__ import annotations
import logging
import threading
import traceback
from contextlib import contextmanager
from dataclasses import dataclass
from psycopg2.pool import ThreadedConnectionPool
from .settings import CONNECTION_OPEN_WARN_THRESHOLD
from .utils import Timer
logger = logging.getLogger()
class Counter(object):
 def __init__(self, name):
 self.name = name
 self.value = 0
 self.max = 0
 self._lock = threading.Lock()
 def increment(self):
 with self._lock:
 self.value += 1
 if self.value > self.max:
 self.max = self.value
 def decrement(self):
 with self._lock:
 self.value -= 1
 def __str__(self) -> str:
 return f"{id(self):#02x} - {self.name}: current={self.value}, max={self.max}"
@dataclass
class PostgreSQLSimplePool:
 pool: ThreadedConnectionPool
 counter = Counter("Active DB connections")
 @contextmanager
 def transact_session(self, commit: bool = False):
 with Timer() as t:
 conn = self.pool.getconn()
 self.counter.increment()
 logger.debug(self.counter)
 try:
 yield conn
 if commit:
 conn.commit()
 except Exception:
 conn.rollback()
 raise
 finally:
 self.counter.decrement()
 self.pool.putconn(conn)
 if t.interval > CONNECTION_OPEN_WARN_THRESHOLD:
 logger.warning(f"DB Connection was held open for {t.interval} seconds:")
 traceback.print_stack()

Config

threaded_pool = ThreadedConnectionPool(
 minconn=1,
 maxconn=20,
 dsn="", # This relies on standard env vars
)
repo = PostgreSQLSimplePool(pool=threaded_pool)

Usage

with repo.transact_session() as connection:
 with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
 # do some stuff with the cursor... 
 

Answer 1

• Counter should not inherit from object in Python 3
• Add PEP484 type-hinting to name, as well as all of your return values
• I'm not sure what you were hoping to accomplish with a field width of 2 in {id(self):#02x}. IDs will almost certainly be over two characters long when formatted.
• counter = Counter("Active DB connections") does not do what you think it does. It doesn't make an instance member on the dataclass, it makes a static member. You'll want to use field passing a default, or perhaps don't write it as a dataclass member at all and instead add it during __post_init__.
• You have a counter mechanism that has no knowledge of the pool's size. In my imagination the counter's max and the pool's maxconn should be the same thing. Currently you have a "statistical" max that only reports on the historical peak connection usage. You could have - instead, or in addition - an "enforced" max. In other words: rather than rolling your own counter with a lock, perhaps you want a semaphore that blocks if the max count is hit.