Proper management of database resources: cursor and connection - mysql

Question

I am creating a test Flask API, and have created a Database class that I use from my main app. I am using pymysql to access my MySQL DB but I am having trouble figuring out when to close the cursor and connection. Right now I have
import pymysql
class Database:
def __init__(self):
host = '127.0.0.1'
user = 'root'
password = ''
db = 'API'
self.con = pymysql.connect(host=host, user=user, password=password, db=db, cursorclass=pymysql.cursors.DictCursor, autocommit=True)
self.cur = self.con.cursor()
def getUser(self, id):
sql = 'SELECT * from users where id = %d'
self.cur.execute(sql, (id))
result = self.cur.fetchall()
return result
def getAllUsers(self):
sql = 'SELECT * from users'
self.cur.execute(sql)
result = self.cur.fetchall()
return result
def AddUser(self, firstName, lastName, email):
sql = "INSERT INTO `users` (`firstName`, `lastName`, `email`) VALUES (%s, %s, %s)"
self.cur.execute(sql, (firstName, lastName, email))
I have tried adding self.cur.close() and self.con.close() after each execution of the cursor in the functions but then I get an error the next time I call a function saying the cursor is closed, or after I do an insert statement it won't show the new value even though it was inserted correctly into MySQL. How do I know when to close the cursor, and how to start it back up properly with each call to a method?

Answer 1

This sounds like a great use case for a python context manager. Context Managers allow you to properly manage resources, such as a database connection, by allowing you to specify how your resource's set-up and tear down methods should work. You can create your own custom context manager in one of two ways: First, by wrapping your database class, and implementing the required methods for the context manager: __init__(), __enter__(), and __exit__(). Second, by utilizing a #contextmanager decorator on a function definition and creating a generator for your database resource within said function definition. I will show both approaches and let you decide which one is your preference. The __init__() method is the initialization method for your custom context manager, similar to the initialization method used for custom python classes. The __enter__() method is your setup code for your custom context manager. Lastly, the __exit()__ method is your teardown code for your custom context manager. Both approaches utilize these methods with the main difference being the first method will explicitly state these methods within your class definition. Where as in the second approach, all the code up to your generator's yield statement is your initialization and setup code and all the code after the yield statement is your teardown code. I would also consider extracting out your user based database actions into a user model class as well. Something along the lines of:
custom context manager: (class based approach):
import pymysql
class MyDatabase():
def __init__(self):
self.host = '127.0.0.1'
self.user = 'root'
self.password = ''
self.db = 'API'
self.con = None
self.cur = None
def __enter__(self):
# connect to database
self.con = pymysql.connect(host=self.host, user=self.user, password=self.password, db=self.db, cursorclass=pymysql.cursors.DictCursor, autocommit=True)
self.cur = self.con.cursor()
return self.cur
def __exit__(self, exc_type, exc_val, traceback):
# params after self are for dealing with exceptions
self.con.close()
user.py (refactored):'
# import your custom context manager created from the step above
# if you called your custom context manager file my_database.py: from my_database import MyDatabase
import <custom_context_manager>
class User:
def getUser(self, id):
sql = 'SELECT * from users where id = %d'
with MyDatabase() as db:
db.execute(sql, (id))
result = db.fetchall()
return result
def getAllUsers(self):
sql = 'SELECT * from users'
with MyDatabase() as db:
db.execute(sql)
result = db.fetchall()
return result
def AddUser(self, firstName, lastName, email):
sql = "INSERT INTO `users` (`firstName`, `lastName`, `email`) VALUES (%s, %s, %s)"
with MyDatabase() as db:
db.execute(sql, (firstName, lastName, email))
context manager (decorator approach):
from contextlib import contextmanager
import pymysql
#contextmanager
def my_database():
try:
host = '127.0.0.1'
user = 'root'
password = ''
db = 'API'
con = pymysql.connect(host=host, user=user, password=password, db=db, cursorclass=pymysql.cursors.DictCursor, autocommit=True)
cur = con.cursor()
yield cur
finally:
con.close()
Then within your User class you could use the context manager by first importing the file and then using it similar to as before:
with my_database() as db:
sql = <whatever sql stmt you wish to execute>
#db action
db.execute(sql)
Hopefully that helps!