feat(bigtable): update admin.go to support instance edition (#14479)

Adding support for new Editions type in admin.go

Author: NickASC4

bigtable/admin.go

@@ -1370,13 +1370,26 @@ const (
 	DEVELOPMENT              = InstanceType(btapb.Instance_DEVELOPMENT)
 )
 
+// Edition is the edition of the instance.
+type Edition int32
+
+const (
+	// EditionUnspecified defaults to ENTERPRISE
+	EditionUnspecified Edition = Edition(btapb.Instance_EDITION_UNSPECIFIED)
+	// Enterprise edition is the default edition.
+	Enterprise Edition = Edition(btapb.Instance_ENTERPRISE)
+	// EnterprisePlus edition is the edition with higher limits and more features than Enterprise edition.
+	EnterprisePlus Edition = Edition(btapb.Instance_ENTERPRISE_PLUS)
+)
+
 // InstanceInfo represents information about an instance
 type InstanceInfo struct {
 	Name          string // name of the instance
 	DisplayName   string // display name for UIs
 	InstanceState InstanceState
 	InstanceType  InstanceType
 	Labels        map[string]string
+	Edition       Edition
 }
 
 // InstanceConf contains the information necessary to create an Instance
@@ -1387,6 +1400,7 @@ type InstanceConf struct {
 	StorageType  StorageType
 	InstanceType InstanceType
 	Labels       map[string]string
+	Edition      Edition
 
 	// AutoscalingConfig configures the autoscaling properties on the cluster
 	// created with the instance. It is optional.
@@ -1409,6 +1423,7 @@ type InstanceWithClustersConfig struct {
 	Clusters                []ClusterConfig
 	InstanceType            InstanceType
 	Labels                  map[string]string
+	Edition                 Edition
 	// Tags maps TagKey resource names (e.g., "tagKeys/123") to TagValue
 	// resource names (e.g., "tagValues/456") to be associated with the instance.
 	Tags map[string]string
@@ -1426,6 +1441,7 @@ func (iac *InstanceAdminClient) CreateInstance(ctx context.Context, conf *Instan
 		InstanceType: conf.InstanceType,
 		Labels:       conf.Labels,
 		Tags:         conf.Tags,
+		Edition:      conf.Edition,
 		Clusters: []ClusterConfig{
 			{
 				InstanceID:        conf.InstanceId,
@@ -1458,6 +1474,7 @@ func (iac *InstanceAdminClient) CreateInstanceWithClusters(ctx context.Context,
 			Type:        btapb.Instance_Type(conf.InstanceType),
 			Labels:      conf.Labels,
 			Tags:        conf.Tags,
+			Edition:     btapb.Instance_Edition(conf.Edition),
 		},
 		Clusters: clusters,
 	}
@@ -1493,6 +1510,10 @@ func (iac *InstanceAdminClient) updateInstance(ctx context.Context, conf *Instan
 		ireq.Instance.Type = btapb.Instance_Type(conf.InstanceType)
 		mask.Paths = append(mask.Paths, "type")
 	}
+	if conf.Edition != EditionUnspecified {
+		ireq.Instance.Edition = btapb.Instance_Edition(conf.Edition)
+		mask.Paths = append(mask.Paths, "edition")
+	}
 	if conf.Labels != nil {
 		ireq.Instance.Labels = conf.Labels
 		mask.Paths = append(mask.Paths, "labels")
@@ -1601,6 +1622,7 @@ func (iac *InstanceAdminClient) Instances(ctx context.Context) ([]*InstanceInfo,
 			InstanceState: InstanceState(i.State),
 			InstanceType:  InstanceType(i.Type),
 			Labels:        i.Labels,
+			Edition:       Edition(i.Edition),
 		})
 	}
 	if len(res.FailedLocations) > 0 {
@@ -1637,6 +1659,7 @@ func (iac *InstanceAdminClient) InstanceInfo(ctx context.Context, instanceID str
 		InstanceState: InstanceState(res.State),
 		InstanceType:  InstanceType(res.Type),
 		Labels:        res.Labels,
+		Edition:       Edition(res.Edition),
 	}, nil
 }
 

bigtable/admin_test.go

@@ -1240,6 +1240,58 @@ func TestInstanceAdmin_CreateInstance_WithAutoscaling(t *testing.T) {
 	}
 }
 
+func TestInstanceAdmin_CreateInstance_WithEdition(t *testing.T) {
+	mock := &mockAdminClock{}
+	c := setupClient(t, mock)
+
+	err := c.CreateInstance(context.Background(), &InstanceConf{
+		InstanceId:   "myinst",
+		DisplayName:  "myinst",
+		InstanceType: PRODUCTION,
+		ClusterId:    "mycluster",
+		Zone:         "us-central1-a",
+		StorageType:  SSD,
+		Edition:      EnterprisePlus,
+	})
+	if err != nil {
+		t.Fatalf("CreateInstance failed: %v", err)
+	}
+
+	got := mock.createInstanceReq.Instance.Edition
+	want := btapb.Instance_ENTERPRISE_PLUS
+	if got != want {
+		t.Errorf("got edition %v, want %v", got, want)
+	}
+}
+
+func TestInstanceAdmin_CreateInstanceWithClusters_WithEdition(t *testing.T) {
+	mock := &mockAdminClock{}
+	c := setupClient(t, mock)
+
+	err := c.CreateInstanceWithClusters(context.Background(), &InstanceWithClustersConfig{
+		InstanceID:   "myinst",
+		DisplayName:  "myinst",
+		InstanceType: PRODUCTION,
+		Edition:      EnterprisePlus,
+		Clusters: []ClusterConfig{
+			{
+				ClusterID:   "mycluster",
+				Zone:        "us-central1-a",
+				StorageType: SSD,
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("CreateInstanceWithClusters failed: %v", err)
+	}
+
+	got := mock.createInstanceReq.Instance.Edition
+	want := btapb.Instance_ENTERPRISE_PLUS
+	if got != want {
+		t.Errorf("got edition %v, want %v", got, want)
+	}
+}
+
 // Test that CreateInstance with a tags argument completes without error.
 // We cannot verify tag creation itself, as tags are not stored in the Instance metadata
 // and thus are not observable here.

bigtable/integration_test.go

@@ -2378,6 +2378,7 @@ func TestIntegration_AdminCreateInstance(t *testing.T) {
 		InstanceType: DEVELOPMENT,
 		Labels:       map[string]string{"test-label-key": "test-label-value"},
 		Tags:         map[string]string{"tagKeys/12345": "tagValues/6789"},
+		Edition:      Enterprise,
 	}
 
 	// CreateInstance can be flaky; retry before marking as failing.
@@ -2403,6 +2404,9 @@ func TestIntegration_AdminCreateInstance(t *testing.T) {
 	if got, want := iInfo.Labels, conf.Labels; !cmp.Equal(got, want) {
 		t.Fatalf("Labels: %v, want: %v", got, want)
 	}
+	if iInfo.Edition != Enterprise {
+		t.Fatalf("Instance edition is not Enterprise: %v", iInfo.Edition)
+	}
 
 	// Update everything we can about the instance in one call.
 	confWithClusters := &InstanceWithClustersConfig{