added count query type to builder

red-crown · Mar 7, 2016 · d2d69ec · d2d69ec
1 parent 7609a6b
commit d2d69ec
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/rcquerybuilder/builder.py b/rcquerybuilder/builder.py
@@ -107,6 +107,11 @@ def insert(self):
 
         return self
 
+    def count(self):
+        self.query['type'] = QueryTypes.TYPE_COUNT
+
+        return self
+
     def build(self, **kwargs):
         query = self.query
         query['query'] = self.expression.query
@@ -168,6 +173,10 @@ def __init__(self, collection, query=(), options=()):
     def execute(self):
         options = self.options
 
+        if self.query['type'] == QueryTypes.TYPE_COUNT:
+            return self.collection.find(self.query['query'],
+                                        self.query['select'] if self.query.setdefault('select', []) else None).count()
+
         # FIND
         if self.query['type'] == QueryTypes.TYPE_FIND:
             cursor = self.collection.find(

diff --git a/tests/test_builder.py b/tests/test_builder.py
@@ -1,6 +1,7 @@
 import unittest
 
 from rcquerybuilder.builder import Builder
+from rcquerybuilder.builder import QueryTypes
 
 
 class BasicTestSuite(unittest.TestCase):
@@ -59,6 +60,20 @@ def test_insert_query(self):
                 'age': 21,
                 'attributes': [0, 1, 2, 3]} == insert_query
 
+    def test_count_query(self):
+        qb = Builder(collection=None)
+
+        qb.count() \
+            .field('name').ne('awesome') \
+            .field('age').lte(45)
+
+        query_list = qb.get_query_list()
+
+        assert {'name': {'$ne': 'awesome'},
+                'age': {'$lte': 45}} == query_list
+
+        assert qb.build().query['type'] == QueryTypes.TYPE_COUNT
+
 
 if __name__ == '__main__':
     unittest.main()