Fix SeekGE in DDB

pkg/gateway/testutil/gateway_setup.go

@@ -15,8 +15,8 @@ import (
 	"github.com/treeverse/lakefs/pkg/config"
 	"github.com/treeverse/lakefs/pkg/gateway"
 	"github.com/treeverse/lakefs/pkg/gateway/multipart"
+	"github.com/treeverse/lakefs/pkg/kv"
 	"github.com/treeverse/lakefs/pkg/kv/kvparams"
-	"github.com/treeverse/lakefs/pkg/kv/kvtest"
 	_ "github.com/treeverse/lakefs/pkg/kv/mem"
 	"github.com/treeverse/lakefs/pkg/logging"
 	"github.com/treeverse/lakefs/pkg/stats"
@@ -34,7 +34,8 @@ func GetBasicHandler(t *testing.T, authService *FakeAuthService, repoName string
 	ctx := context.Background()
 	viper.Set(config.BlockstoreTypeKey, block.BlockstoreTypeMem)
 
-	store := kvtest.MakeStoreByName("mem", kvparams.Config{})(t, ctx)
+	store, err := kv.Open(ctx, kvparams.Config{Type: "mem"})
+	testutil.MustDo(t, "open kv store", err)
 	defer store.Close()
 	multipartTracker := multipart.NewTracker(store)
 

pkg/kv/cosmosdb/store_test.go

@@ -1,15 +1,25 @@
 package cosmosdb_test
 
 import (
-	"testing"
-
+	"context"
+	"github.com/treeverse/lakefs/pkg/kv"
 	"github.com/treeverse/lakefs/pkg/kv/cosmosdb"
 	"github.com/treeverse/lakefs/pkg/kv/kvparams"
+	"testing"
+
 	"github.com/treeverse/lakefs/pkg/kv/kvtest"
 )
 
 func TestCosmosDB(t *testing.T) {
 	t.Skip("CosmosDB tests are flaky due to the emulator. If you plan on running those, make sure to assign at least 3CPUs and" +
 		" 4GB of memory to the container running the emulator.")
-	kvtest.DriverTest(t, cosmosdb.DriverName, kvparams.Config{CosmosDB: testParams})
+	kvtest.DriverTest(t, func(t testing.TB, ctx context.Context) kv.Store {
+		t.Helper()
+		store, err := kv.Open(ctx, kvparams.Config{CosmosDB: testParams, Type: cosmosdb.DriverName})
+		if err != nil {
+			t.Fatalf("failed to open kv '%s' store: %s", cosmosdb.DriverName, err)
+		}
+		t.Cleanup(store.Close)
+		return store
+	})
 }

pkg/kv/dynamodb/main_test.go

@@ -10,22 +10,15 @@ import (
 )
 
 var testParams *kvparams.DynamoDB
+var databaseURI string
 
 func TestMain(m *testing.M) {
-	databaseURI, cleanupFunc, err := testutil.GetDynamoDBInstance()
+	var err error
+	var cleanupFunc func()
+	databaseURI, cleanupFunc, err = testutil.GetDynamoDBInstance()
 	if err != nil {
 		log.Fatalf("Could not connect to Docker: %s", err)
 	}
-
-	testParams = &kvparams.DynamoDB{
-		TableName:          testutil.UniqueKVTableName(),
-		ScanLimit:          10,
-		Endpoint:           databaseURI,
-		AwsRegion:          "us-east-1",
-		AwsAccessKeyID:     "fakeMyKeyId",
-		AwsSecretAccessKey: "fakeSecretAccessKey",
-	}
-
 	code := m.Run()
 	cleanupFunc()
 	os.Exit(code)

pkg/kv/dynamodb/store.go

@@ -380,6 +380,7 @@ func (s *Store) DropTable() error {
 func (e *EntriesIterator) SeekGE(key []byte) {
 	if !e.isInRange(key) {
 		e.startKey = key
+		e.exclusiveStartKey = nil
 		e.runQuery()
 		return
 	}

pkg/kv/dynamodb/store_test.go

@@ -1,13 +1,33 @@
 package dynamodb_test
 
 import (
+	"context"
+	"github.com/treeverse/lakefs/pkg/kv"
+	"github.com/treeverse/lakefs/pkg/kv/dynamodb"
+	"github.com/treeverse/lakefs/pkg/testutil"
 	"testing"
 
-	"github.com/treeverse/lakefs/pkg/kv/dynamodb"
 	"github.com/treeverse/lakefs/pkg/kv/kvparams"
 	"github.com/treeverse/lakefs/pkg/kv/kvtest"
 )
 
 func TestDynamoKV(t *testing.T) {
-	kvtest.DriverTest(t, dynamodb.DriverName, kvparams.Config{DynamoDB: testParams})
+	kvtest.DriverTest(t, func(t testing.TB, ctx context.Context) kv.Store {
+		t.Helper()
+		testParams = &kvparams.DynamoDB{
+			TableName:          testutil.UniqueKVTableName(),
+			ScanLimit:          10,
+			Endpoint:           databaseURI,
+			AwsRegion:          "us-east-1",
+			AwsAccessKeyID:     "fakeMyKeyId",
+			AwsSecretAccessKey: "fakeSecretAccessKey",
+		}
+
+		store, err := kv.Open(ctx, kvparams.Config{DynamoDB: testParams, Type: dynamodb.DriverName})
+		if err != nil {
+			t.Fatalf("failed to open kv '%s' store: %s", dynamodb.DriverName, err)
+		}
+		t.Cleanup(store.Close)
+		return store
+	})
 }

pkg/kv/kvtest/iterators.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/go-test/deep"
+	"github.com/stretchr/testify/require"
 	"github.com/treeverse/lakefs/pkg/graveler"
 	"github.com/treeverse/lakefs/pkg/kv"
 )
@@ -46,60 +47,120 @@ func testPartitionIterator(t *testing.T, ms MakeStore) {
 		}
 	}
 
-	t.Run("listing all values of partition", func(t *testing.T) {
-		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), firstPartitionKey, 0)
+	t.Run("listing all values of partition", testPartitionIteratorListAll(ctx, store))
+	t.Run("listing values SeekGE", listPartitionIteratorWithSeekGE(ctx, store))
+	t.Run("count scans on successive SeekGE operations", testPartitionIteratorCountScansOnSeekGE(store, ctx))
+	t.Run("failed SeekGE partition not found", testPartitionIteratorSeekGEOnPartitionNotFound(ctx, store))
+	t.Run("SeekGE past end", testPartitionIteratorSeekGEPastEnd(ctx, store))
+	t.Run("SeekGE seek back", testPartitionIteratorSeekGESeekBack(ctx, store))
+	t.Run("listing values SeekGE with pagination", testPartitionIteratorSeekGEWithPagination(ctx, store))
+}
+
+func testPartitionIteratorSeekGEWithPagination(ctx context.Context, store kv.Store) func(t *testing.T) {
+	return func(t *testing.T) {
+		// load much more data
+		moreModelNames := []string{
+			"da", "db", "dc", "dd", "de", "df", "dg", "dh", "di", "dj",
+			"dk", "dl", "dm", "dn", "do", "dp", "dq", "dr", "ds", "dt",
+			"du", "dv", "dw", "dx", "dy", "dz",
+			"ea", "eb", "ec", "ed", "ee", "ef", "eg", "eh", "ei", "ej",
+			"ek", "el", "em", "en", "eo", "ep", "eq", "er", "es", "et",
+			"eu", "ev", "ew", "ex", "ey", "ez",
+			"fa", "fb", "fc", "fd", "fe", "ff", "fg", "fh", "fi", "fj",
+			"fk", "fl", "fm", "fn", "fo", "fp", "fq", "fr", "fs", "ft",
+			"fu", "fv", "fw", "fx", "fy", "fz",
+			"ga", "gb", "gc", "gd", "ge", "gf", "gg", "gh", "gi", "gj",
+			"gk", "gl", "gm", "gn", "go", "gp", "gq", "gr", "gs", "gt",
+			"gu", "gv", "gw", "gx", "gy", "gz",
+			"z",
+		}
+		for _, name := range moreModelNames {
+			model := TestModel{Name: []byte(name)}
+			for _, partitionKey := range []string{firstPartitionKey, secondPartitionKey} {
+				err := kv.SetMsg(ctx, store, partitionKey, model.Name, &model)
+				if err != nil {
+					t.Fatalf("failed to set model (partition %s, name %s): %s", partitionKey, name, err)
+				}
+			}
+		}
+
+		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), secondPartitionKey, 0)
 		if itr == nil {
 			t.Fatalf("failed to create partition iterator")
 		}
 		defer itr.Close()
+
+		itr.SeekGE([]byte("b"))
+		require.True(t, itr.Next())
+
+		itr.SeekGE([]byte("z"))
+		require.True(t, itr.Next())
+
+		itr.SeekGE([]byte("d1"))
 		names := scanPartitionIterator(t, itr, func(_ []byte, model *TestModel) string { return string(model.Name) })
-		if diffs := deep.Equal(names, []string{"a", "aa", "b", "c", "d"}); diffs != nil {
+		if diffs := deep.Equal(names, moreModelNames); diffs != nil {
 			t.Fatalf("got wrong list of names: %v", diffs)
 		}
-	})
+	}
+}
 
-	t.Run("listing values SeekGE", func(t *testing.T) {
-		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), secondPartitionKey, 0)
+func testPartitionIteratorSeekGESeekBack(ctx context.Context, store kv.Store) func(t *testing.T) {
+	return func(t *testing.T) {
+		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), firstPartitionKey, 0)
 		if itr == nil {
 			t.Fatalf("failed to create partition iterator")
 		}
 		defer itr.Close()
-		for _, seekValue := range []string{"b", "aaa", "b"} {
-			itr.SeekGE([]byte(seekValue))
-			names := scanPartitionIterator(t, itr, func(_ []byte, model *TestModel) string { return string(model.Name) })
-			if diffs := deep.Equal(names, []string{"b", "c", "d"}); diffs != nil {
-				t.Fatalf("got wrong list of names: %v", diffs)
-			}
+		itr.SeekGE([]byte("z"))
+		if itr.Next() {
+			t.Fatal("expected Next to be false")
 		}
-	})
+		if err := itr.Err(); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		itr.SeekGE([]byte("a"))
+		if !itr.Next() {
+			t.Fatalf("expected Next to be true")
+		}
+		if err := itr.Err(); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		e := itr.Entry()
+		model := e.Value.(*TestModel)
+		if string(model.Name) != "a" {
+			t.Fatalf("expected value a from iterator")
+		}
+	}
+}
 
-	t.Run("count scans on successive SeekGE operations", func(t *testing.T) {
-		store := NewStoreWithCounter(store)
-		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), secondPartitionKey, 0)
+func testPartitionIteratorSeekGEPastEnd(ctx context.Context, store kv.Store) func(t *testing.T) {
+	return func(t *testing.T) {
+		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), firstPartitionKey, 0)
 		if itr == nil {
 			t.Fatalf("failed to create partition iterator")
 		}
 		defer itr.Close()
-		for _, seekValue := range []string{"b", "c", "d"} {
-			itr.SeekGE([]byte(seekValue))
-			if !itr.Next() {
-				t.Fatalf("Expected iterator to have a value")
-			}
-			if err := itr.Err(); err != nil {
-				t.Fatalf("unexpected error: %v", err)
-			}
-			k := itr.Entry().Key
-			if string(k) != seekValue {
-				t.Fatalf("Expected to find value %s. Found %s", seekValue, k)
-			}
+		itr.SeekGE([]byte("b"))
+		if !itr.Next() {
+			t.Fatal("expected Next to be true")
 		}
-		scanCalls := atomic.LoadInt64(&store.ScanCalls)
-		if scanCalls != 1 {
-			t.Fatalf("Expected exactly 1 call to Scan. got: %d", scanCalls)
+		e := itr.Entry()
+		model := e.Value.(*TestModel)
+		if string(model.Name) != "b" {
+			t.Fatalf("expected value b from iterator")
 		}
-	})
+		itr.SeekGE(graveler.UpperBoundForPrefix([]byte("d1")))
+		if itr.Next() {
+			t.Fatalf("expected Next to be false")
+		}
+		if err := itr.Err(); err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+	}
+}
 
-	t.Run("failed SeekGE partition not found", func(t *testing.T) {
+func testPartitionIteratorSeekGEOnPartitionNotFound(ctx context.Context, store kv.Store) func(t *testing.T) {
+	return func(t *testing.T) {
 		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), "", 0)
 		if itr == nil {
 			t.Fatalf("failed to create partition iterator")
@@ -118,58 +179,66 @@ func testPartitionIterator(t *testing.T, ms MakeStore) {
 		if err := itr.Err(); !errors.Is(err, kv.ErrMissingPartitionKey) {
 			t.Fatalf("expected error: %s, got %v", kv.ErrMissingPartitionKey, err)
 		}
-	})
+	}
+}
 
-	t.Run("SeekGE past end", func(t *testing.T) {
-		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), firstPartitionKey, 0)
+func testPartitionIteratorCountScansOnSeekGE(store kv.Store, ctx context.Context) func(t *testing.T) {
+	return func(t *testing.T) {
+		store := NewStoreWithCounter(store)
+		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), secondPartitionKey, 0)
 		if itr == nil {
 			t.Fatalf("failed to create partition iterator")
 		}
 		defer itr.Close()
-		itr.SeekGE([]byte("b"))
-		if !itr.Next() {
-			t.Fatal("expected Next to be true")
+		for _, seekValue := range []string{"b", "c", "d"} {
+			itr.SeekGE([]byte(seekValue))
+			if !itr.Next() {
+				t.Fatalf("Expected iterator to have a value")
+			}
+			if err := itr.Err(); err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			k := itr.Entry().Key
+			if string(k) != seekValue {
+				t.Fatalf("Expected to find value %s. Found %s", seekValue, k)
+			}
 		}
-		e := itr.Entry()
-		model := e.Value.(*TestModel)
-		if string(model.Name) != "b" {
-			t.Fatalf("expected value b from iterator")
+		scanCalls := atomic.LoadInt64(&store.ScanCalls)
+		if scanCalls != 1 {
+			t.Fatalf("Expected exactly 1 call to Scan. got: %d", scanCalls)
 		}
-		itr.SeekGE(graveler.UpperBoundForPrefix([]byte("d1")))
-		if itr.Next() {
-			t.Fatalf("expected Next to be false")
+	}
+}
+
+func listPartitionIteratorWithSeekGE(ctx context.Context, store kv.Store) func(t *testing.T) {
+	return func(t *testing.T) {
+		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), secondPartitionKey, 0)
+		if itr == nil {
+			t.Fatalf("failed to create partition iterator")
 		}
-		if err := itr.Err(); err != nil {
-			t.Fatalf("unexpected error: %v", err)
+		defer itr.Close()
+		for _, seekValue := range []string{"b", "aaa", "b"} {
+			itr.SeekGE([]byte(seekValue))
+			names := scanPartitionIterator(t, itr, func(_ []byte, model *TestModel) string { return string(model.Name) })
+			if diffs := deep.Equal(names, []string{"b", "c", "d"}); diffs != nil {
+				t.Fatalf("got wrong list of names: %v", diffs)
+			}
 		}
-	})
+	}
+}
 
-	t.Run("SeekGE seek back", func(t *testing.T) {
+func testPartitionIteratorListAll(ctx context.Context, store kv.Store) func(t *testing.T) {
+	return func(t *testing.T) {
 		itr := kv.NewPartitionIterator(ctx, store, (&TestModel{}).ProtoReflect().Type(), firstPartitionKey, 0)
 		if itr == nil {
 			t.Fatalf("failed to create partition iterator")
 		}
 		defer itr.Close()
-		itr.SeekGE([]byte("z"))
-		if itr.Next() {
-			t.Fatal("expected Next to be false")
-		}
-		if err := itr.Err(); err != nil {
-			t.Fatalf("unexpected error: %s", err)
-		}
-		itr.SeekGE([]byte("a"))
-		if !itr.Next() {
-			t.Fatalf("expected Next to be true")
-		}
-		if err := itr.Err(); err != nil {
-			t.Fatalf("unexpected error: %s", err)
-		}
-		e := itr.Entry()
-		model := e.Value.(*TestModel)
-		if string(model.Name) != "a" {
-			t.Fatalf("expected value a from iterator")
+		names := scanPartitionIterator(t, itr, func(_ []byte, model *TestModel) string { return string(model.Name) })
+		if diffs := deep.Equal(names, []string{"a", "aa", "b", "c", "d"}); diffs != nil {
+			t.Fatalf("got wrong list of names: %v", diffs)
 		}
-	})
+	}
 }
 
 // scanPartitionIterator scans the iterator and returns a slice of the results of applying fn to each model.