provisionally raising an error if count_documents exceeds its max

astrapy/idiomatic/collection.py

@@ -303,7 +303,11 @@ def count_documents(
     ) -> int:
         cd_response = self._astra_db_collection.count_documents(filter=filter)
         if "count" in cd_response.get("status", {}):
-            return cd_response["status"]["count"]  # type: ignore[no-any-return]
+            count: int = cd_response["status"]["count"]
+            if cd_response["status"].get("moreData", False):
+                raise ValueError(f"Document count exceeds {count}")
+            else:
+                return count
         else:
             raise ValueError(
                 "Could not complete a count_documents operation. "
@@ -810,7 +814,11 @@ async def count_documents(
     ) -> int:
         cd_response = await self._astra_db_collection.count_documents(filter=filter)
         if "count" in cd_response.get("status", {}):
-            return cd_response["status"]["count"]  # type: ignore[no-any-return]
+            count: int = cd_response["status"]["count"]
+            if cd_response["status"].get("moreData", False):
+                raise ValueError(f"Document count exceeds {count}")
+            else:
+                return count
         else:
             raise ValueError(
                 "Could not complete a count_documents operation. "

tests/idiomatic/integration/test_dml_async.py

@@ -46,6 +46,17 @@ async def test_collection_count_documents_async(
         assert await async_empty_collection.count_documents(filter={}) == 3
         assert await async_empty_collection.count_documents(filter={"group": "A"}) == 2
 
+    @pytest.mark.describe("test of overflowing collection count_documents, async")
+    async def test_collection_overflowing_count_documents_async(
+        self,
+        async_empty_collection: AsyncCollection,
+    ) -> None:
+        await async_empty_collection.insert_many([{"a": i} for i in range(999)])
+        assert await async_empty_collection.count_documents(filter={}) == 999
+        await async_empty_collection.insert_many([{"b": i} for i in range(2)])
+        with pytest.raises(ValueError):
+            assert await async_empty_collection.count_documents(filter={})
+
     @pytest.mark.describe("test of collection insert_one, async")
     async def test_collection_insert_one_async(
         self,

tests/idiomatic/integration/test_dml_sync.py

@@ -42,6 +42,17 @@ def test_collection_count_documents_sync(
         assert sync_empty_collection.count_documents(filter={}) == 3
         assert sync_empty_collection.count_documents(filter={"group": "A"}) == 2
 
+    @pytest.mark.describe("test of overflowing collection count_documents, sync")
+    def test_collection_overflowing_count_documents_sync(
+        self,
+        sync_empty_collection: Collection,
+    ) -> None:
+        sync_empty_collection.insert_many([{"a": i} for i in range(999)])
+        assert sync_empty_collection.count_documents(filter={}) == 999
+        sync_empty_collection.insert_many([{"b": i} for i in range(2)])
+        with pytest.raises(ValueError):
+            assert sync_empty_collection.count_documents(filter={})
+
     @pytest.mark.describe("test of collection insert_one, sync")
     def test_collection_insert_one_sync(
         self,