Minimal implementation of aggregate

TurboGears · Nov 30, 2018 · e42a43d · e42a43d
1 parent 433b94a
commit e42a43d
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 2 deletions.
diff --git a/ming/mim.py b/ming/mim.py
@@ -669,6 +669,27 @@ def bulk_write(self, requests, ordered=True,
                     "MIM currently doesn't support %s operations" % type(step)
                 )
 
+    def aggregate(self, pipeline, **kwargs):
+        steps = {}
+        for step in pipeline:
+            if set(step.keys()) & set(steps.keys()):
+                raise ValueError(
+                    'MIM currently supports a single step per type. Duplicate %s' % step
+                )
+            if set(step.keys()) - {'$match', '$project', '$sort', '$limit'}:
+                raise ValueError(
+                    'MIM currently only supports $match,$project,$sort,$limit steps.'
+                )
+            steps.update(step)
+
+        sort = steps.get('$sort', None)
+        if isinstance(sort, (bson.SON, dict)):
+            sort = sort.items()
+        return self.find(filter=steps.get('$match', {}),
+                         sort=sort,
+                         projection=steps.get('$project', None),
+                         limit=steps.get('$limit', None))
+
 
 class Cursor(object):
     def __init__(self, collection, _iterator_gen,

diff --git a/ming/tests/odm/test_mapper.py b/ming/tests/odm/test_mapper.py
@@ -247,7 +247,7 @@ def test_imap(self):
         self.session.expunge(doc)
         self.session.expunge(doc)
 
-    @patch('pymongo.collection.Collection.aggregate')
+    @patch('ming.mim.Collection.aggregate')
     def test_aggregate(self, pymongo_aggregate):
         self.Basic.query.aggregate([])
         assert pymongo_aggregate.called

diff --git a/ming/tests/test_mim.py b/ming/tests/test_mim.py
@@ -9,6 +9,7 @@
 from nose import SkipTest
 from mock import patch
 
+
 class TestDatastore(TestCase):
 
     def setUp(self):
@@ -973,4 +974,40 @@ def test_update_one(self):
         ])
 
         data = sorted([a['dme-o'] for a in coll.find({'dme-o': {'$exists': True}})])
-        self.assertEqual(data, [1, 2])
+        self.assertEqual(data, [1, 2])
+
+
+class TestAggregate(TestCase):
+
+    def setUp(self):
+        self.bind = create_datastore('mim:///testdb')
+        self.bind.conn.drop_all()
+        self.bind.db.coll.insert({'_id':'foo', 'a':2, 'c':[1,2,3],
+                                  'z': {'egg': 'spam', 'spam': 'egg'}})
+        for r in range(4):
+            self.bind.db.rcoll.insert({'_id':'r%s' % r, 'd':r})
+
+    def test_aggregate_match(self):
+        res = self.bind.db.rcoll.aggregate([{'$match': {'d': {'$lt': 2}}}])
+        self.assertEqual(len(list(res)), 2)
+
+    def test_aggregate_match_sort(self):
+        res = self.bind.db.rcoll.aggregate([{'$match': {'d': {'$lt': 2}}},
+                                            {'$sort': {'d': -1}}])
+        res = list(res)
+        self.assertEqual(len(res), 2)
+        self.assertEqual(res[0]['d'], 1)
+
+        res = self.bind.db.rcoll.aggregate([{'$match': {'d': {'$lt': 2}}},
+                                            {'$sort': {'d': 1}}])
+        res = list(res)
+        self.assertEqual(len(res), 2)
+        self.assertEqual(res[0]['d'], 0)
+
+    def test_aggregate_match_sort_limit(self):
+        res = self.bind.db.rcoll.aggregate([{'$match': {'d': {'$lt': 2}}},
+                                            {'$sort': {'d': -1}},
+                                            {'$limit': 1}])
+        res = list(res)
+        self.assertEqual(len(res), 1)
+        self.assertEqual(res[0]['d'], 1)