Fix SeekGE in DDB

pkg/kv/cosmosdb/main_test.go

@@ -1,14 +1,32 @@
 package cosmosdb_test
 
 import (
+	"context"
 	"os"
 	"testing"
 
+	"github.com/treeverse/lakefs/pkg/kv"
+	"github.com/treeverse/lakefs/pkg/kv/cosmosdb"
 	"github.com/treeverse/lakefs/pkg/kv/kvparams"
+	"github.com/treeverse/lakefs/pkg/kv/kvtest"
 )
 
 var testParams *kvparams.CosmosDB
 
+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, 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
+	})
+}
+
 func TestMain(m *testing.M) {
 	code := m.Run()
 	os.Exit(code)

pkg/kv/dynamodb/store_test.go

@@ -16,7 +16,7 @@ func TestDynamoKV(t *testing.T) {
 		t.Helper()
 		testParams = &kvparams.DynamoDB{
 			TableName:          testutil.UniqueKVTableName(),
-			ScanLimit:          10,
+			ScanLimit:          kvtest.MaxPageSize,
 			Endpoint:           databaseURI,
 			AwsRegion:          "us-east-1",
 			AwsAccessKeyID:     "fakeMyKeyId",

pkg/kv/kvtest/iterators.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 
 	"github.com/go-test/deep"
-	"github.com/stretchr/testify/require"
 	"github.com/treeverse/lakefs/pkg/graveler"
 	"github.com/treeverse/lakefs/pkg/kv"
 )
@@ -22,6 +21,9 @@ type StoreWithCounter struct {
 	ScanCalls int64
 }
 
+// MaxPageSize is the maximum page size for pagination tests
+const MaxPageSize = 10
+
 func NewStoreWithCounter(store kv.Store) *StoreWithCounter {
 	return &StoreWithCounter{Store: store}
 }
@@ -63,15 +65,6 @@ func testPartitionIteratorSeekGEWithPagination(ctx context.Context, store kv.Sto
 			"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 {
@@ -90,11 +83,22 @@ func testPartitionIteratorSeekGEWithPagination(ctx context.Context, store kv.Sto
 		}
 		defer itr.Close()
 
+		if !itr.Next() {
+			t.Fatal("expected Next to be true")
+		}
+
 		itr.SeekGE([]byte("b"))
-		require.True(t, itr.Next())
+		for i := 0; i < MaxPageSize+1; i++ {
+			// force pagination
+			if !itr.Next() {
+				t.Fatal("expected Next to be true")
+			}
+		}
 
 		itr.SeekGE([]byte("z"))
-		require.True(t, itr.Next())
+		if !itr.Next() {
+			t.Fatal("expected Next to be true")
+		}
 
 		itr.SeekGE([]byte("d1"))
 		names := scanPartitionIterator(t, itr, func(_ []byte, model *TestModel) string { return string(model.Name) })

pkg/kv/local/store_test.go

@@ -2,9 +2,9 @@ package local_test
 
 import (
 	"context"
-	"github.com/treeverse/lakefs/pkg/kv"
 	"testing"
 
+	"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/local"

pkg/kv/mem/store_test.go

@@ -2,9 +2,9 @@ package mem_test
 
 import (
 	"context"
-	"github.com/treeverse/lakefs/pkg/kv"
 	"testing"
 
+	"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"

pkg/kv/postgres/store_test.go

@@ -2,29 +2,45 @@ package postgres_test
 
 import (
 	"context"
-	"github.com/treeverse/lakefs/pkg/kv"
-	"github.com/treeverse/lakefs/pkg/testutil"
+	"fmt"
 	"testing"
 
+	"github.com/jackc/pgx/v5"
+	"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/postgres"
+	"github.com/treeverse/lakefs/pkg/testutil"
 )
 
 func TestPostgresKV(t *testing.T) {
+	databaseURI, cleanup := runDBInstance(pool, testutil.UniqueKVTableName())
+	t.Cleanup(cleanup)
+
 	kvtest.DriverTest(t, func(t testing.TB, ctx context.Context) kv.Store {
 		t.Helper()
-		databaseURI, cleanup := runDBInstance(pool, testutil.UniqueKVTableName())
+
+		conn, err := pgx.Connect(ctx, databaseURI)
+		if err != nil {
+			t.Fatalf("Unable to connect to database: %v", err)
+		}
+		defer conn.Close(context.Background())
+
+		// create a new schema per test
+		schemaName := "test_schema" + testutil.UniqueName()
+		_, err = conn.Exec(context.Background(), fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s;", schemaName))
+		if err != nil {
+			t.Fatalf("Error creating schema: %v", err)
+		}
 
 		store, err := kv.Open(ctx, kvparams.Config{
 			Type:     postgres.DriverName,
-			Postgres: &kvparams.Postgres{ConnectionString: databaseURI, ScanPageSize: 10},
+			Postgres: &kvparams.Postgres{ConnectionString: fmt.Sprintf("%s&search_path=%s", databaseURI, schemaName), ScanPageSize: kvtest.MaxPageSize},
 		})
 		if err != nil {
 			t.Fatalf("failed to open kv '%s' store: %s", postgres.DriverName, err)
 		}
 		t.Cleanup(store.Close)
-		t.Cleanup(cleanup)
 		return store
 	})
 }

pkg/testutil/dynamodb.go

@@ -78,7 +78,11 @@ func GetDynamoDBInstance() (string, func(), error) {
 }
 
 func UniqueKVTableName() string {
-	return "kvstore_" + nanoid.MustGenerate(chars, charsSize)
+	return "kvstore_" + UniqueName()
+}
+
+func UniqueName() string {
+	return nanoid.MustGenerate(chars, charsSize)
 }
 
 func GetDynamoDBProd(ctx context.Context, tb testing.TB) kv.Store {