Add support for count (#397)

* add count

* update poetry

Author: dvora-h

aredis_om/model/model.py

@@ -345,6 +345,7 @@ def __init__(
         limit: int = DEFAULT_PAGE_SIZE,
         page_size: int = DEFAULT_PAGE_SIZE,
         sort_fields: Optional[List[str]] = None,
+        nocontent: bool = False,
     ):
         if not has_redisearch(model.db()):
             raise RedisModelError(
@@ -358,6 +359,7 @@ def __init__(
         self.offset = offset
         self.limit = limit
         self.page_size = page_size
+        self.nocontent = nocontent
 
         if sort_fields:
             self.sort_fields = self.validate_sort_fields(sort_fields)
@@ -377,6 +379,7 @@ def dict(self) -> Dict[str, Any]:
             limit=self.limit,
             expressions=copy(self.expressions),
             sort_fields=copy(self.sort_fields),
+            nocontent=self.nocontent,
         )
 
     def copy(self, **kwargs):
@@ -716,18 +719,23 @@ def resolve_redisearch_query(cls, expression: ExpressionOrNegated) -> str:
 
         return result
 
-    async def execute(self, exhaust_results=True):
+    async def execute(self, exhaust_results=True, return_raw_result=False):
         args = ["ft.search", self.model.Meta.index_name, self.query, *self.pagination]
         if self.sort_fields:
             args += self.resolve_redisearch_sort_fields()
 
+        if self.nocontent:
+            args.append("NOCONTENT")
+
         # Reset the cache if we're executing from offset 0.
         if self.offset == 0:
             self._model_cache.clear()
 
         # If the offset is greater than 0, we're paginating through a result set,
         # so append the new results to results already in the cache.
         raw_result = await self.model.db().execute_command(*args)
+        if return_raw_result:
+            return raw_result
         count = raw_result[0]
         results = self.model.from_redis(raw_result)
         self._model_cache += results
@@ -759,6 +767,11 @@ async def first(self):
             raise NotFoundError()
         return results[0]
 
+    async def count(self):
+        query = self.copy(offset=0, limit=0, nocontent=True)
+        result = await query.execute(exhaust_results=True, return_raw_result=True)
+        return result[0]
+
     async def all(self, batch_size=DEFAULT_PAGE_SIZE):
         if batch_size != self.page_size:
             query = self.copy(page_size=batch_size, limit=batch_size)
@@ -1175,7 +1188,7 @@ def validate_primary_key(cls):
         if primary_keys == 0:
             raise RedisModelError("You must define a primary key for the model")
         elif primary_keys == 2:
-            cls.__fields__.pop('pk')
+            cls.__fields__.pop("pk")
         elif primary_keys > 2:
             raise RedisModelError("You must define only one primary key for a model")
 

poetry.lock

@@ -706,7 +706,6 @@ class Address(m.BaseHashModel):
 
 @py_test_mark_asyncio
 async def test_primary_key_model_error(m):
-
     class Customer(m.BaseHashModel):
         id: int = Field(primary_key=True, index=True)
         first_name: str = Field(primary_key=True, index=True)
@@ -715,18 +714,19 @@ class Customer(m.BaseHashModel):
 
     await Migrator().run()
 
-    with pytest.raises(RedisModelError, match="You must define only one primary key for a model"):
+    with pytest.raises(
+        RedisModelError, match="You must define only one primary key for a model"
+    ):
         _ = Customer(
             id=0,
             first_name="Mahmoud",
             last_name="Harmouch",
-            bio="Python developer, wanna work at Redis, Inc."
+            bio="Python developer, wanna work at Redis, Inc.",
         )
 
 
 @py_test_mark_asyncio
 async def test_primary_pk_exists(m):
-
     class Customer1(m.BaseHashModel):
         id: int
         first_name: str
@@ -745,10 +745,10 @@ class Customer2(m.BaseHashModel):
         id=0,
         first_name="Mahmoud",
         last_name="Harmouch",
-        bio="Python developer, wanna work at Redis, Inc."
+        bio="Python developer, wanna work at Redis, Inc.",
     )
 
-    assert 'pk' in customer.__fields__
+    assert "pk" in customer.__fields__
 
     customer = Customer2(
         id=1,
@@ -757,4 +757,19 @@ class Customer2(m.BaseHashModel):
         bio="This is member 2 who can be quite anxious until you get to know them.",
     )
 
-    assert 'pk' not in customer.__fields__
+    assert "pk" not in customer.__fields__
+
+
+@py_test_mark_asyncio
+async def test_count(members, m):
+    # member1, member2, member3 = members
+    actual_count = await m.Member.find(
+        (m.Member.first_name == "Andrew") & (m.Member.last_name == "Brookins")
+        | (m.Member.last_name == "Smith")
+    ).count()
+    assert actual_count == 2
+
+    actual_count = await m.Member.find(
+        m.Member.first_name == "Kim", m.Member.last_name == "Brookins"
+    ).count()
+    assert actual_count == 1