diff --git a/oracle/soda.py b/oracle/soda.py index bb82a7d..380bc37 100644 --- a/oracle/soda.py +++ b/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() diff --git a/tests/test_integration.py b/tests/test_integration.py index 327afb7..16bfd29 100644 --- a/tests/test_integration.py +++ b/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",