feat(bigtable): add update table metadata support (#6746)

* add updateTable function to support deletion protection
* add  setupTableClient to add tests for updateTable and createTable

bigtable/admin.go

@@ -275,6 +275,73 @@ func (ac *AdminClient) CreateColumnFamily(ctx context.Context, table, family str
 	return err
 }
 
+// DeletionProtection indicates whether the table is protected against data loss
+// i.e. when set to protected, deleting the table, the column families in the table,
+// and the instance containing the table would be prohibited.
+type DeletionProtection int
+
+// None indicates that deletion protection is unset
+// Protected indicates that deletion protection is enabled
+// Unprotected indicates that deletion protection is disabled
+const (
+	None DeletionProtection = iota
+	Protected
+	Unprotected
+)
+
+// UpdateTableConf contains all of the information necessary to update a table with column families.
+type UpdateTableConf struct {
+	tableID string
+	// deletionProtection can be unset, true or false
+	// set to true to make the table protected against data loss
+	deletionProtection DeletionProtection
+}
+
+// UpdateTableWithDeletionProtection updates a table with the given table ID and deletion protection parameter.
+func (ac *AdminClient) UpdateTableWithDeletionProtection(ctx context.Context, tableID string, deletionProtection DeletionProtection) error {
+	return ac.updateTableWithConf(ctx, &UpdateTableConf{tableID, deletionProtection})
+}
+
+// updateTableWithConf updates a table in the instance from the given configuration.
+// only deletion protection can be updated at this period.
+// table ID is required.
+func (ac *AdminClient) updateTableWithConf(ctx context.Context, conf *UpdateTableConf) error {
+	if conf.tableID == "" {
+		return errors.New("TableID is required")
+	}
+
+	if conf.deletionProtection == None {
+		return errors.New("deletion protection is required")
+	}
+
+	ctx = mergeOutgoingMetadata(ctx, ac.md)
+
+	updateMask := &field_mask.FieldMask{
+		Paths: []string{
+			"deletion_protection",
+		},
+	}
+
+	deletionProtection := true
+	if conf.deletionProtection == Unprotected {
+		deletionProtection = false
+	}
+
+	req := &btapb.UpdateTableRequest{
+		Table: &btapb.Table{
+			Name:               conf.tableID,
+			DeletionProtection: deletionProtection,
+		},
+		UpdateMask: updateMask,
+	}
+	lro, err := ac.tClient.UpdateTable(ctx, req)
+	if err != nil {
+		return err
+	}
+	// ignore the response table proto by passing in nil
+	return longrunning.InternalNewOperation(ac.lroClient, lro).Wait(ctx, nil)
+}
+
 // DeleteTable deletes a table and all of its data.
 func (ac *AdminClient) DeleteTable(ctx context.Context, table string) error {
 	ctx = mergeOutgoingMetadata(ctx, ac.md)

bigtable/admin_test.go

@@ -16,6 +16,8 @@ package bigtable
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -25,6 +27,95 @@ import (
 	"google.golang.org/protobuf/types/known/anypb"
 )
 
+type mockTableAdminClock struct {
+	btapb.BigtableTableAdminClient
+
+	updateTableReq   *btapb.UpdateTableRequest
+	updateTableError error
+}
+
+func (c *mockTableAdminClock) UpdateTable(
+	ctx context.Context, in *btapb.UpdateTableRequest, opts ...grpc.CallOption,
+) (*longrunning.Operation, error) {
+	c.updateTableReq = in
+	return &longrunning.Operation{
+		Done: true,
+		Result: &longrunning.Operation_Response{
+			Response: &anypb.Any{TypeUrl: "google.bigtable.admin.v2.Table"},
+		},
+	}, c.updateTableError
+}
+
+func setupTableClient(t *testing.T, ac btapb.BigtableTableAdminClient) *AdminClient {
+	ctx := context.Background()
+	c, err := NewAdminClient(ctx, "my-cool-project", "my-cool-instance")
+	if err != nil {
+		t.Fatalf("NewAdminClient failed: %v", err)
+	}
+	c.tClient = ac
+	return c
+}
+
+func TestTableAdmin_UpdateTable(t *testing.T) {
+	mock := &mockTableAdminClock{}
+	c := setupTableClient(t, mock)
+	deletionProtection := Protected
+
+	// Check if the deletion protection updates correctly
+	err := c.UpdateTableWithDeletionProtection(context.Background(), "My-table", deletionProtection)
+	if err != nil {
+		t.Errorf("UpdateTable failed: %v", err)
+	}
+	updateTableReq := mock.updateTableReq
+	if !cmp.Equal(updateTableReq.Table.Name, "My-table") {
+		t.Errorf("UpdateTableRequest does not match, TableID: %v", updateTableReq.Table.Name)
+	}
+	if !cmp.Equal(updateTableReq.Table.DeletionProtection, true) {
+		t.Errorf("UpdateTableRequest does not match, TableID: %v", updateTableReq.Table.Name)
+	}
+	if !cmp.Equal(updateTableReq.UpdateMask.Paths[0], "deletion_protection") {
+		t.Errorf("UpdateTableRequest does not match, TableID: %v", updateTableReq.Table.Name)
+	}
+}
+
+func TestTableAdmin_UpdateTable_WithError(t *testing.T) {
+	mock := &mockTableAdminClock{updateTableError: errors.New("update table failure error")}
+	c := setupTableClient(t, mock)
+	deletionProtection := Protected
+
+	// Check if the update fails when update table returns an error
+	err := c.UpdateTableWithDeletionProtection(context.Background(), "My-table", deletionProtection)
+
+	if fmt.Sprint(err) != "update table failure error" {
+		t.Errorf("UpdateTable updated by mistake: %v", err)
+	}
+}
+
+func TestTableAdmin_UpdateTable_TableID_NotProvided(t *testing.T) {
+	mock := &mockTableAdminClock{}
+	c := setupTableClient(t, mock)
+	deletionProtection := Protected
+
+	// Check if the update fails when TableID is not provided
+	err := c.UpdateTableWithDeletionProtection(context.Background(), "", deletionProtection)
+	if fmt.Sprint(err) != "TableID is required" {
+		t.Errorf("UpdateTable failed: %v", err)
+	}
+}
+
+func TestTableAdmin_UpdateTable_DeletionProtection_NotProvided(t *testing.T) {
+	mock := &mockTableAdminClock{}
+	c := setupTableClient(t, mock)
+	deletionProtection := None
+
+	// Check if the update fails when deletion protection is not provided
+	err := c.UpdateTableWithDeletionProtection(context.Background(), "My-table", deletionProtection)
+
+	if fmt.Sprint(err) != "deletion protection is required" {
+		t.Errorf("UpdateTable failed: %v", err)
+	}
+}
+
 type mockAdminClock struct {
 	btapb.BigtableInstanceAdminClient