feat(bigtable): add create table metadata support (#6813)

* add updateTable function to support deletion protection

Co-authored-by: Eric Schmidt <[email protected]>

bigtable/admin.go

@@ -211,18 +211,35 @@ func (ac *AdminClient) Tables(ctx context.Context) ([]string, error) {
 	return names, nil
 }
 
+// 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
+)
+
 // TableConf contains all of the information necessary to create a table with column families.
 type TableConf struct {
 	TableID   string
 	SplitKeys []string
 	// Families is a map from family name to GCPolicy
 	Families map[string]GCPolicy
+	// DeletionProtection can be none, protected or unprotected
+	// set to protected to make the table protected against data loss
+	DeletionProtection DeletionProtection
 }
 
 // CreateTable creates a new table in the instance.
 // This method may return before the table's creation is complete.
 func (ac *AdminClient) CreateTable(ctx context.Context, table string) error {
-	return ac.CreateTableFromConf(ctx, &TableConf{TableID: table})
+	return ac.CreateTableFromConf(ctx, &TableConf{TableID: table, DeletionProtection: None})
 }
 
 // CreatePresplitTable creates a new table in the instance.
@@ -231,17 +248,27 @@ func (ac *AdminClient) CreateTable(ctx context.Context, table string) error {
 // spanning the key ranges: [, s1), [s1, s2), [s2, ).
 // This method may return before the table's creation is complete.
 func (ac *AdminClient) CreatePresplitTable(ctx context.Context, table string, splitKeys []string) error {
-	return ac.CreateTableFromConf(ctx, &TableConf{TableID: table, SplitKeys: splitKeys})
+	return ac.CreateTableFromConf(ctx, &TableConf{TableID: table, SplitKeys: splitKeys, DeletionProtection: None})
 }
 
 // CreateTableFromConf creates a new table in the instance from the given configuration.
 func (ac *AdminClient) CreateTableFromConf(ctx context.Context, conf *TableConf) error {
+	if conf.TableID == "" {
+		return errors.New("TableID is required")
+	}
 	ctx = mergeOutgoingMetadata(ctx, ac.md)
 	var reqSplits []*btapb.CreateTableRequest_Split
 	for _, split := range conf.SplitKeys {
 		reqSplits = append(reqSplits, &btapb.CreateTableRequest_Split{Key: []byte(split)})
 	}
 	var tbl btapb.Table
+	// we'd rather not set anything explicitly if users don't specify a value and let the server set the default value.
+	// if DeletionProtection is not set, currently the API will default it to false.
+	if conf.DeletionProtection == Protected {
+		tbl.DeletionProtection = true
+	} else if conf.DeletionProtection == Unprotected {
+		tbl.DeletionProtection = false
+	}
 	if conf.Families != nil {
 		tbl.ColumnFamilies = make(map[string]*btapb.ColumnFamily)
 		for fam, policy := range conf.Families {
@@ -275,20 +302,6 @@ 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
@@ -326,20 +339,25 @@ func (ac *AdminClient) updateTableWithConf(ctx context.Context, conf *UpdateTabl
 	if conf.deletionProtection == Unprotected {
 		deletionProtection = false
 	}
-
+	prefix := ac.instancePrefix()
 	req := &btapb.UpdateTableRequest{
 		Table: &btapb.Table{
-			Name:               conf.tableID,
+			Name:               prefix + "/tables/" + conf.tableID,
 			DeletionProtection: deletionProtection,
 		},
 		UpdateMask: updateMask,
 	}
 	lro, err := ac.tClient.UpdateTable(ctx, req)
 	if err != nil {
-		return err
+		return fmt.Errorf("error from update: %w", err)
 	}
-	// ignore the response table proto by passing in nil
-	return longrunning.InternalNewOperation(ac.lroClient, lro).Wait(ctx, nil)
+	var tbl btapb.Table
+	op := longrunning.InternalNewOperation(ac.lroClient, lro)
+	err = op.Wait(ctx, &tbl)
+	if err != nil {
+		return fmt.Errorf("error from operation: %v", err)
+	}
+	return nil
 }
 
 // DeleteTable deletes a table and all of its data.
@@ -373,6 +391,10 @@ type TableInfo struct {
 	// DEPRECATED - This field is deprecated. Please use FamilyInfos instead.
 	Families    []string
 	FamilyInfos []FamilyInfo
+	// DeletionProtection indicates whether the table is protected against data loss
+	// DeletionProtection could be None depending on the table view
+	// for example when using NAME_ONLY, the response does not contain DeletionProtection and the value should be None
+	DeletionProtection DeletionProtection
 }
 
 // FamilyInfo represents information about a column family.
@@ -421,6 +443,13 @@ func (ac *AdminClient) TableInfo(ctx context.Context, table string) (*TableInfo,
 			FullGCPolicy: gcRuleToPolicy(fam.GcRule),
 		})
 	}
+	// we expect DeletionProtection to be in the response because Table_SCHEMA_VIEW is being used in this function
+	// but when using NAME_ONLY, the response does not contain DeletionProtection and it could be nil
+	if res.DeletionProtection == true {
+		ti.DeletionProtection = Protected
+	} else {
+		ti.DeletionProtection = Unprotected
+	}
 	return ti, nil
 }
 

bigtable/admin_test.go

@@ -30,10 +30,19 @@ import (
 type mockTableAdminClock struct {
 	btapb.BigtableTableAdminClient
 
+	createTableReq   *btapb.CreateTableRequest
 	updateTableReq   *btapb.UpdateTableRequest
+	createTableResp  *btapb.Table
 	updateTableError error
 }
 
+func (c *mockTableAdminClock) CreateTable(
+	ctx context.Context, in *btapb.CreateTableRequest, opts ...grpc.CallOption,
+) (*btapb.Table, error) {
+	c.createTableReq = in
+	return c.createTableResp, nil
+}
+
 func (c *mockTableAdminClock) UpdateTable(
 	ctx context.Context, in *btapb.UpdateTableRequest, opts ...grpc.CallOption,
 ) (*longrunning.Operation, error) {
@@ -56,18 +65,54 @@ func setupTableClient(t *testing.T, ac btapb.BigtableTableAdminClient) *AdminCli
 	return c
 }
 
-func TestTableAdmin_UpdateTable(t *testing.T) {
+func TestTableAdmin_CreateTableFromConf_DeletionProtection_Protected(t *testing.T) {
+	mock := &mockTableAdminClock{}
+	c := setupTableClient(t, mock)
+
+	deletionProtection := Protected
+	err := c.CreateTableFromConf(context.Background(), &TableConf{TableID: "My-table", DeletionProtection: deletionProtection})
+	if err != nil {
+		t.Fatalf("CreateTableFromConf failed: %v", err)
+	}
+	createTableReq := mock.createTableReq
+	if !cmp.Equal(createTableReq.TableId, "My-table") {
+		t.Errorf("Unexpected table ID: %v, expected %v", createTableReq.TableId, "My-table")
+	}
+	if !cmp.Equal(createTableReq.Table.DeletionProtection, true) {
+		t.Errorf("Unexpected table deletion protection: %v, expected %v", createTableReq.Table.DeletionProtection, true)
+	}
+}
+
+func TestTableAdmin_CreateTableFromConf_DeletionProtection_Unprotected(t *testing.T) {
+	mock := &mockTableAdminClock{}
+	c := setupTableClient(t, mock)
+
+	deletionProtection := Unprotected
+	err := c.CreateTableFromConf(context.Background(), &TableConf{TableID: "My-table", DeletionProtection: deletionProtection})
+	if err != nil {
+		t.Fatalf("CreateTableFromConf failed: %v", err)
+	}
+	createTableReq := mock.createTableReq
+	if !cmp.Equal(createTableReq.TableId, "My-table") {
+		t.Errorf("Unexpected table ID: %v, expected %v", createTableReq.TableId, "My-table")
+	}
+	if !cmp.Equal(createTableReq.Table.DeletionProtection, false) {
+		t.Errorf("Unexpected table deletion protection: %v, expected %v", createTableReq.Table.DeletionProtection, false)
+	}
+}
+
+func TestTableAdmin_UpdateTableWithDeletionProtection(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)
+		t.Fatalf("UpdateTableWithDeletionProtection failed: %v", err)
 	}
 	updateTableReq := mock.updateTableReq
-	if !cmp.Equal(updateTableReq.Table.Name, "My-table") {
+	if !cmp.Equal(updateTableReq.Table.Name, "projects/my-cool-project/instances/my-cool-instance/tables/My-table") {
 		t.Errorf("UpdateTableRequest does not match, TableID: %v", updateTableReq.Table.Name)
 	}
 	if !cmp.Equal(updateTableReq.Table.DeletionProtection, true) {
@@ -86,8 +131,8 @@ func TestTableAdmin_UpdateTable_WithError(t *testing.T) {
 	// 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)
+	if fmt.Sprint(err) != "error from update: update table failure error" {
+		t.Fatalf("UpdateTable updated by mistake: %v", err)
 	}
 }
 
@@ -99,7 +144,7 @@ func TestTableAdmin_UpdateTable_TableID_NotProvided(t *testing.T) {
 	// 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)
+		t.Fatalf("UpdateTable failed: %v", err)
 	}
 }
 
@@ -112,7 +157,7 @@ func TestTableAdmin_UpdateTable_DeletionProtection_NotProvided(t *testing.T) {
 	err := c.UpdateTableWithDeletionProtection(context.Background(), "My-table", deletionProtection)
 
 	if fmt.Sprint(err) != "deletion protection is required" {
-		t.Errorf("UpdateTable failed: %v", err)
+		t.Fatalf("UpdateTable failed: %v", err)
 	}
 }
 

bigtable/integration_test.go

@@ -1148,6 +1148,82 @@ func TestIntegration_SampleRowKeys(t *testing.T) {
 	}
 }
 
+// testing if deletionProtection works properly e.g. when set to Protected, column family and table cannot be deleted;
+// then update the deletionProtection to Unprotected and check if deleting the column family and table works properly.
+func TestIntegration_TableDeletionProtection(t *testing.T) {
+	testEnv, err := NewIntegrationEnv()
+	if err != nil {
+		t.Fatalf("IntegrationEnv: %v", err)
+	}
+	defer testEnv.Close()
+
+	timeout := 2 * time.Second
+	if testEnv.Config().UseProd {
+		timeout = 5 * time.Minute
+	}
+	ctx, _ := context.WithTimeout(context.Background(), timeout)
+
+	adminClient, err := testEnv.NewAdminClient()
+	if err != nil {
+		t.Fatalf("NewAdminClient: %v", err)
+	}
+	defer adminClient.Close()
+
+	tableConf := TableConf{
+		TableID: myTableName,
+		Families: map[string]GCPolicy{
+			"fam1": MaxVersionsPolicy(1),
+			"fam2": MaxVersionsPolicy(2),
+		},
+		DeletionProtection: Protected,
+	}
+
+	if err := adminClient.CreateTableFromConf(ctx, &tableConf); err != nil {
+		t.Fatalf("Create table from config: %v", err)
+	}
+
+	table, err := adminClient.TableInfo(ctx, myTableName)
+	if err != nil {
+		t.Fatalf("Getting table info: %v", err)
+	}
+
+	if table.DeletionProtection != Protected {
+		t.Errorf("Expect table deletion protection to be enabled for table: %v", myTableName)
+	}
+
+	// Check if the deletion protection works properly
+	if err = adminClient.DeleteColumnFamily(ctx, tableConf.TableID, "fam1"); err == nil {
+		t.Errorf("We shouldn't be able to delete the column family fam1 when the deletion protection is enabled for table %v", myTableName)
+	}
+	if err = adminClient.DeleteTable(ctx, tableConf.TableID); err == nil {
+		t.Errorf("We shouldn't be able to delete the table when the deletion protection is enabled for table %v", myTableName)
+	}
+
+	updateTableConf := UpdateTableConf{
+		tableID:            myTableName,
+		deletionProtection: Unprotected,
+	}
+	if err := adminClient.updateTableWithConf(ctx, &updateTableConf); err != nil {
+		t.Fatalf("Update table from config: %v", err)
+	}
+
+	table, err = adminClient.TableInfo(ctx, myTableName)
+	if err != nil {
+		t.Fatalf("Getting table info: %v", err)
+	}
+
+	if table.DeletionProtection != Unprotected {
+		t.Errorf("Expect table deletion protection to be disabled for table: %v", myTableName)
+	}
+
+	if err := adminClient.DeleteColumnFamily(ctx, tableConf.TableID, "fam1"); err != nil {
+		t.Errorf("Delete column family does not work properly while deletion protection bit is disabled: %v", err)
+	}
+	if err = adminClient.DeleteTable(ctx, tableConf.TableID); err != nil {
+		t.Errorf("Deleting the table does not work properly while deletion protection bit is disabled: %v", err)
+	}
+}
+
 func TestIntegration_Admin(t *testing.T) {
 	testEnv, err := NewIntegrationEnv()
 	if err != nil {