A Useful Context Manager

Found all over the Python ecosystem, “Context Managers” are the with something as var: you might use for, say, file I/O:

with open('test.txt', 'w') as f:
    pass

Essentially, this has the behavior of:

  • Entering a context and executing some code (here, opening test.txt in write mode)
  • Doing some stuff (just passing here)
  • Doing cleanup/shutdown behavior (closing the file)

This is similar to how we interact with SQL code, insofar as we:

  • Connect to a database
  • Make some modifications
  • Commit our changes
  • Close our connection

Thus, using the contextlib library, we can hook into this syntax to good use, with the following:

import sqlite3
from contextlib import contextmanager

@contextmanager
def connect_to_db(database):
    path = '.'
    conn = sqlite3.connect(path+database)
    cur = conn.cursor()
    yield cur
    conn.commit()
    conn.close()

Now, all we need to do is call this function and supply a .db database file.

Here, we’ll select from a table that doesn’t exist yet

try:
    with connect_to_db('test.db') as cursor:
        cursor.execute('select * from new_table')
        cursor.fetchone()
except Exception as e:
    print(e)
no such table: new_table

Then we’ll make the database and insert a record into it

with connect_to_db('test.db') as cursor:
    cursor.execute('create table new_table (id integer)')
    cursor.execute('insert into new_table (id) values (1)')

Finally, we rerun the code that works this time

try:
    with connect_to_db('test.db') as cursor:
        cursor.execute('select * from new_table')
        print(cursor.fetchone())
except Exception as e:
    print(e)
(1,)