From 057493c809ac5a6685d0b36015b8cc7722fc2b32 Mon Sep 17 00:00:00 2001
From: Jaremy Creechley <creechley@gmail.com>
Date: Tue, 29 Aug 2023 20:25:47 -0700
Subject: [PATCH] query iterator using items is breaks when the DS isn't
 blocking

---
 datastore/query.nim                  | 22 +++++++++++++++++-----
 tests/datastore/querycommontests.nim |  6 +-----
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/datastore/query.nim b/datastore/query.nim
index bc6435a..1d98e3b 100644
--- a/datastore/query.nim
+++ b/datastore/query.nim
@@ -29,11 +29,23 @@ type
     next*: GetNext
     dispose*: IterDispose
 
-proc collectAllQueries*(q: QueryIter)  =
-  var qr: Future[?!QueryResponse]
-  while not q.finished:
-    qr = q.next()
-    yield qr
+proc waitForAllQueryResults*(qi: Future[?!QueryIter]): Future[?!seq[QueryResponse]] {.async.} =
+  ## for large blocks this would be *expensive*
+  var res: seq[QueryResponse]
+  without iter =? (await qi), err:
+    return failure err
+
+  while not iter.finished:
+    let val = await iter.next()
+    if val.isOk():
+      let qr = val.tryGet()
+      if qr.key.isSome:
+        res.add qr
+    else:
+      return failure val.error()
+  
+  await iter.dispose()
+  return success res
 
 proc defaultDispose(): Future[?!void] {.upraises: [], gcsafe, async.} =
   return success()
diff --git a/tests/datastore/querycommontests.nim b/tests/datastore/querycommontests.nim
index 86092fc..b842ab5 100644
--- a/tests/datastore/querycommontests.nim
+++ b/tests/datastore/querycommontests.nim
@@ -68,10 +68,7 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
     (await ds.put(key3, val3)).tryGet
 
     let
-      iter = (await ds.query(q)).tryGet
-      res = (await allFinished(toSeq(iter)))
-        .mapIt( it.read.tryGet )
-        .filterIt( it.key.isSome )
+      res = tryGet(await ds.query(q).waitForAllQueryResults())
 
     check:
       res.len == 3
@@ -84,7 +81,6 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
       res[2].key.get == key3
       res[2].data.len == 0
 
-    (await iter.dispose()).tryGet
 
   test "Key should not query parent":
     let