Automatically closing database connection with Context Manager
Photo by Taylor Vick on Unsplash
If you don't know what a Context Manager is I recommend you to read about them. Dan Bader wrote a good article. But since you are here, why not read my post?
Now, it's nothing special to write a cursor. You need a driver and credentials to connect to the database. In this example I'll use MySQL driver. My credentials are stored in settings.py
(not in plain text but environment variables) as a dictionary.
First, we need to provide a driver and credentials to our cursor:
import mysql.connector as connector
from settings import DATABASE
class Cursor:
def __init__(self,
host=DATABASE.get('HOST'),
user=DATABASE.get('USER'),
password=DATABASE.get('PASSWORD'),
db_name=DATABASE.get('NAME'),
driver=connector,
):
self.driver = driver
self.connection = self.driver.connect(
host=host,
user=user,
password=password,
database=db_name
)
self.cursor = self.connection.cursor()
Now we need to provide methods of a Context Manager to our class:
class Cursor:
def __init__(...)
def __enter__(self):
return self.cursor
def __exit__(self, ext_type, exc_value, traceback):
self.cursor.close()
if isinstance(exc_value, Exception):
self.connection.rollback()
else:
self.connection.commit()
self.connection.close()
And finally, we need to return something from the database, when it is needed:
class Cursor:
def __init__(...)
def __iter__(self):
for item in self.cursor:
yield item
def __enter__(...)
def __exit__(...)
Done. Usage is a simple with Cursor() as cursor:
I've never bothered to simplify it via @contextmanager
decorator because this implementation works perfectly fine for me. And I'm not sure if we can apply that decorator to a class.
I'd like to use ORM because it makes things so much easier and faster. But sometimes ORM is an overkill and you need to interact with DB manually.
Hope you find this helpful if you ever need to write a custom Cursor.
Top comments (1)
ORM is a middleware / abstraction layer. It can not be "faster" (unless you mess up something in your code badly). It is always more performant to use the underlying driver directly. Sometimes there is a huge difference. The "easier" part might be valid though.