# Reduce cursor creation duplication with context managers
Now that we've reduced the connection-pooling duplication, let's create another context manager that helps us reduce duplication in database.py
.
Every function does something like this:
def get_poll(connection, poll_id: int) -> Poll:
with connection:
with connection.cursor() as cursor:
cursor.execute(SELECT_POLL, (poll_id,))
return cursor.fetchone()
We can do better than having this duplicated code.
So what I'll do is create a context manager that starts a transaction (with with connection
), gets a cursor (with with connection.cursor() as cursor
) and then yields that to the caller function.
Note that I'm doing this because every function just gets one cursor. If you have functions that create multiple cursors (or none at all), then you might not want to do this.
Something else that I want to keep is that model classes should still pass a connection to the functions they call.
This is my context manager:
from contextlib import contextmanager
...
@contextmanager
def get_cursor(connection):
with connection:
with connection.cursor() as cursor:
yield cursor
Then, in every function we will use this context manager:
def get_poll(connection, poll_id: int) -> Poll:
with get_cursor(connection) as cursor:
cursor.execute(SELECT_POLL, (poll_id,))
return cursor.fetchone()
We should apply these changes to all the functions in database.py
.
With that, we've substantially reduced duplication! If later on we want to swap psycopg2
for something else, it'll be slightly easier!