feat: add support for table deletion protection (#2430)

adashyan · web-flow · commit 687b6df14b74 · 2024-11-25T09:25:52.000-05:00
* add support for table deletion protection

* fear: add support for table deletion protection

* stily fix

* stily fix

* styling fix

* rename isProtected to isDeletionProtected
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateTableRequest.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateTableRequest.java
@@ -123,6 +123,12 @@ public CreateTableRequest addChangeStreamRetention(Duration retention) {
     return this;
   }
 
+  /** Configures if the table is deletion protected. */
+  public CreateTableRequest setDeletionProtection(boolean deletionProtection) {
+    requestBuilder.getTableBuilder().setDeletionProtection(deletionProtection);
+    return this;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/Table.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/Table.java
@@ -105,6 +105,7 @@ public com.google.bigtable.admin.v2.Table.ClusterState.ReplicationState toProto(
   private final List<ColumnFamily> columnFamilies;
 
   private final Duration changeStreamRetention;
+  private final boolean deletionProtection;
 
   @InternalApi
   public static Table fromProto(@Nonnull com.google.bigtable.admin.v2.Table proto) {
@@ -135,19 +136,22 @@ public static Table fromProto(@Nonnull com.google.bigtable.admin.v2.Table proto)
         TableName.parse(proto.getName()),
         replicationStates.build(),
         columnFamilies.build(),
-        changeStreamConfig);
+        changeStreamConfig,
+        proto.getDeletionProtection());
   }
 
   private Table(
       TableName tableName,
       Map<String, ReplicationState> replicationStatesByClusterId,
       List<ColumnFamily> columnFamilies,
-      Duration changeStreamRetention) {
+      Duration changeStreamRetention,
+      boolean deletionProtection) {
     this.instanceId = tableName.getInstance();
     this.id = tableName.getTable();
     this.replicationStatesByClusterId = replicationStatesByClusterId;
     this.columnFamilies = columnFamilies;
     this.changeStreamRetention = changeStreamRetention;
+    this.deletionProtection = deletionProtection;
   }
 
   /** Gets the table's id. */
@@ -172,6 +176,11 @@ public Duration getChangeStreamRetention() {
     return changeStreamRetention;
   }
 
+  /** Returns whether this table is deletion protected. */
+  public boolean isDeletionProtected() {
+    return deletionProtection;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
@@ -185,12 +194,18 @@ public boolean equals(Object o) {
         && Objects.equal(instanceId, table.instanceId)
         && Objects.equal(replicationStatesByClusterId, table.replicationStatesByClusterId)
         && Objects.equal(columnFamilies, table.columnFamilies)
-        && Objects.equal(changeStreamRetention, table.changeStreamRetention);
+        && Objects.equal(changeStreamRetention, table.changeStreamRetention)
+        && Objects.equal(deletionProtection, table.deletionProtection);
   }
 
   @Override
   public int hashCode() {
     return Objects.hashCode(
-        id, instanceId, replicationStatesByClusterId, columnFamilies, changeStreamRetention);
+        id,
+        instanceId,
+        replicationStatesByClusterId,
+        columnFamilies,
+        changeStreamRetention,
+        deletionProtection);
   }
 }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/UpdateTableRequest.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/UpdateTableRequest.java
@@ -74,6 +74,13 @@ public UpdateTableRequest disableChangeStreamRetention() {
     return addChangeStreamRetention(Duration.ZERO);
   }
 
+  /** Changes the deletion protection of an existing table. */
+  public UpdateTableRequest setDeletionProtection(boolean deletionProtection) {
+    requestBuilder.getTableBuilder().setDeletionProtection(deletionProtection);
+    requestBuilder.getUpdateMaskBuilder().addPaths("deletion_protection");
+    return this;
+  }
+
   @InternalApi
   public com.google.bigtable.admin.v2.UpdateTableRequest toProto(
       String projectId, String instanceId) {
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java
@@ -299,6 +299,43 @@ public void testCreateTable() {
     assertThat(result).isEqualTo(Table.fromProto(expectedResponse));
   }
 
+  @Test
+  public void testCreateTableWithDeletionProtectionSet() {
+    // Setup
+    Mockito.when(mockStub.createTableCallable()).thenReturn(mockCreateTableCallable);
+
+    com.google.bigtable.admin.v2.CreateTableRequest expectedRequest =
+        com.google.bigtable.admin.v2.CreateTableRequest.newBuilder()
+            .setParent(INSTANCE_NAME)
+            .setTableId(TABLE_ID)
+            .setTable(
+                com.google.bigtable.admin.v2.Table.newBuilder()
+                    .setDeletionProtection(true)
+                    .putColumnFamilies(
+                        "cf1",
+                        ColumnFamily.newBuilder()
+                            .setGcRule(GcRule.getDefaultInstance())
+                            .setValueType(TypeProtos.intSumType())
+                            .build()))
+            .build();
+
+    com.google.bigtable.admin.v2.Table expectedResponse =
+        com.google.bigtable.admin.v2.Table.newBuilder().setName(TABLE_NAME).build();
+
+    Mockito.when(mockCreateTableCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // Execute
+    Table result =
+        adminClient.createTable(
+            CreateTableRequest.of(TABLE_ID)
+                .addFamily("cf1", Type.int64Sum())
+                .setDeletionProtection(true));
+
+    // Verify
+    assertThat(result).isEqualTo(Table.fromProto(expectedResponse));
+  }
+
   @Test
   public void testUpdateTable() {
     // Setup
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/CreateTableRequestTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/CreateTableRequestTest.java
@@ -48,7 +48,8 @@ public void testToProto() {
             .addFamily("another-family", GCRULES.maxAge(100, TimeUnit.HOURS))
             .addSplit(splitKey)
             .addSplit(secondSplitKey)
-            .addChangeStreamRetention(Duration.ofHours(24));
+            .addChangeStreamRetention(Duration.ofHours(24))
+            .setDeletionProtection(true);
 
     com.google.bigtable.admin.v2.CreateTableRequest requestProto =
         com.google.bigtable.admin.v2.CreateTableRequest.newBuilder()
@@ -70,7 +71,8 @@ public void testToProto() {
                         ChangeStreamConfig.newBuilder()
                             .setRetentionPeriod(
                                 com.google.protobuf.Duration.newBuilder().setSeconds(86400))
-                            .build()))
+                            .build())
+                    .setDeletionProtection(true))
             .setParent(NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID))
             .addInitialSplits(
                 com.google.bigtable.admin.v2.CreateTableRequest.Split.newBuilder().setKey(splitKey))
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/TableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/TableTest.java
@@ -67,6 +67,7 @@ public void testFromProto() {
                                     .setSeconds(1)
                                     .setNanos(99)))
                     .build())
+            .setDeletionProtection(true)
             .build();
 
     Table result = Table.fromProto(proto);
@@ -78,6 +79,7 @@ public void testFromProto() {
             "cluster1", Table.ReplicationState.READY,
             "cluster2", Table.ReplicationState.INITIALIZING);
     assertThat(result.getColumnFamilies()).hasSize(3);
+    assertThat(result.isDeletionProtected()).isTrue();
 
     for (Entry<String, ColumnFamily> entry : proto.getColumnFamiliesMap().entrySet()) {
       assertThat(result.getColumnFamilies())
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/UpdateTableRequestTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/models/UpdateTableRequestTest.java
@@ -81,4 +81,36 @@ public void testNoChangeChangeStreamToProto() {
             .build();
     assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
   }
+
+  @Test
+  public void testEnableDeletionProtection() {
+    UpdateTableRequest request = UpdateTableRequest.of(TABLE_ID).setDeletionProtection(true);
+
+    com.google.bigtable.admin.v2.UpdateTableRequest requestProto =
+        com.google.bigtable.admin.v2.UpdateTableRequest.newBuilder()
+            .setTable(
+                Table.newBuilder()
+                    .setName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+                    .setDeletionProtection(true))
+            .setUpdateMask(FieldMask.newBuilder().addPaths("deletion_protection").build())
+            .build();
+
+    assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
+  }
+
+  @Test
+  public void testDisableDeletionProtection() {
+    UpdateTableRequest request = UpdateTableRequest.of(TABLE_ID).setDeletionProtection(false);
+
+    com.google.bigtable.admin.v2.UpdateTableRequest requestProto =
+        com.google.bigtable.admin.v2.UpdateTableRequest.newBuilder()
+            .setTable(
+                Table.newBuilder()
+                    .setName(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+                    .setDeletionProtection(false))
+            .setUpdateMask(FieldMask.newBuilder().addPaths("deletion_protection").build())
+            .build();
+
+    assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
+  }
 }
