In Python sqlite3, the context manager
with sqlite3.connect(database)
would automatically commit or rollback transactions. However, it does not close the connection nor the cursor. So every time I have to write something like this:
with sqlite3.connect(database) as conn:
cur = conn.cursor()
cur.execute(an_sql_statement)
cur.executemany(more_sql_statements)
cur.close()
conn.close()
which becomes quite repetitive.
I want to be able to do something like this:
with AutoCloseDB(database) as cur:
cur.execute(an_sql_statement)
cur.executemany(more_sql_statements)
and the above context manager would automatically commit or rollback transactions, close the cursor and finally close the database altogether upon exiting.
So I came up with the following context manager:
import sqlite3
class AutoCloseDB:
"""A context manager that automatically closes the cursor and the database.
Return a cursor object upon entering.
"""
def __init__(self, database):
self.conn = sqlite3.connect(database)
def __enter__(self):
self.conn = self.conn.__enter__()
self.cur = self.conn.cursor()
return self.cur
def __exit__(self, *exc_info):
result = self.conn.__exit__(*exc_info)
self.cur.close()
self.conn.close()
return result
Is there any problem in the above code? Is there any better solution?