feat(bigtable): add support for autoscaling

bigtable/admin.go

@@ -811,7 +811,7 @@ type InstanceConf struct {
 
 	// AutoscalingConfig configures the autoscaling properties on the cluster
 	// created with the instance. It is optional.
-	AutoscalingConfig AutoscalingConfig
+	AutoscalingConfig *AutoscalingConfig
 }
 
 // InstanceWithClustersConfig contains the information necessary to create an Instance
@@ -828,7 +828,7 @@ var instanceNameRegexp = regexp.MustCompile(`^projects/([^/]+)/instances/([a-z][
 // This method will return when the instance has been created or when an error occurs.
 func (iac *InstanceAdminClient) CreateInstance(ctx context.Context, conf *InstanceConf) error {
 	ctx = mergeOutgoingMetadata(ctx, iac.md)
-	newConfig := InstanceWithClustersConfig{
+	newConfig := &InstanceWithClustersConfig{
 		InstanceID:   conf.InstanceId,
 		DisplayName:  conf.DisplayName,
 		InstanceType: conf.InstanceType,
@@ -844,7 +844,7 @@ func (iac *InstanceAdminClient) CreateInstance(ctx context.Context, conf *Instan
 			},
 		},
 	}
-	return iac.CreateInstanceWithClusters(ctx, &newConfig)
+	return iac.CreateInstanceWithClusters(ctx, newConfig)
 }
 
 // CreateInstanceWithClusters creates a new instance with configured clusters in the project.
@@ -949,9 +949,9 @@ func (iac *InstanceAdminClient) UpdateInstanceWithClusters(ctx context.Context,
 	// Update any clusters
 	for _, cluster := range conf.Clusters {
 		var clusterErr error
-		if !cluster.AutoscalingConfig.isZero() {
-			clusterErr = iac.SetAutoscaling(ctx, conf.InstanceID, cluster.ClusterID, cluster.AutoscalingConfig)
-		} else {
+		if cluster.AutoscalingConfig != nil {
+			clusterErr = iac.SetAutoscaling(ctx, conf.InstanceID, cluster.ClusterID, *cluster.AutoscalingConfig)
+		} else if cluster.NumNodes > 0 {
 			clusterErr = iac.UpdateCluster(ctx, conf.InstanceID, cluster.ClusterID, cluster.NumNodes)
 		}
 		if clusterErr != nil {
@@ -1048,34 +1048,28 @@ func (iac *InstanceAdminClient) InstanceInfo(ctx context.Context, instanceID str
 // AutoscalingConfig contains autoscaling configuration for a cluster.
 // For details, see https://cloud.google.com/bigtable/docs/autoscaling.
 type AutoscalingConfig struct {
-	// MinNodes sets the minumum number of nodes in a cluster.
+	// MinNodes sets the minumum number of nodes in a cluster. MinNodes must
+	// be 1 or greater.
 	MinNodes int
-	// MaxNodes sets the maximum number of nodes in a cluster.
+	// MaxNodes sets the maximum number of nodes in a cluster. MaxNodes must be
+	// equal to or greater than MinNodes.
 	MaxNodes int
 	// CPUTargetPercent sets the CPU utilization target for your cluster's
 	// workload.
 	CPUTargetPercent int
 }
 
-func (a AutoscalingConfig) isZero() bool {
-	return a.MinNodes == 0 && a.MaxNodes == 0 && a.CPUTargetPercent == 0
-}
-
-func (a AutoscalingConfig) proto() *btapb.Cluster_ClusterConfig_ {
-	if a.isZero() {
+func (a *AutoscalingConfig) proto() *btapb.Cluster_ClusterAutoscalingConfig {
+	if a == nil {
 		return nil
 	}
-	return &btapb.Cluster_ClusterConfig_{
-		ClusterConfig: &btapb.Cluster_ClusterConfig{
-			ClusterAutoscalingConfig: &btapb.Cluster_ClusterAutoscalingConfig{
-				AutoscalingLimits: &btapb.AutoscalingLimits{
-					MinServeNodes: int32(a.MinNodes),
-					MaxServeNodes: int32(a.MaxNodes),
-				},
-				AutoscalingTargets: &btapb.AutoscalingTargets{
-					CpuUtilizationPercent: int32(a.CPUTargetPercent),
-				},
-			},
+	return &btapb.Cluster_ClusterAutoscalingConfig{
+		AutoscalingLimits: &btapb.AutoscalingLimits{
+			MinServeNodes: int32(a.MinNodes),
+			MaxServeNodes: int32(a.MaxNodes),
+		},
+		AutoscalingTargets: &btapb.AutoscalingTargets{
+			CpuUtilizationPercent: int32(a.CPUTargetPercent),
 		},
 	}
 }
@@ -1119,19 +1113,27 @@ type ClusterConfig struct {
 
 	// AutoscalingConfig configures the autoscaling properties on a cluster.
 	// One of NumNodes or AutoscalingConfig is required.
-	AutoscalingConfig AutoscalingConfig
+	AutoscalingConfig *AutoscalingConfig
 }
 
 func (cc *ClusterConfig) proto(project string) *btapb.Cluster {
-	return &btapb.Cluster{
+	cl := &btapb.Cluster{
 		ServeNodes:         cc.NumNodes,
 		DefaultStorageType: cc.StorageType.proto(),
 		Location:           "projects/" + project + "/locations/" + cc.Zone,
 		EncryptionConfig: &btapb.Cluster_EncryptionConfig{
 			KmsKeyName: cc.KMSKeyName,
 		},
-		Config: cc.AutoscalingConfig.proto(),
 	}
+
+	if asc := cc.AutoscalingConfig.proto(); asc != nil {
+		cl.Config = &btapb.Cluster_ClusterConfig_{
+			ClusterConfig: &btapb.Cluster_ClusterConfig{
+				ClusterAutoscalingConfig: cc.AutoscalingConfig.proto(),
+			},
+		}
+	}
+	return cl
 }
 
 // ClusterInfo represents information about a cluster.
@@ -1155,7 +1157,7 @@ type ClusterInfo struct {
 	KMSKeyName string
 
 	// AutoscalingConfig are the configured values for a cluster.
-	AutoscalingConfig AutoscalingConfig
+	AutoscalingConfig *AutoscalingConfig
 }
 
 // CreateCluster creates a new cluster in an instance.
@@ -1186,12 +1188,16 @@ func (iac *InstanceAdminClient) DeleteCluster(ctx context.Context, instanceID, c
 }
 
 // SetAutoscaling enables autoscaling on a cluster. To remove autoscaling, use
-// UpdateCluster.
+// UpdateCluster. See AutoscalingConfig documentation for deatils.
 func (iac *InstanceAdminClient) SetAutoscaling(ctx context.Context, instanceID, clusterID string, conf AutoscalingConfig) error {
 	ctx = mergeOutgoingMetadata(ctx, iac.md)
 	cluster := &btapb.Cluster{
-		Name:   "projects/" + iac.project + "/instances/" + instanceID + "/clusters/" + clusterID,
-		Config: conf.proto(),
+		Name: "projects/" + iac.project + "/instances/" + instanceID + "/clusters/" + clusterID,
+		Config: &btapb.Cluster_ClusterConfig_{
+			ClusterConfig: &btapb.Cluster_ClusterConfig{
+				ClusterAutoscalingConfig: conf.proto(),
+			},
+		},
 	}
 	lro, err := iac.iClient.PartialUpdateCluster(ctx, &btapb.PartialUpdateClusterRequest{
 		UpdateMask: &field_mask.FieldMask{
@@ -1262,10 +1268,9 @@ func (iac *InstanceAdminClient) Clusters(ctx context.Context, instanceID string)
 			StorageType: storageTypeFromProto(c.DefaultStorageType),
 			KMSKeyName:  kmsKeyName,
 		}
-		// Use type assertion to handle protobuf oneof type
-		if cfg, ok := c.Config.(*btapb.Cluster_ClusterConfig_); ok {
-			if asc, ok := autoscalingConfig(cfg); ok {
-				ci.AutoscalingConfig = asc
+		if cfg := c.GetClusterConfig(); cfg != nil {
+			if asc, ok := fromClusterConfigProto(cfg); ok {
+				ci.AutoscalingConfig = &asc
 			}
 		}
 		cis = append(cis, ci)
@@ -1309,22 +1314,22 @@ func (iac *InstanceAdminClient) GetCluster(ctx context.Context, instanceID, clus
 		KMSKeyName:  kmsKeyName,
 	}
 	// Use type assertion to handle protobuf oneof type
-	if cfg, ok := c.Config.(*btapb.Cluster_ClusterConfig_); ok {
-		if asc, ok := autoscalingConfig(cfg); ok {
-			ci.AutoscalingConfig = asc
+	if cfg := c.GetClusterConfig(); cfg != nil {
+		if asc, ok := fromClusterConfigProto(cfg); ok {
+			ci.AutoscalingConfig = &asc
 		}
 	}
 	return ci, nil
 }
 
-func autoscalingConfig(c *btapb.Cluster_ClusterConfig_) (AutoscalingConfig, bool) {
-	if c.ClusterConfig == nil {
+func fromClusterConfigProto(c *btapb.Cluster_ClusterConfig) (AutoscalingConfig, bool) {
+	if c == nil {
 		return AutoscalingConfig{}, false
 	}
-	if c.ClusterConfig.ClusterAutoscalingConfig == nil {
+	if c.ClusterAutoscalingConfig == nil {
 		return AutoscalingConfig{}, false
 	}
-	got := c.ClusterConfig.ClusterAutoscalingConfig
+	got := c.ClusterAutoscalingConfig
 	if got.AutoscalingLimits == nil || got.AutoscalingTargets == nil {
 		return AutoscalingConfig{}, false
 	}
@@ -1643,7 +1648,7 @@ func UpdateInstanceAndSyncClusters(ctx context.Context, iac *InstanceAdminClient
 		}
 		delete(existingClusterNames, cluster.ClusterID)
 
-		if cluster.NumNodes <= 0 && cluster.AutoscalingConfig.isZero() {
+		if cluster.NumNodes <= 0 && cluster.AutoscalingConfig == nil {
 			// We only synchronize clusters with a valid number of nodes
 			// or a valid autoscaling config.
 			continue
@@ -1652,9 +1657,9 @@ func UpdateInstanceAndSyncClusters(ctx context.Context, iac *InstanceAdminClient
 		// We update teh clusters autoscaling config, or its number of serve
 		// nodes.
 		var updateErr error
-		if !cluster.AutoscalingConfig.isZero() {
+		if cluster.AutoscalingConfig != nil {
 			updateErr = iac.SetAutoscaling(ctx, conf.InstanceID, cluster.ClusterID,
-				cluster.AutoscalingConfig)
+				*cluster.AutoscalingConfig)
 		} else {
 			updateErr = iac.UpdateCluster(ctx, conf.InstanceID, cluster.ClusterID,
 				cluster.NumNodes)

bigtable/admin_test.go

@@ -1,3 +1,17 @@
+// Copyright 2022 Google LLC
+//
+// 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 bigtable
 
 import (
@@ -87,7 +101,7 @@ func setupClient(t *testing.T, ac btapb.BigtableInstanceAdminClient) *InstanceAd
 func TestInstanceAdmin_GetCluster(t *testing.T) {
 	tcs := []struct {
 		cluster    *btapb.Cluster
-		wantConfig AutoscalingConfig
+		wantConfig *AutoscalingConfig
 		desc       string
 	}{
 		{
@@ -98,7 +112,7 @@ func TestInstanceAdmin_GetCluster(t *testing.T) {
 				State:              btapb.Cluster_READY,
 				DefaultStorageType: btapb.StorageType_SSD,
 			},
-			wantConfig: AutoscalingConfig{},
+			wantConfig: nil,
 		},
 		{
 			desc: "when autoscaling is enabled",
@@ -121,7 +135,7 @@ func TestInstanceAdmin_GetCluster(t *testing.T) {
 					},
 				},
 			},
-			wantConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+			wantConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 		},
 	}
 
@@ -134,7 +148,7 @@ func TestInstanceAdmin_GetCluster(t *testing.T) {
 				t.Fatalf("GetCluster failed: %v", err)
 			}
 
-			if gotConfig := info.AutoscalingConfig; gotConfig != tc.wantConfig {
+			if gotConfig := info.AutoscalingConfig; !cmp.Equal(gotConfig, tc.wantConfig) {
 				t.Fatalf("want autoscaling config = %v, got = %v", tc.wantConfig, gotConfig)
 			}
 		})
@@ -144,7 +158,7 @@ func TestInstanceAdmin_GetCluster(t *testing.T) {
 func TestInstanceAdmin_Clusters(t *testing.T) {
 	tcs := []struct {
 		cluster    *btapb.Cluster
-		wantConfig AutoscalingConfig
+		wantConfig *AutoscalingConfig
 		desc       string
 	}{
 		{
@@ -155,7 +169,7 @@ func TestInstanceAdmin_Clusters(t *testing.T) {
 				State:              btapb.Cluster_READY,
 				DefaultStorageType: btapb.StorageType_SSD,
 			},
-			wantConfig: AutoscalingConfig{},
+			wantConfig: nil,
 		},
 		{
 			desc: "when autoscaling is enabled",
@@ -178,7 +192,7 @@ func TestInstanceAdmin_Clusters(t *testing.T) {
 					},
 				},
 			},
-			wantConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+			wantConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 		},
 	}
 
@@ -195,7 +209,7 @@ func TestInstanceAdmin_Clusters(t *testing.T) {
 			}
 
 			info := infos[0]
-			if gotConfig := info.AutoscalingConfig; gotConfig != tc.wantConfig {
+			if gotConfig := info.AutoscalingConfig; !cmp.Equal(gotConfig, tc.wantConfig) {
 				t.Fatalf("want autoscaling config = %v, got = %v", tc.wantConfig, gotConfig)
 			}
 		})
@@ -274,7 +288,7 @@ func TestInstanceAdmin_CreateInstance_WithAutoscaling(t *testing.T) {
 		ClusterId:         "mycluster",
 		Zone:              "us-central1-a",
 		StorageType:       SSD,
-		AutoscalingConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+		AutoscalingConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 	})
 	if err != nil {
 		t.Fatalf("CreateInstance failed: %v", err)
@@ -314,7 +328,7 @@ func TestInstanceAdmin_CreateInstance_WithAutoscaling(t *testing.T) {
 
 	// omitting autoscaling config results in a nil config in the request
 	mycc = mock.createInstanceReq.Clusters["mycluster"]
-	if cc := mycc.Config.(*btapb.Cluster_ClusterConfig_); cc != nil {
+	if cc := mycc.GetClusterConfig(); cc != nil {
 		t.Fatalf("want config = nil, got = %v", gotConfig)
 	}
 }
@@ -332,7 +346,7 @@ func TestInstanceAdmin_CreateInstanceWithClusters_WithAutoscaling(t *testing.T)
 				ClusterID:         "mycluster",
 				Zone:              "us-central1-a",
 				StorageType:       SSD,
-				AutoscalingConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+				AutoscalingConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 			},
 		},
 	})
@@ -368,7 +382,7 @@ func TestInstanceAdmin_CreateCluster_WithAutoscaling(t *testing.T) {
 		ClusterID:         "mycluster",
 		Zone:              "us-central1-a",
 		StorageType:       SSD,
-		AutoscalingConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+		AutoscalingConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 	})
 	if err != nil {
 		t.Fatalf("CreateCluster failed: %v", err)
@@ -403,11 +417,37 @@ func TestInstanceAdmin_CreateCluster_WithAutoscaling(t *testing.T) {
 	}
 
 	// omitting autoscaling config results in a nil config in the request
-	if cc := mock.createClusterReq.Cluster.Config.(*btapb.Cluster_ClusterConfig_); cc != nil {
+	if cc := mock.createClusterReq.Cluster.GetClusterConfig(); cc != nil {
 		t.Fatalf("want config = nil, got = %v", gotConfig)
 	}
 }
 
+func TestInstanceAdmin_UpdateInstanceWithClusters_IgnoresInvalidClusters(t *testing.T) {
+	mock := &mockAdminClock{}
+	c := setupClient(t, mock)
+
+	err := c.UpdateInstanceWithClusters(context.Background(), &InstanceWithClustersConfig{
+		InstanceID:  "myinst",
+		DisplayName: "myinst",
+		Clusters: []ClusterConfig{
+			{
+				ClusterID: "mycluster",
+				Zone:      "us-central1-a",
+				// Cluster has no autoscaling or num nodes
+				// It should be ignored
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("UpdateInstanceWithClusters failed: %v", err)
+	}
+
+	if mock.partialUpdateClusterReq != nil {
+		t.Fatalf("PartialUpdateCluster should not have been called, got = %v",
+			mock.partialUpdateClusterReq)
+	}
+}
+
 func TestInstanceAdmin_UpdateInstanceWithClusters_WithAutoscaling(t *testing.T) {
 	mock := &mockAdminClock{}
 	c := setupClient(t, mock)
@@ -419,7 +459,7 @@ func TestInstanceAdmin_UpdateInstanceWithClusters_WithAutoscaling(t *testing.T)
 			{
 				ClusterID:         "mycluster",
 				Zone:              "us-central1-a",
-				AutoscalingConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+				AutoscalingConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 			},
 		},
 	})
@@ -498,7 +538,7 @@ func TestInstanceAdmin_UpdateInstanceAndSyncClusters_WithAutoscaling(t *testing.
 			{
 				ClusterID:         "mycluster",
 				Zone:              "us-central1-a",
-				AutoscalingConfig: AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
+				AutoscalingConfig: &AutoscalingConfig{MinNodes: 1, MaxNodes: 2, CPUTargetPercent: 10},
 			},
 		},
 	})