Cherry-pick #229 to 1.3.x (backend: 3.3.x)

store/mongo/migration_2_0_0.go → store/mongo/migration_2_0_1.go

@@ -1,4 +1,4 @@
-// Copyright 2021 Northern.tech AS
+// Copyright 2022 Northern.tech AS
 //
 //    Licensed under the Apache License, Version 2.0 (the "License");
 //    you may not use this file except in compliance with the License.
@@ -31,30 +31,19 @@ const (
 	findBatchSize = 255
 )
 
-type migration_2_0_0 struct {
+type migration_2_0_1 struct {
 	client *mongo.Client
 	db     string
 }
 
-func (m *migration_2_0_0) Up(from migrate.Version) error {
+func (m *migration_2_0_1) Up(from migrate.Version) error {
 	ctx := context.Background()
 	client := m.client
 
 	collections := map[string]struct {
 		Indexes []mongo.IndexModel
 	}{
-		DevicesCollectionName: {
-			Indexes: []mongo.IndexModel{
-				{
-					Keys: bson.D{
-						{Key: mstore.FieldTenantID, Value: 1},
-						{Key: dbFieldID, Value: 1},
-					},
-					Options: mopts.Index().
-						SetName(mstore.FieldTenantID + "_" + dbFieldID),
-				},
-			},
-		},
+		DevicesCollectionName: {},
 		SessionsCollectionName: {
 			Indexes: []mongo.IndexModel{
 				{
@@ -104,50 +93,57 @@ func (m *migration_2_0_0) Up(from migrate.Version) error {
 	ctx = identity.WithContext(ctx, &identity.Identity{
 		Tenant: tenantID,
 	})
+	writes := make([]mongo.WriteModel, 0, findBatchSize)
 
-	for collection := range collections {
-		// get all the documents in the collection
+	for collection, idxes := range collections {
+		writes = writes[:0]
 		findOptions := mopts.Find().
 			SetBatchSize(findBatchSize).
 			SetSort(bson.D{{Key: dbFieldID, Value: 1}})
-		coll := client.Database(m.db).Collection(collection)
 		collOut := client.Database(DbName).Collection(collection)
-		_, err := collOut.Indexes().CreateMany(ctx, collections[collection].Indexes)
-		if err != nil {
-			return err
+		if m.db == DbName {
+			if len(idxes.Indexes) > 0 {
+				_, err := collOut.Indexes().CreateMany(ctx, collections[collection].Indexes)
+				if err != nil {
+					return err
+				}
+			}
+			_, err := collOut.UpdateMany(ctx, bson.D{
+				{Key: mstore.FieldTenantID, Value: bson.D{
+					{Key: "$exists", Value: false},
+				}},
+			}, bson.D{{Key: "$set", Value: bson.D{
+				{Key: mstore.FieldTenantID, Value: ""},
+			}}},
+			)
+			if err != nil {
+				return err
+			}
+			continue
 		}
 
+		coll := client.Database(m.db).Collection(collection)
+		// get all the documents in the collection
 		cur, err := coll.Find(ctx, bson.D{}, findOptions)
 		if err != nil {
 			return err
 		}
 		defer cur.Close(ctx)
 
-		writes := make([]mongo.WriteModel, 0, findBatchSize)
-
 		// migrate the documents
 		for cur.Next(ctx) {
-			item := bson.D{}
-			err := cur.Decode(&item)
-			if err != nil {
+			id := cur.Current.Lookup(dbFieldID)
+			var item bson.D
+			if err = cur.Decode(&item); err != nil {
 				return err
 			}
-
-			item = mstore.WithTenantID(ctx, item)
-			if m.db == DbName {
-				filter := bson.D{}
-				for _, i := range item {
-					if i.Key == dbFieldID {
-						filter = append(filter, i)
-					}
-				}
-				writes = append(writes,
-					mongo.NewReplaceOneModel().
-						SetFilter(filter).SetReplacement(item))
-			} else {
-				writes = append(writes, mongo.NewInsertOneModel().SetDocument(item))
-			}
+			writes = append(writes, mongo.
+				NewReplaceOneModel().
+				SetFilter(bson.D{{Key: dbFieldID, Value: id}}).
+				SetUpsert(true).
+				SetReplacement(mstore.WithTenantID(ctx, item)))
 			if len(writes) == findBatchSize {
+				_, err = collOut.BulkWrite(ctx, writes)
 				if err != nil {
 					return err
 				}
@@ -165,6 +161,6 @@ func (m *migration_2_0_0) Up(from migrate.Version) error {
 	return nil
 }
 
-func (m *migration_2_0_0) Version() migrate.Version {
-	return migrate.MakeVersion(2, 0, 0)
+func (m *migration_2_0_1) Version() migrate.Version {
+	return migrate.MakeVersion(2, 0, 1)
 }

store/mongo/migration_2_0_1_test.go

@@ -0,0 +1,125 @@
+// Copyright 2022 Northern.tech AS
+//
+//    Licensed under the Apache License, Version 2.0 (the "License");
+//    you may not use this file except in compliance with the License.
+//    You may obtain a copy of the License at
+//
+//        http://www.apache.org/licenses/LICENSE-2.0
+//
+//    Unless required by applicable law or agreed to in writing, software
+//    distributed under the License is distributed on an "AS IS" BASIS,
+//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//    See the License for the specific language governing permissions and
+//    limitations under the License.
+
+package mongo
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.mongodb.org/mongo-driver/bson"
+
+	mstore "github.com/mendersoftware/go-lib-micro/store/v2"
+)
+
+type index struct {
+	Keys        map[string]int `bson:"key"`
+	Name        string         `bson:"name"`
+	ExpireAfter *int           `bson:"expireAfterSeconds"`
+}
+
+func TestMigration_2_0_1(t *testing.T) {
+	db := db.Client().Database(DbName)
+
+	collDevs := db.Collection(DevicesCollectionName)
+	ctx := context.Background()
+
+	_, err := collDevs.InsertMany(ctx, func() []interface{} {
+		ret := make([]interface{}, findBatchSize+1)
+		for i := range ret {
+			ret[i] = bson.D{}
+		}
+		return ret
+	}())
+	if !assert.NoError(t, err) {
+		t.FailNow()
+	}
+	_, err = db.Client().
+		Database(DbName+"-123456789012345678901234").
+		Collection(DevicesCollectionName).
+		InsertMany(ctx, func() []interface{} {
+			ret := make([]interface{}, findBatchSize+1)
+			for i := range ret {
+				ret[i] = bson.D{}
+			}
+			return ret
+		}())
+	require.NoError(t, err)
+
+	// NOTE: We're using 2.0.1 since it's the same migration, thus verifying
+	//       that the migration is reentrant.
+	err = Migrate(ctx, DbName, "2.0.1", db.Client(), true)
+	require.NoError(t, err)
+
+	collNames, err := db.ListCollectionNames(ctx, bson.D{})
+	for _, collName := range collNames {
+		idxes, err := db.Collection(collName).
+			Indexes().
+			ListSpecifications(ctx)
+		require.NoError(t, err)
+		expectedLen := 1
+		switch collName {
+		case RecordingsCollectionName:
+			expectedLen = 3
+		case SessionsCollectionName, ControlCollectionName:
+			expectedLen = 2
+		}
+		require.Len(t, idxes, expectedLen)
+		for _, idx := range idxes {
+			if idx.Name == "_id_" {
+				// Skip default index
+				continue
+			}
+			switch idx.Name {
+			case IndexNameLogsExpire:
+				var keys bson.D
+				assert.True(t, idx.ExpireAfterSeconds != nil, "Expected index to have TTL")
+				err := bson.Unmarshal(idx.KeysDocument, &keys)
+				require.NoError(t, err)
+				assert.Equal(t, bson.D{
+					{Key: "expire_ts", Value: int32(1)},
+				}, keys)
+
+			case mstore.FieldTenantID + "_" + dbFieldSessionID:
+				var keys bson.D
+				err := bson.Unmarshal(idx.KeysDocument, &keys)
+				require.NoError(t, err)
+				assert.Equal(t, bson.D{
+					{Key: mstore.FieldTenantID, Value: int32(1)},
+					{Key: dbFieldSessionID, Value: int32(1)},
+				}, keys)
+
+			case mstore.FieldTenantID + "_" + dbFieldDeviceID:
+				var keys bson.D
+				err := bson.Unmarshal(idx.KeysDocument, &keys)
+				require.NoError(t, err)
+				assert.Equal(t, bson.D{
+					{Key: mstore.FieldTenantID, Value: int32(1)},
+					{Key: dbFieldDeviceID, Value: int32(1)},
+				}, keys)
+
+			default:
+				assert.Failf(t, "index test failed", "Index name \"%s\" not recognized", idx.Name)
+			}
+		}
+	}
+
+	actual, err := collDevs.CountDocuments(ctx, bson.D{
+		{Key: mstore.FieldTenantID, Value: bson.D{{Key: "$exists", Value: false}}},
+	})
+	require.NoError(t, err)
+	assert.Zerof(t, actual, "%d documents are not indexed by tenant_id", actual)
+}