import sqlite3


class SqliteTransaction:
    def __init__(self, conn: sqlite3.Connection):
        self.conn = conn

    def __enter__(self) -> sqlite3.Cursor:
        # Start the transaction. Sqlite3 does not require explicit transaction start, so just return a new cursor:
        return self.conn.cursor()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.conn.commit()
        else:
            self.conn.rollback()


def main():
    with sqlite3.connect(':memory:') as conn:  # type: sqlite3.Connection
        with SqliteTransaction(conn) as cursor:
            cursor.execute("CREATE TABLE countries (iso_code text, label txt)")
            cursor.executemany("INSERT INTO countries VALUES(?, ?)", [
                ('CZ', 'Czech Republic'),
                ('SK', 'Slovakia'),
            ])
            print(f'Inside Context Manager: in_transaction = {conn.in_transaction}')

        for row in conn.execute("SELECT * FROM countries"):
            print(row)
        print(f'After Context manager: in_transaction = {conn.in_transaction}')


if __name__ == '__main__':
    main()
