Creating an PostgreSQL database connection pool in Python

Directory

You must create a connection pool before using the database. Even for single-threaded applications, as long as there are multiple methods that need database connections, those who establish 1 or 2 connections will consider pooling them first. Connection pooling has many benefits,

1) If creating a connection repeatedly is quite time consuming, 2) For an application where a single connection is used to the end, when there is a connection pool, it avoids the transfer of database connection objects. 3) Forgot to close the connection, and the connection pool can help to close it after 1 time. Of course, the application of dense connection will run out of connection. 4) The number of open connections per application is controllable

Come into contact with Python After using PostgreSQL It is also natural to consider creating a connection pool, taking it from the pool when you use it and returning it when you use it, instead of creating a physical connection every time you need it. Python Connect PostgreSQL There are mainly two bags, py-postgresql And psycopg2 And the example in this article will use the latter.

Psycopg In psycopg2.pool Module provides the implementation of two connection pools, both of which inherit from psycopg2.pool.AbstractConnectionPool,

The basic method of this abstract class is

getconn(key=None): Get a connection PostgreSQL 0 Return Connection closeall(): Close all connections in the connection pool

The implementation classes for the two connection pools are

psycopg2.pool.SimpleConnectionPool(minconn, maxconn, *args, **kwars) For single-threaded applications psycopg2.pool.ThreadedConnectionPool(minconn, maxconn, *args, **kwars) Multithreaded is safer, in fact, in the getconn() And putconn() A lock is added to control

So the safest way to insure is to use ThreadedConnectionPool In single-threaded applications, SimpleConnectionPool  It's not necessarily better than ThreadedConnectionPool How efficient is it?

Let's look at a specific connection pool implementation, which uses Context Manager , combined with when using with Keyword is more convenient, and no explicit call is needed after use putconn() Return Connection

db_helper.py

from psycopg2 import pool
from psycopg2.extras import RealDictCursor
from contextlib import contextmanager
import atexit


class DBHelper:
    def __init__(self):
        self._connection_pool = None

    def initialize_connection_pool(self):
        db_dsn = 'postgresql://admin:password@localhost/testdb?connect_timeout=5'
        self._connection_pool = pool.ThreadedConnectionPool(1, 3 , db_dsn)

    @contextmanager
    def get_resource(self, autocommit=True):
        if self._connection_pool is None:
            self.initialize_connection_pool()

        conn = self._connection_pool.getconn()
        conn.autocommit = autocommit
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        try:
            yield cursor, conn
        finally:
            cursor.close()
            self._connection_pool.putconn(conn)

    def shutdown_connection_pool(self):
        if self._connection_pool is not None:
            self._connection_pool.closeall()


db_helper = DBHelper()


@atexit.register
def shutdown_connection_pool():
    db_helper.shutdown_connection_pool()
from psycopg2 import pool

from psycopg2 . extras import RealDictCursor

from contextlib import contextmanager

import atexit

class DBHelper :

     def __init__ ( self ) :

         self . _connection_pool = None

     def initialize_connection_pool ( self ) :

         db_dsn = 'postgresql://admin:password@localhost/testdb?connect_timeout=5'

         self . _connection_pool = pool . ThreadedConnectionPool ( 1 , 3  ,  db_dsn )

     @ contextmanager

     def get_resource ( self , autocommit = True ) :

         if self . _connection_pool is None :

             self . initialize_connection_pool ( )

         conn = self . _connection_pool . getconn ( )

         conn . autocommit = autocommit

         cursor = conn . cursor ( cursor_factory = RealDictCursor )

         try :

             yield cursor , conn

         finally :

             cursor . close ( )

             self . _connection_pool . putconn ( conn )

     def shutdown_connection_pool ( self ) :

         if self . _connection_pool is not None :

             self . _connection_pool . closeall ( )

db_helper = DBHelper ( )

@ atexit . register

def shutdown_connection_pool ( ) :

     db_helper . shutdown_connection_pool ( )

A few notes:

Only on the first call get_resource() Instead of creating a connection pool in the from db_helper import db_helper Connection pool is created when referenced Context Manager Returns two objects, cursor And connection , need to use connection Use it when managing things Default time cursor Records returned are dictionaries, not arrays The connection is automatically submitted by default Last @atexit.register That PostgreSQL 0 It may be redundant, and the connection is closed when the process exits. PostgreSQL 1 It should take a little longer

Usage:

If you don't use things

from db_helper import db_helper


with db_helper.get_resource() as (cursor, _):
    cursor.execute('select * from users')
    for record in cursor.fetchall():
        ... process record, record['name'] ...
from db_helper import db_helper

with db_helper . get_resource ( ) as ( cursor , _ ) :

     cursor . execute ( 'select * from users' )

     for record in cursor . fetchall ( ) :

         . . . process record , record [ 'name' ] . . .

If you need to use things,

with db_helper.get_resource(autocommit=False) as (cursor, _):
    try:
        cursor.execute('update users set name = %s where id = %s', ('new_name', 1))
        cursor.execute('delete from orders where user_id = %s', (1,))
        conn.commit()
    except:
        conn.rollback()
with db_helper . get_resource ( autocommit = False ) as ( cursor , _ ) :

     try :

         cursor . execute ( 'update users set name = %s where id = %s' , ( 'new_name' , 1 ) )

         cursor . execute ( 'delete from orders where user_id = %s' , ( 1 , ) )

         conn . commit ( )

     except :

         conn . rollback ( )

When writing this article, when I looked at the official website of psycopg, I found that PostgreSQL 2 The official version was released on October 13, 2021 (Psycopg 3.0 released), which better supports async. Support for asynchronism began with Psycopg 2 2.2. It is also noted that Psycopg The main part of C is implemented with C, which makes it more efficient. It is no wonder that it is often used pip install psycopg2 Installation was unsuccessful, and you want to use pip install psycopg2-binary To install the reason.

Adding parameters when creating connection pools keepalivesXxx Can let the server break the dead link in time, otherwise in Linux By default, it will take 2 hours to disconnect. Dead link occurs when the client exits abnormally (such as when the power is cut off), and the previously established link becomes dead link.

 pool.ThreadedConnectionPool(1, 3, db_dsn, keepalives=1, keepalives_idle=30, keepalives_interval=10, keepalives_count=5) 

PostgreSQL The server will be idle for the connection tcp_keepalives_idle Seconds later, actively send tcp_keepalives_count 个 tcp_keeplive Detective pack, every detective pack is in tcp_keepalives_interval If there is no response within seconds, it is considered as a dead connection, so it is cut off.