Add upgrade steps to count and record the models for each cloud

MalloZup · May 17, 2018 · 7cc7072 · 7cc7072
1 parent 1c16b12
commit 7cc7072
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 0 deletions.
diff --git a/state/upgrades.go b/state/upgrades.go
@@ -1476,6 +1476,38 @@ func RemoveVotingMachineIds(st *State) error {
 	return nil
 }
 
+// AddCloudModelCounts updates cloud docs to ensure the model count field is set.
+func AddCloudModelCounts(st *State) error {
+	cloudsColl, closer := st.db().GetCollection(cloudsC)
+	defer closer()
+
+	var clouds []cloudDoc
+	err := cloudsColl.Find(nil).All(&clouds)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	modelsColl, closer := st.db().GetCollection(modelsC)
+	defer closer()
+
+	var updateOps []txn.Op
+	for _, c := range clouds {
+		n, err := modelsColl.Find(bson.D{{"cloud", c.Name}}).Count()
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if n != c.ModelCount {
+			updateOps = append(updateOps, txn.Op{
+				C:      cloudsC,
+				Id:     c.Name,
+				Assert: txn.DocExists,
+				Update: bson.D{{"$set", bson.D{{"modelcount", n}}}},
+			})
+		}
+	}
+	return st.db().RunTransaction(updateOps)
+}
+
 // UpgradeDefaultContainerImageStreamConfig ensures that the config value for
 // container-image-stream is set to its default value, "released".
 func UpgradeContainerImageStreamDefault(st *State) error {

diff --git a/state/upgrades_test.go b/state/upgrades_test.go
@@ -2509,3 +2509,60 @@ func (s *upgradesSuite) TestRemoveContainerImageStreamFromNonModelSettings(c *gc
 		expectUpgradedData{settingsColl, expectedSettings},
 	)
 }
+
+func (s *upgradesSuite) TestAddCloudModelCounts(c *gc.C) {
+	modelsColl, closer := s.state.db().GetRawCollection(modelsC)
+	defer closer()
+
+	err := modelsColl.Insert(
+		modelDoc{
+			Type:           ModelTypeIAAS,
+			UUID:           "0000-dead-beaf-0001",
+			Owner:          "user-admin@local",
+			Name:           "controller",
+			ControllerUUID: "deadbeef-1bad-500d-9000-4b1d0d06f00d",
+			Cloud:          "cloud-foo",
+		},
+		modelDoc{
+			Type:           ModelTypeIAAS,
+			UUID:           "0000-dead-beaf-0002",
+			Owner:          "user-mary@external",
+			Name:           "default",
+			ControllerUUID: "deadbeef-1bad-500d-9000-4b1d0d06f00d",
+			Cloud:          "cloud-foo",
+		},
+	)
+	c.Assert(err, jc.ErrorIsNil)
+
+	cloudsColl, closer := s.state.db().GetRawCollection(cloudsC)
+	defer closer()
+
+	err = cloudsColl.Insert(
+		bson.M{
+			"_id":        "cloud-foo",
+			"name":       "cloud-foo",
+			"type":       "dummy",
+			"auth-types": []string{"empty"},
+			"endpoint":   "here",
+		},
+	)
+	c.Assert(err, jc.ErrorIsNil)
+
+	expected := []bson.M{{
+		"_id":        "cloud-foo",
+		"name":       "cloud-foo",
+		"type":       "dummy",
+		"auth-types": []interface{}{"empty"},
+		"endpoint":   "here",
+		"modelcount": 2,
+	}, {
+		"_id":        "dummy",
+		"name":       "dummy",
+		"type":       "dummy",
+		"auth-types": []interface{}{"empty"},
+		"regions":    bson.M{"dummy-region": bson.M{}},
+		"endpoint":   "",
+		"modelcount": 1, // unchanged
+	}}
+	s.assertUpgradedData(c, AddCloudModelCounts, expectUpgradedData{cloudsColl, expected})
+}
diff --git a/upgrades/backend.go b/upgrades/backend.go
@@ -41,6 +41,7 @@ type StateBackend interface {
 	MoveMongoSpaceToHASpaceConfig() error
 	CreateMissingApplicationConfig() error
 	RemoveVotingMachineIds() error
+	AddCloudModelCounts() error
 }
 
 // Model is an interface providing access to the details of a model within the
@@ -155,6 +156,10 @@ func (s stateBackend) RemoveVotingMachineIds() error {
 	return state.RemoveVotingMachineIds(s.st)
 }
 
+func (s stateBackend) AddCloudModelCounts() error {
+	return state.AddCloudModelCounts(s.st)
+}
+
 type modelShim struct {
 	st *state.State
 	m  *state.Model

diff --git a/upgrades/steps_24.go b/upgrades/steps_24.go
@@ -45,6 +45,13 @@ func stateStepsFor24() []Step {
 				return context.State().RemoveVotingMachineIds()
 			},
 		},
+		&upgradeStep{
+			description: "add cloud model counts",
+			targets:     []Target{DatabaseMaster},
+			run: func(context Context) error {
+				return context.State().AddCloudModelCounts()
+			},
+		},
 	}
 }
 

diff --git a/upgrades/steps_24_test.go b/upgrades/steps_24_test.go
@@ -43,3 +43,9 @@ func (s *steps24Suite) TestRemoveVotingMachineIds(c *gc.C) {
 	// Logic for step itself is tested in state package.
 	c.Assert(step.Targets(), jc.DeepEquals, []upgrades.Target{upgrades.DatabaseMaster})
 }
+
+func (s *steps24Suite) TestCloudModelCounts(c *gc.C) {
+	step := findStateStep(c, v24, "add cloud model counts")
+	// Logic for step itself is tested in state package.
+	c.Assert(step.Targets(), jc.DeepEquals, []upgrades.Target{upgrades.DatabaseMaster})
+}