"use_slave" decorator & context manager

multidb/pinning.py

@@ -6,7 +6,7 @@
 
 __all__ = ['this_thread_is_pinned', 'pin_this_thread', 'unpin_this_thread',
            'use_master', 'db_write', 'set_db_write_for_this_thread',
-           'unset_db_write_for_this_thread',
+           'use_slave', 'unset_db_write_for_this_thread',
            'this_thread_has_db_write_set',
            'set_db_write_for_this_thread_if_needed']
 
@@ -95,6 +95,23 @@ def __exit__(self, type, value, tb):
 use_master = UseMaster()
 
 
+class UseSlave(UseMaster):
+    """A contextmanager/decorator to use the slave database."""
+    "Use this in cases where the usual behavior would be to pin to master,"
+    "such as when the request method is POST, but you know you're not doing any writing."
+    old = False
+
+    def __enter__(self):
+        self.old = this_thread_is_pinned()
+        unpin_this_thread()
+
+    def __exit__(self, type, value, tb):
+        if self.old:
+            pin_this_thread()
+
+use_slave = UseSlave()
+
+
 def mark_as_write(response):
     """Mark a response as having done a DB write."""
     response._db_write = True

multidb/tests/test_all.py

@@ -11,7 +11,7 @@
 from multidb.conf import settings
 from multidb.middleware import PinningRouterMiddleware
 from multidb.pinning import (this_thread_is_pinned, pin_this_thread,
-                             unpin_this_thread, use_master, db_write)
+                             unpin_this_thread, use_master, use_slave, db_write)
 
 
 def expire_cookies(cookies):
@@ -327,3 +327,48 @@ def test_context_manager_exception(self):
                 self.assertTrue(this_thread_is_pinned())
                 raise ValueError
         self.assertFalse(this_thread_is_pinned())
+
+    def test_slave_decorator(self):
+
+        @use_slave
+        def check():
+            self.assertFalse(this_thread_is_pinned())
+
+        pin_this_thread()
+        self.assertTrue(this_thread_is_pinned())
+        check()
+        self.assertTrue(this_thread_is_pinned())
+
+    def test_slave_decorator_resets(self):
+
+        @use_slave
+        def check():
+            self.assertFalse(this_thread_is_pinned())
+
+        unpin_this_thread()
+        self.assertFalse(this_thread_is_pinned())
+        check()
+        self.assertFalse(this_thread_is_pinned())
+
+    def test_slave_context_manager(self):
+        pin_this_thread()
+        self.assertTrue(this_thread_is_pinned())
+        with use_slave:
+            self.assertFalse(this_thread_is_pinned())
+        self.assertTrue(this_thread_is_pinned())
+
+    def test_slave_context_manager_resets(self):
+        unpin_this_thread()
+        self.assertFalse(this_thread_is_pinned())
+        with use_slave:
+            self.assertFalse(this_thread_is_pinned())
+        self.assertFalse(this_thread_is_pinned())
+
+    def test_slave_context_manager_exception(self):
+        pin_this_thread()
+        self.assertTrue(this_thread_is_pinned())
+        with self.assertRaises(ValueError):
+            with use_slave:
+                self.assertFalse(this_thread_is_pinned())
+                raise ValueError
+        self.assertTrue(this_thread_is_pinned())