2
\$\begingroup\$

This is my database connection class in order to handle multiple connections with the use of context managers. Brief explanation:

  • The __init__ method reads database configurations from an INI file in order to call the connection and set two dictionaries, one for connections and one for cursors.
  • The __enter__ method takes the specified database connection and cursor returning at the end the cursor to execute queries.
  • The __exit__ method commits the query and close the cursor if no exception has been identified.

I'm using PyMySQL as MySQL module but I think that the code can be generalized for all MySQL modules, the functions are always those.

Any advice is welcome, I'm new with Python!

import os.path
import configparser
import pymysql.cursors

class Database:
    def __init__(self):
        config = configparser.ConfigParser()

        dbConfigFile = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'dbconfig.ini'))
        config.read(dbConfigFile)

        # Connections (database:connection method)
        self.connections = {
            'db1': self.connection(dict(config.items('db1info'))),
            'db2': self.connection(dict(config.items('db2info')))
        }

        # Cursors list (database:cursor method)
        self.cursors = {
            'db1': self.cursor(self.connections['db1']),
            'db2': self.cursor(self.connections['db2'])
        }

    def connection(self, config):
        return pymysql.connect(**config)

    def cursor(self, connection):
        return connection.cursor

    def __call__(self, database):
        self.database = database
        return self

    def __enter__(self):
        self.db = self.connections[self.database]
        self.cursor = self.cursors[self.database](pymysql.cursors.SSDictCursor)
        return self.cursor

    def __exit__(self, exc_type, exc_value, tb):
        if not exc_type:
            self.db.commit()
            self.cursor.close()

Usage:

db = Database()
with db('db1') as cursor:
    cursor.execute("SELECT 1 FROM table_name")
\$\endgroup\$
3
\$\begingroup\$

Neat idea! Some suggestions:

  • __exit__ should try to commit or roll back based on whether there is an unhandled exception, should then close the cursor, and should finally unconditionally close the connection. This is one of the main reasons for using a context manager - it is expected to always clean up after itself, leaving the relevant parts of the system in the same state as before using it.
  • I would reuse the configuration section name as the keys in the dict. That way you don't need to maintain a mapping in your head or the code - what's in the configuration is what you get when you use the context manager.
  • Rather than opening all connections and then using only one of them, it should open the connection and the cursor with the name passed in. Otherwise you're wasting resources.
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.