# 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
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
With that, we've substantially reduced duplication! If later on we want to swap
psycopg2 for something else, it'll be slightly easier!