The Complete Python/PostgreSQL Course 2.0

# Reduce pooling duplication with context managers

Now all our model class methods have essentially identical code:

...
connection = pool.getconn()
...
pool.putconn(connection)
...

While it's not a lot, it does mean that later on if we want to change how the connection pool works (e.g. if we stop using psycopg2), we're going to have to go in there and change every method.

We can do better!

We'll create a context manager[1] that retrieves the connection for us from the pool, and puts it back when we're done.

If you're not familiar with context managers or the yield[2] keyword, please check out the resources linked.

# Creating our context manager

This is our new context manager, which I'm writing inside connection_pool.py:

from contextlib import contextmanager

..

@contextmanager
def get_connection():
    connection = pool.getconn()

    try:
        yield connection
    finally:
        pool.putconn(connection)

The new context manager is called get_connection.

EXPLANATION

# Using our context manager in the models

Let's go to the Poll model class and use this context manager.

Remember to import it, and we no longer need to import the pool:

-from connection_pool import pool
+from connection_pool import get_connection

Then let's change the method. This:

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

Becomes this:

def save(self):
    with get_connection() as connection:
        new_poll_id = database.create_poll(connection, self.title, self.owner)
        self.id = new_poll_id

We should do the same for every other method in any model that uses the pool.

# Using our context manager in app.py

Let's do the same in app.py!

-from connection_pool import pool
+from connection_pool import get_connection

...

-connection = pool.getconn()
-database.create_tables(connection)
-pool.putconn(connection)
+with get_connection() as connection:
+   database.create_tables(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. The Curious Case of Python's Context Manager (opens new window) ↩︎

  2. What does the yield keyword do? (StackOverflow) (opens new window) ↩︎

← Reduce cursor creation duplication with context managers →