From f19ffada53d45919e872bec7089f0a540a35755d Mon Sep 17 00:00:00 2001 From: Kimia Hemmatirad <110636670+kimihrr@users.noreply.github.com> Date: Fri, 7 Oct 2022 17:51:32 -0400 Subject: [PATCH] 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 | 67 +++++++++++++++++++++++++++++++ bigtable/admin_test.go | 91 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+) diff --git a/bigtable/admin.go b/bigtable/admin.go index 647f692a21fc..0fe58360232b 100644 --- a/bigtable/admin.go +++ b/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) diff --git a/bigtable/admin_test.go b/bigtable/admin_test.go index 4c708e5f514a..471bada719ba 100644 --- a/bigtable/admin_test.go +++ b/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