Connection pooling with psycopg2

In this chapter we'll look at replacing our create_connection() function with a connection pool.

Our model class methods that use create_connection() will instead get a connection from the pool, and put it back into the pool when they're done (instead of closing the connection).

Creating the connection pool

I'm going to start by creating a file called connection_pool.py.

This is how we can create a simple connection pool that is not thread-safe^[1]:

import os
from psycopg2.pool import SimpleConnectionPool
from dotenv import load_dotenv

DATABASE_PROMPT = "Enter the DATABASE_URI value or leave empty to load from .env file: "

database_uri = input(DATABASE_PROMPT)
if not database_uri:
    load_dotenv()
    database_uri = os.environ["DATABASE_URI"]

pool = SimpleConnectionPool(minconn=1, maxconn=10, dsn=database_uri)

As you can see, I've:

• imported SimpleConnectionPool from psycopg2.pool
• Re-added the database prompt, since now we'll only need to create the connections in one place, and therefore we'll only ask this once.
• Created the pool using pool = SimpleConnectionPool().

Arguments to SimpleConnectionPool

• minconn is the minimum number of connections. These are the connections that will be created immediately when the pool is created.
• maxconn is the maximum number of connections. If we request a connection from the pool but we have ran out, a new connection will be created unless we already have maxconn connections. If that's the case, an error will be raised.
• **kwargs get passed to psycopg2.connect, so SimpleConnectionPool can actually accept any keyword arguments that psycopg2.connect would accept in order to connect to the database. Here we're passing dsn= which is the database URI.

Using the connection pool

Using the pool is very easy. Our model classes will call pool.getconn() and pool.putconn(), like so:

def save(self):
    connection = pool.getconn()
    new_poll_id = database.create_poll(connection, self.title, self.owner)
    pool.putconn(connection)
    self.id = new_poll_id

Note that pool.getconn() returns a connection. We assign that to a variable, and then pass it to pool.putconn(connection) at the end.

We should do this to all model class methods, so that they're no longer calling create_connection().

Remember to change the imports as well:

-from connections import create_connection
+from connection_pool import pool

Changes to app.py

In app.py we're creating our tables, so we need to make a small change there too to use the connection pool:

-connection = create_connection()
+connection = pool.getconn()
database.create_tables(connection)
+pool.putconn(connection)

Enjoyed this article?
You'll love the complete video course!

• Complete video lessons for each topic
• Get the most out of your learning
• Master Advanced PostgreSQL
• Get one-on-one help from a real person

Get the video course

---

1. psycopg2.pool.ThreadedConnectionPool (opens new window) ↩︎