Iterator checks context

- Check context for error on each iteration.
- Test that canceling context stops iteration.

Author: gammazero

datastore.go

@@ -124,6 +124,10 @@ func QueryIter(ctx context.Context, ds Read, q query.Query) iter.Seq2[query.Entr
 		defer results.Close()
 
 		for result := range results.Next() {
+			if ctx.Err() != nil {
+				yield(query.Entry{}, ctx.Err())
+				return
+			}
 			if result.Error != nil {
 				yield(query.Entry{}, result.Error)
 				return

test/basic_tests.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/rand"
+	"errors"
 	"fmt"
 	"reflect"
 	"strings"
@@ -403,7 +404,7 @@ func randValue() []byte {
 }
 
 func subtestQuery(t *testing.T, ds dstore.Datastore, q dsq.Query, count int) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
 
 	var input []dsq.Entry
 	for i := 0; i < count; i++ {
@@ -538,6 +539,31 @@ func subtestQuery(t *testing.T, ds dstore.Datastore, q dsq.Query, count int) {
 		}
 	}
 
+	const cancelAt = 1
+	if len(actual) > cancelAt {
+		// Test that query iterator stops when context is canceled.
+		var i int
+		for ent, err := range dstore.QueryIter(ctx, ds, q) {
+			if err != nil {
+				if !errors.Is(err, context.Canceled) {
+					t.Fatal("query result error: ", err)
+				}
+				t.Log("err at:", i, err)
+				continue
+			}
+			if ent.Key == "" {
+				t.Fatal("entry has empty key")
+			}
+			i++
+			if i == cancelAt {
+				cancel()
+			}
+		}
+		if i != cancelAt {
+			t.Fatal("expected iteration to be canceled at", cancelAt, "canceled at", i)
+		}
+	}
+
 	t.Log("deleting all keys")
 	for _, e := range input {
 		if err := ds.Delete(ctx, dstore.RawKey(e.Key)); err != nil {