Search code examples
pythoncommitmysql-connectorrollbackcontextmanager

How to commit or rollback a change in database when using a custom context manager class with mysql connector


What is the best way to handle committing or rolling back a change in the database when using mysql connector through a custom context manager class. For example, say the custom class takes the form:

class DatabaseConnection:
    def __init__(self, host, user, password, database_name)
        self.host = host
        self.user = user
        self.passwd = password
        self.database_name = database_name
        self.connection = None

    def __enter__(self) -> Union[mysql.cursor.MySQLCursor, None]:
       try:
          if self.connection is None or not sef.connection.is_connected():
              self.connection = mysql.connect(
                  host = self.host, user=self.user, psswd=self.psswd, db=self.database_name
              )
          cursor = self.connection.cursor()
          return cursor
      except Exception:
          return None

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.connection is not None: 
            self.connection.close()

    def commit(self):
        if self.connection is not None and self.connection.is_connected():
             try:
                 self.connection.commit()
                 return True
             except:
                 self.connection.rollback()

    def rollback(self):
        if self.connection is not None:
           self.connection.rollback()

And you then want to use that using the with statement so that you have the cursor. How would you manage commiting or rollback the changes, especially when errors are raised. For example:

test = True
database_connection = DatabaseConnection(....)
try:
    with database_connection() as cursor:
        # perform some operations

        if test:
            database_connection.rollback()
        else:
            database_connection.commit()
except Exception as e:
     database_connection.rollback()

Ignoring the bad practice with respect to error handling, would the commit and rolling back work as expected or would you have to organize it in a different way?

EDIT

The above pattern seems to work as expected (if test, rollback, if not then commit and if exception then rollback). I'm not entirely sure why so any enlightenment on this regard would be helpful. Thanks!


Solution

  • It works because you have written it to handle all possibilities. It is odd however that you don't make use of the object you create in the context manager.

    Consider the possibilities:

    1.Test succeeds:
    commit is performed. (If commit throws an exception it will be caught and handled in your commit function)
    2. Test fails:
    rollback is performed. (note however you have done this outside the context manager which is a little odd because the connection (for cursor at least) is closed )
    3. An exception is thrown:
    The exception is caught and a rollback is performed.

    So all cases are handled. The one thing to note is that you probably want to log something to let a user know that you failed to commit and performed a rollback.
    It is also notable that the syntax of: database_connection = DatabaseConnection(...) doesn't me make sense if you want to use the context manager as you do on the next line. You should be able to condense the two lines to simply with DatabaseConnection(...) as cursor: and then replace usages of database_connection in the context manager with cursor. That's the point to using the context manager in the first place. Don't mix and match between two different objects that are the 'same'-ish connection to the database where one is context managed and the other isn't. Thats bad coding and extremely unpythonic and could yield very unexpected results.

    Which then raises @snakecharmerb's point that you should look at putting the rollback in the __exit__ code. As it stands you are basically not using the context manager.