diff --git a/multidb/pinning.py b/multidb/pinning.py index 6daec5f..b74bcd4 100644 --- a/multidb/pinning.py +++ b/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 diff --git a/multidb/tests/test_all.py b/multidb/tests/test_all.py index 78e6317..91fe0a2 100644 --- a/multidb/tests/test_all.py +++ b/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())