forked from jbalogh/django-multidb-router
/
pinning.py
59 lines (40 loc) · 1.36 KB
/
pinning.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
"""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()
def this_thread_is_pinned():
"""Return whether the current thread should send all its reads to the
master DB."""
return getattr(_locals, 'pinned', False)
def pin_this_thread():
"""Mark this thread as "stuck" to the master for all DB access."""
_locals.pinned = True
def unpin_this_thread():
"""Unmark this thread as "stuck" to the master for all DB access.
If the thread wasn't marked, do nothing.
"""
try:
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()