diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 563c25ac..b44a3847 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -456,7 +456,18 @@ class QuerySet(object): :param *fields: fields to include """ - self._loaded_fields = list(fields) + self._loaded_fields = [] + for field in fields: + if '.' in field: + raise InvalidQueryError('Subfields cannot be used as ' + 'arguments to QuerySet.only') + # Translate field name + field_name = QuerySet._lookup_field(self._document, field)[-1].name + self._loaded_fields.append(field_name) + + # _cls is needed for polymorphism + if self._document._meta.get('allow_inheritance'): + self._loaded_fields += ['_cls'] return self def order_by(self, *keys): diff --git a/tests/queryset.py b/tests/queryset.py index 61cb0537..944baf64 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -258,6 +258,39 @@ class QuerySetTest(unittest.TestCase): BlogPost.drop_collection() + def test_only(self): + """Ensure that QuerySet.only only returns the requested fields. + """ + person = self.Person(name='test', age=25) + person.save() + + obj = self.Person.objects.only('name').get() + self.assertEqual(obj.name, person.name) + self.assertEqual(obj.age, None) + + obj = self.Person.objects.only('age').get() + self.assertEqual(obj.name, None) + self.assertEqual(obj.age, person.age) + + obj = self.Person.objects.only('name', 'age').get() + self.assertEqual(obj.name, person.name) + self.assertEqual(obj.age, person.age) + + # Check polymorphism still works + class Employee(self.Person): + salary = IntField(name='wage') + + employee = Employee(name='test employee', age=40, salary=30000) + employee.save() + + obj = self.Person.objects(id=employee.id).only('age').get() + self.assertTrue(isinstance(obj, Employee)) + + # Check field names are looked up properly + obj = Employee.objects(id=employee.id).only('salary').get() + self.assertEqual(obj.salary, employee.salary) + self.assertEqual(obj.name, None) + def test_find_embedded(self): """Ensure that an embedded document is properly returned from a query. """