Fix SeekGE in DDB (#6561)

* Fix SeekGE in DDB * Fix test * Use callbacks in kv store creation for tests isolation * make linter happy * Make the tests really fail * Fix comments * Update pkg/kv/postgres/store_test.go Co-authored-by: Barak Amar <[email protected]> * More fixes --------- Co-authored-by: Barak Amar <[email protected]>
treeverse · Sep 10, 2023 · 1b4a32d · 1b4a32d
1 parent ae85f0a
commit 1b4a32d
Show file tree

Hide file tree

Showing 14 changed files with 267 additions and 132 deletions.
diff --git a/pkg/gateway/testutil/gateway_setup.go b/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)
 

diff --git a/pkg/kv/cosmosdb/main_test.go b/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)

diff --git a/pkg/kv/cosmosdb/store_test.go b/pkg/kv/cosmosdb/store_test.go
diff --git a/pkg/kv/dynamodb/main_test.go b/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)

diff --git a/pkg/kv/dynamodb/store.go b/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
 	}

diff --git a/pkg/kv/dynamodb/store_test.go b/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:          kvtest.MaxPageSize,
+			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
+	})
 }