Allowed to switch databases for a specific query.

This commit is contained in:
Omer Katz 2014-06-25 18:22:39 +03:00
parent f099dc6a37
commit 5ae588833b
4 changed files with 43 additions and 11 deletions

View File

@ -1,6 +1,5 @@
from mongoengine.common import _import_class
from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
from mongoengine.queryset import QuerySet
__all__ = ("switch_db", "switch_collection", "no_dereference",
@ -162,12 +161,6 @@ class no_sub_classes(object):
return self.cls
class QuerySetNoDeRef(QuerySet):
"""Special no_dereference QuerySet"""
def __dereference(items, max_depth=1, instance=None, name=None):
return items
class query_counter(object):
""" Query_counter context manager to get the number of queries. """

View File

@ -13,11 +13,11 @@ import pymongo
from pymongo.common import validate_read_preference
from mongoengine import signals
from mongoengine.context_managers import switch_db
from mongoengine.common import _import_class
from mongoengine.base.common import get_document
from mongoengine.errors import (OperationError, NotUniqueError,
InvalidQueryError, LookUpError)
from mongoengine.queryset import transform
from mongoengine.queryset.field_list import QueryFieldList
from mongoengine.queryset.visitor import Q, QNode
@ -522,6 +522,19 @@ class BaseQuerySet(object):
return self
def using(self, alias):
"""This method is for controlling which database the QuerySet will be evaluated against if you are using more than one database.
:param alias: The database alias
.. versionadded:: 0.8
"""
with switch_db(self._document, alias) as cls:
collection = cls._get_collection()
return self.clone_into(self.__class__(self._document, collection))
def clone(self):
"""Creates a copy of the current
:class:`~mongoengine.queryset.QuerySet`
@ -1450,6 +1463,7 @@ class BaseQuerySet(object):
# type of this field and use the corresponding
# .to_python(...)
from mongoengine.fields import EmbeddedDocumentField
obj = self._document
for chunk in path.split('.'):
obj = getattr(obj, chunk, None)
@ -1460,6 +1474,7 @@ class BaseQuerySet(object):
if obj and data is not None:
data = obj.to_python(data)
return data
return clean(row)
def _sub_js_fields(self, code):
@ -1468,6 +1483,7 @@ class BaseQuerySet(object):
substituted for the MongoDB name of the field (specified using the
:attr:`name` keyword argument in a field's constructor).
"""
def field_sub(match):
# Extract just the field name, and look up the field objects
field_name = match.group(1).split('.')

View File

@ -155,3 +155,10 @@ class QuerySetNoCache(BaseQuerySet):
queryset = self.clone()
queryset.rewind()
return queryset
class QuerySetNoDeRef(QuerySet):
"""Special no_dereference QuerySet"""
def __dereference(items, max_depth=1, instance=None, name=None):
return items

View File

@ -29,6 +29,7 @@ class QuerySetTest(unittest.TestCase):
def setUp(self):
connect(db='mongoenginetest')
connect(db='mongoenginetest2', alias='test2')
class PersonMeta(EmbeddedDocument):
weight = IntField()
@ -2957,6 +2958,21 @@ class QuerySetTest(unittest.TestCase):
Number.drop_collection()
def test_using(self):
"""Ensure that switching databases for a queryset is possible
"""
class Number2(Document):
n = IntField()
Number2.drop_collection()
for i in xrange(1, 10):
t = Number2(n=i)
t.switch_db('test2')
t.save()
self.assertEqual(len(Number2.objects.using('test2')), 9)
def test_unset_reference(self):
class Comment(Document):
text = StringField()