Count limits

docs/changelog.rst

@@ -1,6 +1,13 @@
 Changelog
 =========
 
+UNRELEASED
+----------
+
+Release Date: UNRELEASED
+
+* Added the `limit` parameter to :py:meth:`aiodynamo.client.Client.count`
+
 21.8
 ----
 

src/aiodynamo/client.py

@@ -271,13 +271,15 @@ async def count(
         start_key: Optional[Dict[str, Any]] = None,
         filter_expression: Optional[Condition] = None,
         index: Optional[str] = None,
+        limit: Optional[int] = None,
     ) -> int:
         return await self.client.count(
             self.name,
             key_condition,
             start_key=start_key,
             filter_expression=filter_expression,
             index=index,
+            limit=limit,
         )
 
     async def update_item(
@@ -745,6 +747,7 @@ async def count(
         start_key: Optional[Dict[str, Any]] = None,
         filter_expression: Optional[Condition] = None,
         index: Optional[str] = None,
+        limit: Optional[int] = None,
     ) -> int:
         params = Parameters()
 
@@ -764,7 +767,7 @@ async def count(
         payload.update(params.to_request_payload())
 
         count_sum = 0
-        async for result in self._depaginate("Query", payload):
+        async for result in self._depaginate("Query", payload, limit):
             count_sum += result["Count"]
         return count_sum
 
@@ -907,6 +910,7 @@ async def _depaginate(
         task: Optional[asyncio.Task[Dict[str, Any]]] = asyncio.create_task(
             self.send_request(action=action, payload=payload)
         )
+        is_count = payload.get("Select") == "COUNT"
         try:
             while task:
                 result = await task
@@ -919,7 +923,10 @@ async def _depaginate(
                     task = None
                 else:
                     if limit is not None:
-                        limit -= len(result["Items"])
+                        consumed: int = (
+                            result["Count"] if is_count else len(result["Items"])
+                        )
+                        limit -= consumed
                         if limit > 0:
                             payload["Limit"] = limit
                         else:

tests/integration/test_client.py

@@ -92,6 +92,30 @@ async def test_count(client: Client, table: TableName):
     )
 
 
+async def test_count_with_limit(client: Client, high_throughput_table: TableName):
+    big = "x" * 20_000
+    hk = HashKey("h", "h")
+
+    assert await client.count(high_throughput_table, hk, limit=1) == 0
+
+    await client.put_item(high_throughput_table, {"h": "h", "r": "0", "big": big})
+
+    assert await client.count(high_throughput_table, hk, limit=1) == 1
+
+    await client.put_item(high_throughput_table, {"h": "h", "r": "1", "big": big})
+
+    assert await client.count(high_throughput_table, hk, limit=1) == 1
+
+    await asyncio.gather(
+        *(
+            client.put_item(high_throughput_table, {"h": "h", "r": str(i), "big": big})
+            for i in range(2, 100)
+        )
+    )
+
+    assert await client.count(high_throughput_table, hk, limit=90) == 90
+
+
 async def test_update_item(client: Client, table: TableName):
     item = {
         "h": "hkv",