From 5cac52720ca256c15919f0ae9e4d0e4ef9db3e44 Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Wed, 27 Jan 2010 15:57:11 +0000 Subject: [PATCH 1/4] Fixed querying on ReferenceFields using primary key --- mongoengine/fields.py | 15 +++++++-------- tests/fields.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/mongoengine/fields.py b/mongoengine/fields.py index 13fa6689..beb8ae00 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -209,20 +209,19 @@ class ReferenceField(BaseField): return super(ReferenceField, self).__get__(instance, owner) def to_mongo(self, document): - if isinstance(document, (str, unicode, pymongo.objectid.ObjectId)): - # document may already be an object id - id_ = document - else: + id_field_name = self.document_type._meta['id_field'] + id_field = self.document_type._fields[id_field_name] + + if isinstance(document, Document): # We need the id from the saved object to create the DBRef id_ = document.id if id_ is None: raise ValidationError('You can only reference documents once ' 'they have been saved to the database') + else: + id_ = document - # id may be a string rather than an ObjectID object - if not isinstance(id_, pymongo.objectid.ObjectId): - id_ = pymongo.objectid.ObjectId(id_) - + id_ = id_field.to_mongo(id_) collection = self.document_type._meta['collection'] return pymongo.dbref.DBRef(collection, id_) diff --git a/tests/fields.py b/tests/fields.py index affb0c9b..0ebb143d 100644 --- a/tests/fields.py +++ b/tests/fields.py @@ -259,6 +259,40 @@ class FieldTest(unittest.TestCase): User.drop_collection() BlogPost.drop_collection() + def test_reference_query_conversion(self): + """Ensure that ReferenceFields can be queried using objects and values + of the type of the primary key of the referenced object. + """ + class Member(Document): + user_num = IntField(primary_key=True) + + class BlogPost(Document): + title = StringField() + author = ReferenceField(Member) + + Member.drop_collection() + BlogPost.drop_collection() + + m1 = Member(user_num=1) + m1.save() + m2 = Member(user_num=2) + m2.save() + + post1 = BlogPost(title='post 1', author=m1) + post1.save() + + post2 = BlogPost(title='post 2', author=m2) + post2.save() + + post = BlogPost.objects(author=m1.id).first() + self.assertEqual(post.id, post1.id) + + post = BlogPost.objects(author=m2.id).first() + self.assertEqual(post.id, post2.id) + + Member.drop_collection() + BlogPost.drop_collection() + if __name__ == '__main__': unittest.main() From 8d953f0bcb2484c1dce937efac07ae949f732052 Mon Sep 17 00:00:00 2001 From: Florian Schlachter Date: Sat, 30 Jan 2010 21:12:46 +0100 Subject: [PATCH 2/4] Added get_or_create-method --- mongoengine/queryset.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 6c8572c6..c40523ae 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -10,6 +10,9 @@ __all__ = ['queryset_manager', 'Q', 'InvalidQueryError', # The maximum number of items to display in a QuerySet.__repr__ REPR_OUTPUT_SIZE = 20 +class MultipleObjectsReturned(Exception): + pass + class InvalidQueryError(Exception): pass @@ -292,6 +295,23 @@ class QuerySet(object): return mongo_query + def get_or_create(self, **kwargs): + """Retreive unique object or create with paras, if it doesn't exist + """ + dataset = self.filter(**kwargs) + cnt = dataset.count() + if cnt == 0: + if kwargs.has_key('defaults'): + kwargs.update(kwargs.get('defaults')) + del kwargs['defaults'] + doc = self._document(**kwargs) + doc.save() + return doc + elif cnt == 1: + return dataset.first() + else: + raise MultipleObjectsReturned(u'%d items returned, instead of 1' % cnt) + def first(self): """Retrieve the first object matching the query. """ From b3cc2f990a9ad612f66f4f7def23972074d77292 Mon Sep 17 00:00:00 2001 From: Florian Schlachter Date: Sat, 30 Jan 2010 22:01:43 +0100 Subject: [PATCH 3/4] improved get_or_create --- mongoengine/queryset.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index c40523ae..bec6da2c 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -296,14 +296,16 @@ class QuerySet(object): return mongo_query def get_or_create(self, **kwargs): - """Retreive unique object or create with paras, if it doesn't exist + """Retreive unique object or create, if it doesn't exist """ + defaults = kwargs.get('defaults', {}) + if kwargs.has_key('defaults'): + del kwargs['defaults'] + dataset = self.filter(**kwargs) cnt = dataset.count() if cnt == 0: - if kwargs.has_key('defaults'): - kwargs.update(kwargs.get('defaults')) - del kwargs['defaults'] + kwargs.update(defaults) doc = self._document(**kwargs) doc.save() return doc From 7d6e117f687462709e19fd8ce343ac3534d3a825 Mon Sep 17 00:00:00 2001 From: Florian Schlachter Date: Sun, 31 Jan 2010 01:11:37 +0100 Subject: [PATCH 4/4] added get-method to fetch exactly one document from the collection. catching pymongo's ObjectId-errors and raising mongoengine's ValidationError instead. --- mongoengine/base.py | 6 ++++-- mongoengine/queryset.py | 16 +++++++++++++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/mongoengine/base.py b/mongoengine/base.py index 83fd34ee..8c038c8c 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -6,7 +6,6 @@ import pymongo class ValidationError(Exception): pass - class BaseField(object): """A base class for fields in a MongoDB document. Instances of this class may be added to subclasses of `Document` to define a document's schema. @@ -76,7 +75,10 @@ class ObjectIdField(BaseField): def to_mongo(self, value): if not isinstance(value, pymongo.objectid.ObjectId): - return pymongo.objectid.ObjectId(str(value)) + try: + return pymongo.objectid.ObjectId(str(value)) + except Exception, e: + raise ValidationError(e.message) return value def prepare_query_value(self, value): diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index bec6da2c..69bec002 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -10,6 +10,10 @@ __all__ = ['queryset_manager', 'Q', 'InvalidQueryError', # The maximum number of items to display in a QuerySet.__repr__ REPR_OUTPUT_SIZE = 20 +class DoesNotExist(Exception): + pass + + class MultipleObjectsReturned(Exception): pass @@ -313,7 +317,17 @@ class QuerySet(object): return dataset.first() else: raise MultipleObjectsReturned(u'%d items returned, instead of 1' % cnt) - + + def get(self, **kwargs): + dataset = self.filter(**kwargs) + cnt = dataset.count() + if cnt == 1: + return dataset.first() + elif cnt > 1: + raise MultipleObjectsReturned(u'%d items returned, instead of 1' % cnt) + else: + raise DoesNotExist('Document not found') + def first(self): """Retrieve the first object matching the query. """