1
\$\begingroup\$

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... 
        
\$\endgroup\$

1 Answer 1

1
\$\begingroup\$
  • 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.
\$\endgroup\$
0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.