Context manager for database connection

[Jason Friedman]

I want to be able to write code like this:

with Database() as mydb:
conn = mydb.get_connection()
cursor = conn.get_cursor()
cursor.execute("update table1 set x = 1 where y = 2")
cursor.close()
cursor = conn.get_cursor()
cursor.execute("update table2 set a = 1 where b = 2")
cursor.close()

I'd like for both statements to succeed and commit, or if either fails to
stop and for all to rollback.

Is what I have below correct?


import jaydebeapi as jdbc
class Database:
    database_connection = None

    def __init__(self, auto_commit: bool = False):
        self.database_connection = jdbc.connect(...)
        self.database_connection.jconn.setAutoCommit(auto_commit)

    def __enter__(self) -> jdbc.Connection:
        return self

    def __exit__(self, exception_type: Optional[Type[BaseException]],
                 exception_value: Optional[BaseException],
                 traceback: Optional[types.TracebackType]) -> bool:
        if exception_type:
            self.database_connection.rollback()
        else:
            self.database_connection.commit()
        self.database_connection.close()

    def get_connection(self) -> jdbc.Connection:
        return self.database_connection