Track branch with -n Option

api/facadeversions.go

@@ -75,7 +75,7 @@ var facadeVersions = map[string]int{
 	"MigrationStatusWatcher":       1,
 	"MigrationTarget":              1,
 	"ModelConfig":                  2,
-	"ModelGeneration":              2,
+	"ModelGeneration":              3,
 	"ModelManager":                 8,
 	"ModelUpgrader":                1,
 	"NotifyWatcher":                1,

api/modelgeneration/modelgeneration.go

@@ -71,10 +71,11 @@ func (c *Client) CommitBranch(branchName string) (int, error) {
 
 // TrackBranch sets the input units and/or applications
 // to track changes made under the input branch name.
-func (c *Client) TrackBranch(branchName string, entities []string) error {
+func (c *Client) TrackBranch(branchName string, entities []string, numUnits int) error {
 	var result params.ErrorResults
 	arg := params.BranchTrackArg{
 		BranchName: branchName,
+		NumUnits:   numUnits,
 	}
 	if len(entities) == 0 {
 		return errors.New("no units or applications specified")

api/modelgeneration/modelgeneration_test.go

@@ -86,15 +86,15 @@ func (s *modelGenerationSuite) TestTrackBranchSuccess(c *gc.C) {
 	s.fCaller.EXPECT().FacadeCall("TrackBranch", arg, gomock.Any()).SetArg(2, resultsSource).Return(nil)
 
 	api := modelgeneration.NewStateFromCaller(s.fCaller)
-	err := api.TrackBranch(s.branchName, []string{"mysql/0", "mysql"})
+	err := api.TrackBranch(s.branchName, []string{"mysql/0", "mysql"}, 0)
 	c.Assert(err, gc.IsNil)
 }
 
 func (s *modelGenerationSuite) TestTrackBranchError(c *gc.C) {
 	defer s.setUpMocks(c).Finish()
 
 	api := modelgeneration.NewStateFromCaller(s.fCaller)
-	err := api.TrackBranch(s.branchName, []string{"mysql/0", "mysql", "machine-3"})
+	err := api.TrackBranch(s.branchName, []string{"mysql/0", "mysql", "machine-3"}, 0)
 	c.Assert(err, gc.ErrorMatches, `"machine-3" is not an application or a unit`)
 }
 

apiserver/allfacades.go

@@ -251,6 +251,7 @@ func AllFacades() *facade.Registry {
 	reg("ModelConfig", 2, modelconfig.NewFacadeV2)
 	reg("ModelGeneration", 1, modelgeneration.NewModelGenerationFacade)
 	reg("ModelGeneration", 2, modelgeneration.NewModelGenerationFacadeV2)
+	reg("ModelGeneration", 3, modelgeneration.NewModelGenerationFacadeV3)
 	reg("ModelManager", 2, modelmanager.NewFacadeV2)
 	reg("ModelManager", 3, modelmanager.NewFacadeV3)
 	reg("ModelManager", 4, modelmanager.NewFacadeV4)

apiserver/facades/client/modelgeneration/interface.go

@@ -39,6 +39,7 @@ type Generation interface {
 	Created() int64
 	CreatedBy() string
 	AssignAllUnits(string) error
+	AssignUnits(string, int) error
 	AssignUnit(string) error
 	AssignedUnits() map[string][]string
 	Commit(string) (int, error)

apiserver/facades/client/modelgeneration/modelgeneration.go

@@ -31,12 +31,16 @@ type API struct {
 	modelCache        ModelCache
 }
 
-type APIV1 struct {
+type APIV2 struct {
 	*API
 }
 
-// NewModelGenerationFacadeV2 provides the signature required for facade registration.
-func NewModelGenerationFacadeV2(ctx facade.Context) (*API, error) {
+type APIV1 struct {
+	*APIV2
+}
+
+// NewModelGenerationFacadeV3 provides the signature required for facade registration.
+func NewModelGenerationFacadeV3(ctx facade.Context) (*API, error) {
 	authorizer := ctx.Auth()
 	st := &stateShim{State: ctx.State()}
 	m, err := st.Model()
@@ -50,6 +54,15 @@ func NewModelGenerationFacadeV2(ctx facade.Context) (*API, error) {
 	return NewModelGenerationAPI(st, authorizer, m, &modelCacheShim{Model: mc})
 }
 
+// NewModelGenerationFacadeV2 provides the signature required for facade registration.
+func NewModelGenerationFacadeV2(ctx facade.Context) (*APIV2, error) {
+	v3, err := NewModelGenerationFacadeV3(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &APIV2{v3}, nil
+}
+
 // NewModelGenerationFacade provides the signature required for facade registration.
 func NewModelGenerationFacade(ctx facade.Context) (*APIV1, error) {
 	v2, err := NewModelGenerationFacadeV2(ctx)
@@ -116,6 +129,14 @@ func (api *API) AddBranch(arg params.BranchArg) (params.ErrorResult, error) {
 	return result, nil
 }
 
+// TrackBranch marks the input units and/or applications as tracking the input
+// branch, causing them to realise changes made under that branch.
+func (api *APIV2) TrackBranch(arg params.BranchTrackArg) (params.ErrorResults, error) {
+	// For backwards compatibility, ensure we always set the NumUnits to 0
+	arg.NumUnits = 0
+	return api.API.TrackBranch(arg)
+}
+
 // TrackBranch marks the input units and/or applications as tracking the input
 // branch, causing them to realise changes made under that branch.
 func (api *API) TrackBranch(arg params.BranchTrackArg) (params.ErrorResults, error) {
@@ -126,6 +147,13 @@ func (api *API) TrackBranch(arg params.BranchTrackArg) (params.ErrorResults, err
 	if !isModelAdmin && !api.isControllerAdmin {
 		return params.ErrorResults{}, common.ErrPerm
 	}
+	// Ensure we guard against the numUnits being greater than 0 and the number
+	// units/applications greater than 1. This is because we don't know how to
+	// topographically distribute between all the applications and units,
+	// especially if an error occurs whilst assigning the units.
+	if arg.NumUnits > 0 && len(arg.Entities) > 1 {
+		return params.ErrorResults{}, errors.Errorf("number of units and unit IDs can not be specified at the same time")
+	}
 
 	branch, err := api.model.Branch(arg.BranchName)
 	if err != nil {
@@ -143,7 +171,7 @@ func (api *API) TrackBranch(arg params.BranchTrackArg) (params.ErrorResults, err
 		}
 		switch tag.Kind() {
 		case names.ApplicationTagKind:
-			result.Results[i].Error = common.ServerError(branch.AssignAllUnits(tag.Id()))
+			result.Results[i].Error = common.ServerError(branch.AssignUnits(tag.Id(), arg.NumUnits))
 		case names.UnitTagKind:
 			result.Results[i].Error = common.ServerError(branch.AssignUnit(tag.Id()))
 		default:

apiserver/facades/client/modelgeneration/modelgeneration_test.go

@@ -68,7 +68,7 @@ func (s *modelGenerationSuite) TestAddBranchSuccess(c *gc.C) {
 
 func (s *modelGenerationSuite) TestTrackBranchEntityTypeError(c *gc.C) {
 	defer s.setupModelGenerationAPI(c).Finish()
-	s.expectAssignAllUnits("ghost")
+	s.expectAssignUnits("ghost", 0)
 	s.expectAssignUnit("mysql/0")
 	s.expectBranch()
 
@@ -91,7 +91,7 @@ func (s *modelGenerationSuite) TestTrackBranchEntityTypeError(c *gc.C) {
 
 func (s *modelGenerationSuite) TestTrackBranchSuccess(c *gc.C) {
 	defer s.setupModelGenerationAPI(c).Finish()
-	s.expectAssignAllUnits("ghost")
+	s.expectAssignUnits("ghost", 0)
 	s.expectAssignUnit("mysql/0")
 	s.expectBranch()
 
@@ -110,6 +110,22 @@ func (s *modelGenerationSuite) TestTrackBranchSuccess(c *gc.C) {
 	})
 }
 
+func (s *modelGenerationSuite) TestTrackBranchWithTooManyNumUnits(c *gc.C) {
+	defer s.setupModelGenerationAPI(c).Finish()
+
+	arg := params.BranchTrackArg{
+		BranchName: s.newBranchName,
+		Entities: []params.Entity{
+			{Tag: names.NewUnitTag("mysql/0").String()},
+			{Tag: names.NewApplicationTag("ghost").String()},
+		},
+		NumUnits: 1,
+	}
+	result, err := s.api.TrackBranch(arg)
+	c.Assert(err, gc.ErrorMatches, "num of units allowed when specifying multiple entities")
+	c.Check(result.Results, gc.DeepEquals, []params.ErrorResult(nil))
+}
+
 func (s *modelGenerationSuite) TestCommitBranchSuccess(c *gc.C) {
 	defer s.setupModelGenerationAPI(c).Finish()
 	s.expectCommit()
@@ -263,6 +279,10 @@ func (s *modelGenerationSuite) expectAssignAllUnits(appName string) {
 	s.mockGen.EXPECT().AssignAllUnits(appName).Return(nil)
 }
 
+func (s *modelGenerationSuite) expectAssignUnits(appName string, numUnits int) {
+	s.mockGen.EXPECT().AssignUnits(appName, numUnits).Return(nil)
+}
+
 func (s *modelGenerationSuite) expectAssignUnit(unitName string) {
 	s.mockGen.EXPECT().AssignUnit(unitName).Return(nil)
 }

apiserver/facades/schema.json

@@ -23037,7 +23037,7 @@
     },
     {
         "Name": "ModelGeneration",
-        "Version": 2,
+        "Version": 3,
         "Schema": {
             "type": "object",
             "properties": {
@@ -23166,6 +23166,9 @@
                             "items": {
                                 "$ref": "#/definitions/Entity"
                             }
+                        },
+                        "num-units": {
+                            "type": "integer"
                         }
                     },
                     "additionalProperties": false,

apiserver/params/params.go

@@ -1198,6 +1198,7 @@ type BranchInfoArgs struct {
 type BranchTrackArg struct {
 	BranchName string   `json:"branch"`
 	Entities   []Entity `json:"entities"`
+	NumUnits   int      `json:"num-units,omitempty"`
 }
 
 // GenerationApplication represents changes to an application