diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index c12f9553..f64b0075 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -974,8 +974,33 @@ class BaseQuerySet(object): son_data = json_util.loads(json_data) return [self._document._from_son(data) for data in son_data] - # JS functionality + def aggregate(self, *pipeline, **kwargs): + """ + Perform a aggreggate function based in your queryset params + :param pipeline: list of agreggation commands, + see: http://docs.mongodb.org/manual/core/aggregation-pipeline/ + .. versionadded:: 0.9 + """ + initial_pipeline = [] + + if self._query: + initial_pipeline.append({'$match': self._query}) + + if self._ordering: + initial_pipeline.append({'$sort': dict(self._ordering)}) + + if self._limit is not None: + initial_pipeline.append({'$limit': self._limit}) + + if self._skip is not None: + initial_pipeline.append({'$skip': self._skip}) + + pipeline = initial_pipeline + list(pipeline) + + return self._collection.aggregate(pipeline, cursor={}, **kwargs) + + # JS functionality def map_reduce(self, map_f, reduce_f, output, finalize_f=None, limit=None, scope=None): """Perform a map/reduce query using the current query spec diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 1d3e8048..0eadd7ed 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -4341,6 +4341,55 @@ class QuerySetTest(unittest.TestCase): self.assertTrue(Person.objects._has_data(), 'Cursor has data and returned False') + def test_queryset_aggregation_framework(self): + class Person(Document): + name = StringField() + age = IntField() + + Person.drop_collection() + + p1 = Person(name="Isabella Luanna", age=16) + p1.save() + + p2 = Person(name="Wilson Junior", age=21) + p2.save() + + p3 = Person(name="Sandra Mara", age=37) + p3.save() + + data = Person.objects(age__lte=22).aggregate( + {'$project': {'name': {'$toUpper': '$name'}}} + ) + + self.assertEqual(list(data), [ + {'_id': p1.pk, 'name': "ISABELLA LUANNA"}, + {'_id': p2.pk, 'name': "WILSON JUNIOR"} + ]) + + data = Person.objects(age__lte=22).order_by('-name').aggregate( + {'$project': {'name': {'$toUpper': '$name'}}} + ) + + self.assertEqual(list(data), [ + {'_id': p2.pk, 'name': "WILSON JUNIOR"}, + {'_id': p1.pk, 'name': "ISABELLA LUANNA"} + ]) + + data = Person.objects( + age__gte=17, age__lte=40).order_by('-age').aggregate( + {'$group': { + '_id': None, + 'total': {'$sum': 1}, + 'avg': {'$avg': '$age'} + } + } + + ) + + self.assertEqual(list(data), [ + {'_id': None, 'avg': 29, 'total': 2} + ]) + if __name__ == '__main__': unittest.main()