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!
