Cloned DMOJ

judge/dblock.py

@@ -0,0 +1,23 @@
+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()