NDOJ/judge/dblock.py

from itertools import chain

from django.db import connection, transaction


class LockModel(object):
    def __init__(self, write, read=()):
        self.tables = ", ".join(
            chain(
                ("`%s` WRITE" % model._meta.db_table for model in write),
                ("`%s` READ" % model._meta.db_table for model in read),
            )
        )
        self.cursor = connection.cursor()

    def __enter__(self):
        self.cursor.execute("LOCK TABLES " + self.tables)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            transaction.commit()
        else:
            transaction.rollback()
        self.cursor.execute("UNLOCK TABLES")
        self.cursor.close()