Permalink
Browse files

Add use_master context decorator.

Adds `multidb.pinning.use_master`, which is both a context manager and a
decorator, i.e., you can do:

    >>> with use_master:
    ...     do_stuff()

or you can do:

    >>> @use_master
    ... def do_stuff():
  • Loading branch information...
1 parent 25ca0c6 commit c5286208ac6c95b7f39a961faec5ef82958c0bef James Socol committed Mar 11, 2011
Showing with 96 additions and 1 deletion.
  1. +21 −0 README.rst
  2. +29 −0 multidb/pinning.py
  3. +46 −1 multidb/tests/__init__.py
View
21 README.rst
@@ -86,3 +86,24 @@ If you need to change the name of the cookie, use the ``MULTIDB_PINNING_COOKIE``
setting::
MULTIDB_PINNING_COOKIE = 'multidb_pin_writes'
+
+
+``use_master``
+==============
+
+``multidb.pinning.use_master`` is both a context manager and a decorator for
+wrapping code to use the master database. You can use it as a context manager::
+
+ from multidb.pinning import use_master
+
+ with use_master:
+ touch_the_database()
+ touch_another_database()
+
+or as a decorator::
+
+ from multidb.pinning import use_master
+
+ @use_master
+ def func(*args, **kw):
+ """Touches the master database."""
View
29 multidb/pinning.py
@@ -1,9 +1,14 @@
"""An encapsulated thread-local variable that indicates whether future DB
writes should be "stuck" to the master."""
+from functools import wraps
import threading
+__all__ = ['this_thread_is_pinned', 'pin_this_thread', 'unpin_this_thread',
+ 'use_master']
+
+
_locals = threading.local()
@@ -28,3 +33,27 @@ def unpin_this_thread():
del _locals.pinned
except AttributeError:
pass
+
+
+class UseMaster(object):
+ """A contextmanager/decorator to use the master database."""
+ old = False
+
+ def __call__(self, func):
+ @wraps(func)
+ def decorator(*args, **kw):
+ with self:
+ return func(*args, **kw)
+ return decorator
+
+ def __enter__(self):
+ self.old = this_thread_is_pinned()
+ pin_this_thread()
+
+ def __exit__(self, type, value, tb):
+ if not self.old:
+ unpin_this_thread()
+ if any((type, value, tb)):
+ raise type, value, tb
+
+use_master = UseMaster()
View
47 multidb/tests/__init__.py
@@ -8,7 +8,7 @@
from multidb.middleware import (PINNING_COOKIE, PINNING_SECONDS,
PinningRouterMiddleware)
from multidb.pinning import (this_thread_is_pinned,
- pin_this_thread, unpin_this_thread)
+ pin_this_thread, unpin_this_thread, use_master)
class UnpinningTestCase(TestCase):
@@ -115,3 +115,48 @@ def test_process_response(self):
assert PINNING_COOKIE in response.cookies
eq_(response.cookies[PINNING_COOKIE]['max-age'],
PINNING_SECONDS)
+
+
+class ContextDecoratorTests(TestCase):
+ def test_decorator(self):
+ @use_master
+ def check():
+ assert this_thread_is_pinned()
+ unpin_this_thread()
+ assert not this_thread_is_pinned()
+ check()
+ assert not this_thread_is_pinned()
+
+ def test_decorator_resets(self):
+ @use_master
+ def check():
+ assert this_thread_is_pinned()
+ pin_this_thread()
+ assert this_thread_is_pinned()
+ check()
+ assert this_thread_is_pinned()
+
+ def test_context_manager(self):
+ unpin_this_thread()
+ assert not this_thread_is_pinned()
+ with use_master:
+ assert this_thread_is_pinned()
+ assert not this_thread_is_pinned()
+
+ def text_context_manager_resets(self):
+ pin_this_thread()
+ assert this_thread_is_pinned()
+ with use_master:
+ assert this_thread_is_pinned()
+ assert this_thread_is_pinned()
+
+ def test_context_manager_exception(self):
+ unpin_this_thread()
+ try:
+ assert not this_thread_is_pinned()
+ with use_master:
+ assert this_thread_is_pinned()
+ raise ValueError
+ except ValueError:
+ pass
+ assert not this_thread_is_pinned()

0 comments on commit c528620

Please sign in to comment.