Repository: adamchainz/django-mysql

Issue:
Table locking context manager
For MyISAM tables especially, taking a table lock can be sometimes useful. A simple context manager could be implemented to help lock them. There are some snippets about already:
- http://shiningpanda.com/mysql-table-lock-django.html
- https://djangosnippets.org/snippets/2039/

This would need to work, or error sanely, with:
- [x] plain models
- [x] proxy models
- [x] abstract models
- [ ] swapped models
- [x] inherited models

Also it should work as both a context manager and a function decorator.



Discussion:

Resolving patch:
diff --git a/django_mysql/locks.py b/django_mysql/locks.py
--- a/django_mysql/locks.py
+++ b/django_mysql/locks.py
@@ -1,5 +1,8 @@
 # -*- coding:utf-8 -*-
+from collections import OrderedDict
+
 from django.db import connections
+from django.db.transaction import TransactionManagementError, atomic
 from django.db.utils import DEFAULT_DB_ALIAS
 
 from django_mysql.exceptions import TimeoutError
@@ -83,3 +86,67 @@ def held_with_prefix(cls, prefix, using=DEFAULT_DB_ALIAS):
                 cls.unmake_name(using, row[0]): row[1]
                 for row in cursor.fetchall()
             }
+
+
+class TableLock(object):
+    def __init__(self, read=None, write=None, using=None):
+        self.read = self._process_names(read)
+        self.write = self._process_names(write)
+        self.db = DEFAULT_DB_ALIAS if using is None else using
+
+    def _process_names(self, names):
+        """
+        Convert a list of models/table names into a list of table names. Deals
+        with cases of model inheritance, etc.
+        """
+        if names is None:
+            return []
+
+        table_names = OrderedDict()  # Preserve order and ignore duplicates
+        while len(names):
+            name = names.pop(0)
+            if hasattr(name, '_meta'):
+                if name._meta.abstract:
+                    raise ValueError("Can't lock abstract model {}"
+                                     .format(name.__name__))
+
+                table_names[name._meta.db_table] = True
+                # Include all parent models - the keys are the model classes
+                if name._meta.parents:
+                    names.extend(name._meta.parents.keys())
+            else:
+                table_names[name] = True
+        return table_names.keys()
+
+    def __enter__(self):
+        self.lock()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.unlock(exc_type, exc_value, traceback)
+
+    def lock(self):
+        connection = connections[self.db]
+        qn = connection.ops.quote_name
+        with connection.cursor() as cursor:
+            if not connection.get_autocommit():
+                raise TransactionManagementError(
+                    "InnoDB requires that we not be in a transaction when "
+                    "gaining a table lock.")
+
+            # Begin transaction - does 'SET autocommit = 0'
+            self._atomic = atomic(using=self.db)
+            self._atomic.__enter__()
+
+            locks = []
+            for name in self.read:
+                locks.append("{} READ".format(qn(name)))
+            for name in self.write:
+                locks.append("{} WRITE".format(qn(name)))
+            cursor.execute("LOCK TABLES {}".format(", ".join(locks)))
+
+    def unlock(self, exc_type, exc_value, traceback):
+        connection = connections[self.db]
+        with connection.cursor() as cursor:
+            self._atomic.__exit__(exc_type, exc_value, traceback)
+            self._atomic = None
+            cursor.execute("UNLOCK TABLES")


Test patch:
diff --git a/tests/django_mysql_tests/models.py b/tests/django_mysql_tests/models.py
--- a/tests/django_mysql_tests/models.py
+++ b/tests/django_mysql_tests/models.py
@@ -130,6 +130,29 @@ class Alphabet(Model):
     g = DecimalField(default=0, decimal_places=2, max_digits=10)
 
 
+class ProxyAlphabet(Alphabet):
+    class Meta:
+        proxy = True
+
+    @property
+    def a_times_b(self):
+        return self.a * self.b
+
+
+class Customer(Model):
+    name = CharField(max_length=32)
+
+
+class AgedCustomer(Customer):
+    age = IntegerField(default=21)
+
+
+class TitledAgedCustomer(AgedCustomer):
+    class Meta:
+        db_table = 'titled aged customer'  # Test name quoting
+    title = CharField(max_length=16)
+
+
 class SizeFieldModel(Model):
     binary1 = SizedBinaryField(size_class=1)
     binary2 = SizedBinaryField(size_class=2)
diff --git a/tests/django_mysql_tests/test_locks.py b/tests/django_mysql_tests/test_locks.py
--- a/tests/django_mysql_tests/test_locks.py
+++ b/tests/django_mysql_tests/test_locks.py
@@ -2,12 +2,17 @@
 from threading import Thread
 
 import pytest
-from django.db import connection
-from django.test import TestCase
+from django.db import OperationalError, connection
+from django.db.transaction import TransactionManagementError, atomic
+from django.test import TestCase, TransactionTestCase
 from django.utils.six.moves import queue
 
 from django_mysql.exceptions import TimeoutError
-from django_mysql.locks import Lock
+from django_mysql.locks import Lock, TableLock
+from django_mysql.models import Model
+from django_mysql_tests.models import (
+    AgedCustomer, Alphabet, Customer, ProxyAlphabet, TitledAgedCustomer
+)
 
 
 class LockTests(TestCase):
@@ -205,3 +210,183 @@ def test_held_with_prefix(self):
 
         assert Lock.held_with_prefix('') == {}
         assert Lock.held_with_prefix('mylock') == {}
+
+
+class TableLockTests(TransactionTestCase):
+    def tearDown(self):
+        Alphabet.objects.all().delete()
+        Alphabet.objects.using('other').all().delete()
+        Customer.objects.all().delete()
+        Customer.objects.using('other').all().delete()
+
+    def test_write(self):
+        Alphabet.objects.create(a=12345)
+        with TableLock(write=[Alphabet]):
+            assert Alphabet.objects.count() == 1
+
+            Alphabet.objects.all().delete()
+            assert Alphabet.objects.count() == 0
+
+        assert Alphabet.objects.count() == 0
+
+    def test_write_with_table_name(self):
+        Alphabet.objects.create(a=71718)
+        with TableLock(write=[Alphabet._meta.db_table]):
+            assert Alphabet.objects.count() == 1
+
+            Alphabet.objects.all().delete()
+            assert Alphabet.objects.count() == 0
+
+        assert Alphabet.objects.count() == 0
+
+    def test_write_with_using(self):
+        Alphabet.objects.using('other').create(a=878787)
+        with TableLock(write=[Alphabet], using='other'):
+            assert Alphabet.objects.using('other').count() == 1
+
+            Alphabet.objects.using('other').all().delete()
+            assert Alphabet.objects.using('other').count() == 0
+
+        assert Alphabet.objects.using('other').count() == 0
+
+    def test_write_twice(self):
+        with TableLock(write=[Alphabet]):
+            Alphabet.objects.create(a=1)
+
+        with TableLock(write=[Alphabet]):
+            Alphabet.objects.create(a=2)
+
+    def test_write_fails_touching_other_table(self):
+        with pytest.raises(OperationalError) as excinfo:
+            with TableLock(write=[Alphabet]):
+                Customer.objects.create(name="Lizzy")
+
+        assert "was not locked with LOCK TABLES" in str(excinfo.value)
+
+    def test_read_and_write(self):
+        Customer.objects.create(name="Fred")
+        with TableLock(read=[Customer], write=[Alphabet]):
+            ab = Alphabet.objects.create(a=Customer.objects.count())
+            assert ab.a == 1
+
+    def test_creates_an_atomic(self):
+        assert connection.get_autocommit() == 1
+        assert not connection.in_atomic_block
+        with TableLock(read=[Alphabet]):
+            assert connection.get_autocommit() == 0
+            assert connection.in_atomic_block
+        assert connection.get_autocommit() == 1
+        assert not connection.in_atomic_block
+
+    def test_fails_in_atomic(self):
+        with atomic(), pytest.raises(TransactionManagementError) as excinfo:
+            with TableLock(read=[Alphabet]):
+                pass
+
+        assert str(excinfo.value).startswith("InnoDB requires that we not be")
+
+    def test_fail_nested(self):
+        with pytest.raises(TransactionManagementError) as excinfo:
+            with TableLock(write=[Alphabet]), TableLock(write=[Customer]):
+                pass
+
+        assert str(excinfo.value).startswith("InnoDB requires that we not be")
+
+    def test_atomic_works_in_lock(self):
+        Alphabet.objects.create(a=4567)
+        with TableLock(write=[Alphabet]):
+            assert Alphabet.objects.count() == 1
+
+            try:
+                with atomic():
+                    raise ValueError("Hi")
+            except ValueError:
+                pass
+
+            Alphabet.objects.all().delete()
+            assert Alphabet.objects.count() == 0
+
+    def test_writes_fail_under_read(self):
+        with TableLock(read=[Alphabet]):
+            with pytest.raises(OperationalError) as excinfo:
+                Alphabet.objects.update(a=2)
+
+            assert (
+                "was locked with a READ lock and can't be updated" in
+                str(excinfo.value)
+            )
+
+    def test_fails_with_abstract_model(self):
+        with pytest.raises(ValueError) as excinfo:
+            with TableLock(read=[Model]):
+                pass
+
+        assert "Can't lock abstract model Model" in str(excinfo.value)
+
+    def test_proxy_model(self):
+        Alphabet.objects.create(a=2, b=3)
+
+        with TableLock(read=[ProxyAlphabet]):
+            ab = ProxyAlphabet.objects.get()
+            assert ab.a_times_b == 6
+
+        with TableLock(write=[ProxyAlphabet]):
+            assert ProxyAlphabet.objects.count() == 1
+            ProxyAlphabet.objects.all().delete()
+            assert ProxyAlphabet.objects.count() == 0
+
+        assert ProxyAlphabet.objects.count() == 0
+
+    def test_inherited_model(self):
+        TitledAgedCustomer.objects.create(title="Sir", name="Knighty")
+
+        with TableLock(write=[TitledAgedCustomer]):
+            assert Customer.objects.count() == 1
+
+            TitledAgedCustomer.objects.create(name="Grandpa Potts", age=99)
+            assert Customer.objects.count() == 2
+            assert TitledAgedCustomer.objects.count() == 2
+
+            TitledAgedCustomer.objects.all().delete()
+            assert Customer.objects.count() == 0
+            assert AgedCustomer.objects.count() == 0
+            assert TitledAgedCustomer.objects.count() == 0
+
+        assert Customer.objects.count() == 0
+        assert AgedCustomer.objects.count() == 0
+        assert TitledAgedCustomer.objects.count() == 0
+
+    def test_inherited_model_top_parent_fails(self):
+        AgedCustomer.objects.create(age=99999, name="Methuselah")
+        with pytest.raises(OperationalError) as excinfo:
+            with TableLock(write=[Customer]):
+                # Django automatically follows down to children which aren't
+                # locked
+                Customer.objects.all().delete()
+        assert "was not locked with LOCK TABLES" in str(excinfo.value)
+
+    def test_read_and_write_with_threads(self):
+        to_me = queue.Queue()
+
+        def get_read_lock():
+            with TableLock(read=[Alphabet]):
+                to_me.put("Reading")
+
+        try:
+            # Thread can obtain lock alone
+            other_thread = Thread(target=get_read_lock)
+            other_thread.start()
+            assert to_me.get() == "Reading"
+
+            with TableLock(write=[Alphabet]):
+                other_thread = Thread(target=get_read_lock)
+                other_thread.start()
+
+                # Doesn't lock now
+                with pytest.raises(queue.Empty):
+                    to_me.get(timeout=0.2)
+
+            # Gets now
+            assert to_me.get() == "Reading"
+        finally:
+            other_thread.join()