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

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

diff --git a/datastore/query.nim b/datastore/query.nim
index e405556..1d98e3b 100644
--- a/datastore/query.nim
+++ b/datastore/query.nim
@@ -29,9 +29,11 @@ type
     next*: GetNext
     dispose*: IterDispose
 
-proc waitForAllQueryResults*(iter: QueryIter): Future[?!seq[QueryResponse]] {.async.} =
+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()
@@ -45,12 +47,6 @@ proc waitForAllQueryResults*(iter: QueryIter): Future[?!seq[QueryResponse]] {.as
   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 663b75a..6153411 100644
--- a/tests/datastore/querycommontests.nim
+++ b/tests/datastore/querycommontests.nim
@@ -9,7 +9,6 @@ import pkg/stew/byteutils
 
 import pkg/datastore
 
-import pretty
 
 template queryTests*(ds: Datastore, extended = true) {.dirty.} =
   var
@@ -91,8 +90,7 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
     (await ds.put(key3, val3)).tryGet
 
     let
-      iter = tryGet(await ds.query(q))
-      res = tryGet(await iter.waitForAllQueryResults())
+      res = tryGet(await ds.query(q).waitForAllQueryResults())
 
     check:
       res.len == 2
@@ -102,8 +100,6 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
       res[1].key.get == key3
       res[1].data == val3
 
-    (await iter.dispose()).tryGet
-
   test "Key should all list all keys at the same level":
     let
       queryKey = Key.init("/a").tryGet
@@ -185,8 +181,7 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
         (await ds.put(key, val)).tryGet
 
       let
-        iter = tryGet(await ds.query(q))
-        res = tryGet(await iter.waitForAllQueryResults())
+        res = tryGet(await ds.query(q).waitForAllQueryResults())
 
       check:
         res.len == 5
@@ -200,8 +195,6 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
           res[i].key.get == key
           res[i].data == val
 
-      (await iter.dispose()).tryGet
-
     test "Should apply sort order - descending":
       let
         key = Key.init("/a").tryGet
@@ -222,8 +215,7 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
 
       kvs = kvs.reversed
       let
-        iter = tryGet(await ds.query(q))
-        res = tryGet(await iter.waitForAllQueryResults())
+        res = tryGet(await ds.query(q).waitForAllQueryResults())
 
       check:
         res.len == 100
@@ -232,5 +224,3 @@ template queryTests*(ds: Datastore, extended = true) {.dirty.} =
         check:
           res[i].key.get == kvs[i].key.get
           res[i].data == kvs[i].data
-
-      (await iter.dispose()).tryGet