Skip to content

Commit

Permalink
Use a global variable in read_preference context manager if outside o…
Browse files Browse the repository at this point in the history
…f request context.
  • Loading branch information
thomasst committed Jun 26, 2013
1 parent 0a568d8 commit 8c7b4d3
Showing 1 changed file with 22 additions and 5 deletions.
27 changes: 22 additions & 5 deletions flask_mongoengine/context_managers.py
@@ -1,13 +1,30 @@
from flask import g
from flask import g, has_request_context
from mongoengine.queryset.queryset import QuerySet

orig_cursor_args = QuerySet._cursor_args.fget

# In case we're not in a request context
_read_preference = None

def _get_read_preference():
if has_request_context():
return getattr(g, 'read_preference', None)
else:
return _read_preference

def _set_read_preference(val):
if has_request_context():
g.read_preference = val
else:
_read_preference = val


def _patched_cursor_args(self):
cursor_args = orig_cursor_args(self)
if getattr(g, 'read_preference', None) is not None:
read_preference = _get_read_preference()
if read_preference is not None:
if not 'read_preference' in cursor_args:
cursor_args['read_preference'] = g.read_preference
cursor_args['read_preference'] = read_preference
del cursor_args['slave_okay']
return cursor_args

Expand All @@ -18,7 +35,7 @@ def __init__(self, read_preference):
self.read_preference = read_preference

def __enter__(self):
g.read_preference = self.read_preference
_set_read_preference(self.read_preference)

def __exit__(self, t, value, traceback):
g.read_preference = None
_set_read_preference(None)

0 comments on commit 8c7b4d3

Please sign in to comment.