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

---
 datastore/query.nim                  | 10 +++++---
 tests/datastore/querycommontests.nim | 36 +++++++---------------------
 2 files changed, 16 insertions(+), 30 deletions(-)

diff --git a/datastore/query.nim b/datastore/query.nim
index 1d98e3b..e405556 100644
--- a/datastore/query.nim
+++ b/datastore/query.nim
@@ -29,11 +29,9 @@ type
     next*: GetNext
     dispose*: IterDispose
 
-proc waitForAllQueryResults*(qi: Future[?!QueryIter]): Future[?!seq[QueryResponse]] {.async.} =
+proc waitForAllQueryResults*(iter: 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()
@@ -47,6 +45,12 @@ proc waitForAllQueryResults*(qi: Future[?!QueryIter]): Future[?!seq[QueryRespons
   await iter.dispose()
   return success res
 
+proc waitForAllQueryResults*(qi: Future[?!QueryIter]): Future[?!seq[QueryResponse]] {.async.} =
+  without iter =? (await qi), err:
+    return failure err
+  await waitForAllQueryResults(iter)
+
+
 proc defaultDispose(): Future[?!void] {.upraises: [], gcsafe, async.} =
   return success()
 
diff --git a/tests/datastore/querycommontests.nim b/tests/datastore/querycommontests.nim
index b842ab5..663b75a 100644
--- a/tests/datastore/querycommontests.nim
+++ b/tests/datastore/querycommontests.nim
@@ -91,10 +91,8 @@ 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 )
+      iter = tryGet(await ds.query(q))
+      res = tryGet(await iter.waitForAllQueryResults())
 
     check:
       res.len == 2
@@ -119,9 +117,7 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
       iter = (await ds.query(q)).tryGet
 
     var
-      res = (await allFinished(toSeq(iter)))
-        .mapIt( it.read.tryGet )
-        .filterIt( it.key.isSome )
+      res = tryGet(await ds.query(q).waitForAllQueryResults())
 
     res.sort do (a, b: QueryResponse) -> int:
       cmp(a.key.get.id, b.key.get.id)
@@ -153,16 +149,11 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
         (await ds.put(key, val)).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 == 10
 
-      (await iter.dispose()).tryGet
-
     test "Should not apply offset":
       let
         key = Key.init("/a").tryGet
@@ -176,16 +167,11 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
         (await ds.put(key, val)).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 == 10
 
-      (await iter.dispose()).tryGet
-
     test "Should not apply offset and limit":
       let
         key = Key.init("/a").tryGet
@@ -199,10 +185,8 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
         (await ds.put(key, val)).tryGet
 
       let
-        iter = (await ds.query(q)).tryGet
-        res = (await allFinished(toSeq(iter)))
-          .mapIt( it.read.tryGet )
-          .filterIt( it.key.isSome )
+        iter = tryGet(await ds.query(q))
+        res = tryGet(await iter.waitForAllQueryResults())
 
       check:
         res.len == 5
@@ -238,10 +222,8 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
 
       kvs = kvs.reversed
       let
-        iter = (await ds.query(q)).tryGet
-        res = (await allFinished(toSeq(iter)))
-          .mapIt( it.read.tryGet )
-          .filterIt( it.key.isSome )
+        iter = tryGet(await ds.query(q))
+        res = tryGet(await iter.waitForAllQueryResults())
 
       check:
         res.len == 100