From 0e6413db32f2a27269602fa88afa762abdb837c0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Barbosa=20Sampaio?= <andrebsamp@gmail.com>
Date: Tue, 14 May 2024 17:10:07 -0400
Subject: [PATCH] feat(bigtable): Add ignore_warnings flag to SetGcPolicy
 (#9372)

* feat(bigtable): Add ignore_warnings flag to SetGcPolicy

* feat(bigtable): Add GcPolicyOptions

* refactor(bigtable): Change options struct to interface

* refactor(bigtable): Add nil checks and rename

---------

Co-authored-by: Baha Aiman <bahaaiman@google.com>
---
 bigtable/admin.go      | 45 +++++++++++++++++++++++++++++++----
 bigtable/admin_test.go | 53 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/bigtable/admin.go b/bigtable/admin.go
index e8ffb6fa56b1..cd9d300326b8 100644
--- a/bigtable/admin.go
+++ b/bigtable/admin.go
@@ -659,23 +659,60 @@ func (ac *AdminClient) TableInfo(ctx context.Context, table string) (*TableInfo,
 	return ti, nil
 }
 
-// SetGCPolicy specifies which cells in a column family should be garbage collected.
-// GC executes opportunistically in the background; table reads may return data
-// matching the GC policy.
-func (ac *AdminClient) SetGCPolicy(ctx context.Context, table, family string, policy GCPolicy) error {
+type gcPolicySettings struct {
+	ignoreWarnings bool
+}
+
+// GCPolicyOption is the interface to change GC policy settings
+type GCPolicyOption interface {
+	apply(s *gcPolicySettings)
+}
+
+type ignoreWarnings bool
+
+func (w ignoreWarnings) apply(s *gcPolicySettings) {
+	s.ignoreWarnings = bool(w)
+}
+
+// IgnoreWarnings returns a gcPolicyOption that ignores safety checks when modifying the column families
+func IgnoreWarnings() GCPolicyOption {
+	return ignoreWarnings(true)
+}
+
+func (ac *AdminClient) setGCPolicy(ctx context.Context, table, family string, policy GCPolicy, opts ...GCPolicyOption) error {
 	ctx = mergeOutgoingMetadata(ctx, ac.md)
 	prefix := ac.instancePrefix()
+
+	s := gcPolicySettings{}
+	for _, opt := range opts {
+		if opt != nil {
+			opt.apply(&s)
+		}
+	}
 	req := &btapb.ModifyColumnFamiliesRequest{
 		Name: prefix + "/tables/" + table,
 		Modifications: []*btapb.ModifyColumnFamiliesRequest_Modification{{
 			Id:  family,
 			Mod: &btapb.ModifyColumnFamiliesRequest_Modification_Update{Update: &btapb.ColumnFamily{GcRule: policy.proto()}},
 		}},
+		IgnoreWarnings: s.ignoreWarnings,
 	}
 	_, err := ac.tClient.ModifyColumnFamilies(ctx, req)
 	return err
 }
 
+// SetGCPolicy specifies which cells in a column family should be garbage collected.
+// GC executes opportunistically in the background; table reads may return data
+// matching the GC policy.
+func (ac *AdminClient) SetGCPolicy(ctx context.Context, table, family string, policy GCPolicy) error {
+	return ac.SetGCPolicyWithOptions(ctx, table, family, policy)
+}
+
+// SetGCPolicyWithOptions is similar to SetGCPolicy but allows passing options
+func (ac *AdminClient) SetGCPolicyWithOptions(ctx context.Context, table, family string, policy GCPolicy, opts ...GCPolicyOption) error {
+	return ac.setGCPolicy(ctx, table, family, policy, opts...)
+}
+
 // DropRowRange permanently deletes a row range from the specified table.
 func (ac *AdminClient) DropRowRange(ctx context.Context, table, rowKeyPrefix string) error {
 	ctx = mergeOutgoingMetadata(ctx, ac.md)
diff --git a/bigtable/admin_test.go b/bigtable/admin_test.go
index 0806857d7697..cb1ae4c376c0 100644
--- a/bigtable/admin_test.go
+++ b/bigtable/admin_test.go
@@ -43,6 +43,8 @@ type mockTableAdminClock struct {
 	copyBackupReq   *btapb.CopyBackupRequest
 	copyBackupError error
 
+	modColumnReq *btapb.ModifyColumnFamiliesRequest
+
 	createAuthorizedViewReq   *btapb.CreateAuthorizedViewRequest
 	createAuthorizedViewError error
 	updateAuthorizedViewReq   *btapb.UpdateAuthorizedViewRequest
@@ -76,6 +78,12 @@ func (c *mockTableAdminClock) CopyBackup(
 	return nil, c.copyBackupError
 }
 
+func (c *mockTableAdminClock) ModifyColumnFamilies(
+	ctx context.Context, in *btapb.ModifyColumnFamiliesRequest, opts ...grpc.CallOption) (*btapb.Table, error) {
+	c.modColumnReq = in
+	return nil, nil
+}
+
 func (c *mockTableAdminClock) CreateAuthorizedView(
 	ctx context.Context, in *btapb.CreateAuthorizedViewRequest, opts ...grpc.CallOption,
 ) (*longrunning.Operation, error) {
@@ -357,6 +365,51 @@ func TestTableAdmin_UpdateTableDisableChangeStream(t *testing.T) {
 	}
 }
 
+func TestTableAdmin_SetGcPolicy(t *testing.T) {
+	for _, test := range []struct {
+		desc string
+		opts GCPolicyOption
+		want bool
+	}{
+		{
+			desc: "IgnoreWarnings: false",
+			want: false,
+		},
+		{
+			desc: "IgnoreWarnings: true",
+			opts: IgnoreWarnings(),
+			want: true,
+		},
+	} {
+
+		mock := &mockTableAdminClock{}
+		c := setupTableClient(t, mock)
+
+		err := c.SetGCPolicyWithOptions(context.Background(), "My-table", "cf1", NoGcPolicy(), test.opts)
+		if err != nil {
+			t.Fatalf("%v: Failed to set GC Policy: %v", test.desc, err)
+		}
+
+		modColumnReq := mock.modColumnReq
+		if modColumnReq.IgnoreWarnings != test.want {
+			t.Errorf("%v: IgnoreWarnings got: %v, want: %v", test.desc, modColumnReq.IgnoreWarnings, test.want)
+		}
+	}
+
+	mock := &mockTableAdminClock{}
+	c := setupTableClient(t, mock)
+
+	err := c.SetGCPolicy(context.Background(), "My-table", "cf1", NoGcPolicy())
+	if err != nil {
+		t.Fatalf("SetGCPolicy: Failed to set GC Policy: %v", err)
+	}
+
+	modColumnReq := mock.modColumnReq
+	if modColumnReq.IgnoreWarnings {
+		t.Errorf("SetGCPolicy: IgnoreWarnings should be set to false")
+	}
+}
+
 func TestTableAdmin_CreateAuthorizedView_DeletionProtection_Protected(t *testing.T) {
 	mock := &mockTableAdminClock{}
 	c := setupTableClient(t, mock)