SODA: streaming getCursor() for unbounded result sets (#213)

oracle/soda.py

@@ -345,18 +345,44 @@ def count(self) -> int:
                              "filter": self._filter, "n": n})
         return int(n.getvalue())
 
-    def getDocuments(self) -> "list[SodaDocument]":
-        """Every matched document. Without `.limit(...)` this materialises up to
-        a fixed cap and raises if more match, rather than truncating."""
+    def _run_get_docs(self, skip: int, lim: int, cap: int) -> dict:
         cur = self._connection.cursor()
-        cap = self._limit if self._limit else _DEFAULT_FETCH_CAP
         b = _new_docs_array_binds(cur, cap)
-        b.update(self._in_binds())
-        b["cap"] = cap
+        b.update({"name": self._collection.name, "key": self._key,
+                  "filter": self._filter, "skip": skip, "lim": lim, "cap": cap})
         cur.execute(_GET_DOCS, b)
+        return b
+
+    def getDocuments(self) -> "list[SodaDocument]":
+        """Every matched document. Without `.limit(...)` this materialises up to
+        a fixed cap and raises if more match, rather than truncating; use
+        `getCursor()` to stream without a cap."""
+        cap = self._limit if self._limit else _DEFAULT_FETCH_CAP
+        b = self._run_get_docs(self._skip, self._limit, cap)
         _check_fetch_overflow(b, cap)
         return _docs_from_arrays(b)
 
+    def getCursor(self, batchSize: int = 100):
+        """Stream the matched documents in batches of `batchSize` (offset
+        pagination), so an arbitrarily large result set fetches without the
+        getDocuments cap. Returns an iterator of SodaDocument; honours `skip` /
+        `limit` if set."""
+        batch = max(1, batchSize)
+        yielded, offset = 0, 0
+        while True:
+            if self._limit and yielded >= self._limit:
+                return
+            lim = batch if not self._limit else min(batch, self._limit - yielded)
+            docs = _docs_from_arrays(
+                self._run_get_docs(self._skip + offset, lim, batch))
+            if not docs:
+                return
+            yield from docs
+            yielded += len(docs)
+            offset += len(docs)
+            if len(docs) < lim:               # short batch -> end of result set
+                return
+
     def replaceOne(self, doc) -> bool:
         """Replace the matched document's content. Returns True if one was
         replaced, False if nothing matched."""
@@ -675,16 +701,40 @@ async def count(self) -> int:
                                    "n": n})
         return int(n.getvalue())
 
-    async def getDocuments(self) -> "list[SodaDocument]":
+    async def _run_get_docs(self, skip: int, lim: int, cap: int) -> dict:
         cur = self._connection.cursor()
-        cap = self._limit if self._limit else _DEFAULT_FETCH_CAP
         b = _new_docs_array_binds(cur, cap)
-        b.update(self._in_binds())
-        b["cap"] = cap
+        b.update({"name": self._collection.name, "key": self._key,
+                  "filter": self._filter, "skip": skip, "lim": lim, "cap": cap})
         await cur.execute(_GET_DOCS, b)
+        return b
+
+    async def getDocuments(self) -> "list[SodaDocument]":
+        cap = self._limit if self._limit else _DEFAULT_FETCH_CAP
+        b = await self._run_get_docs(self._skip, self._limit, cap)
         _check_fetch_overflow(b, cap)
         return _docs_from_arrays(b)
 
+    async def getCursor(self, batchSize: int = 100):
+        """Async streaming counterpart to the sync `getCursor`: `async for doc in
+        op.getCursor(): ...`."""
+        batch = max(1, batchSize)
+        yielded, offset = 0, 0
+        while True:
+            if self._limit and yielded >= self._limit:
+                return
+            lim = batch if not self._limit else min(batch, self._limit - yielded)
+            docs = _docs_from_arrays(
+                await self._run_get_docs(self._skip + offset, lim, batch))
+            if not docs:
+                return
+            for doc in docs:
+                yield doc
+            yielded += len(docs)
+            offset += len(docs)
+            if len(docs) < lim:
+                return
+
     async def replaceOne(self, doc) -> bool:
         _, content, _ = _doc_to_bind(doc)
         cur = self._connection.cursor()

tests/test_integration.py

@@ -2366,6 +2366,22 @@ def test_save_upsert(self):
         self.assertEqual(col.find().key(g.key).getOne().getContent()["v"], 99)
         self.assertEqual(col.find().count(), 2)
 
+    def test_streaming_getcursor(self):
+        col = self.soda.createCollection("it_cur")
+        for i in range(25):
+            col.insertOne({"i": i})
+        # streams every document exactly once across batches (no dup / skip)
+        got = [d.getContent()["i"] for d in col.find().getCursor(batchSize=10)]
+        self.assertEqual(sorted(got), list(range(25)))
+        self.assertEqual(len(got), len(set(got)))
+        # honours filter and limit
+        self.assertEqual(
+            sorted(d.getContent()["i"] for d in col.find()
+                   .filter({"i": {"$gte": 20}}).getCursor(batchSize=2)),
+            [20, 21, 22, 23, 24])
+        self.assertEqual(len(list(col.find().limit(7).getCursor(batchSize=3))), 7)
+        self.assertEqual(list(col.find().filter({"i": 999}).getCursor()), [])
+
     def test_indexing_and_data_guide(self):
         col = self.soda.createCollection("it_idx")
         for age in (20, 30, 40):
@@ -3305,6 +3321,13 @@ async def test_async_soda_collections(self):
             self.assertEqual(sg2.key, sg.key)
             self.assertEqual(
                 (await col.find().key(sg.key).getOne()).getContent()["y"], 2)
+            # streaming getCursor on the async path (#213)
+            for i in range(6):
+                await col.insertOne({"z": i})
+            streamed = [d async for d in
+                        col.find().filter({"z": {"$exists": True}})
+                        .getCursor(batchSize=2)]
+            self.assertEqual(len(streamed), 6)
             # indexing + data guide on the async path (#203)
             self.assertIsNone(await col.getDataGuide())
             await col.createIndex({"name": "AIDX",