feat: add support for UpdateDatabase in Cloud Spanner (#2265)

* feature: Cloud Spanner Drop Database Protection

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Rajat Bhatta <93644539+rajatbhatta@users.noreply.github.com>

Author: 3 people

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Database.java

@@ -194,17 +194,56 @@ private String database() {
   static Database fromProto(
       com.google.spanner.admin.database.v1.Database proto, DatabaseAdminClient client) {
     checkArgument(!proto.getName().isEmpty(), "Missing expected 'name' field");
-    return new Database.Builder(client, DatabaseId.of(proto.getName()))
-        .setState(fromProtoState(proto.getState()))
-        .setCreateTime(Timestamp.fromProto(proto.getCreateTime()))
-        .setRestoreInfo(RestoreInfo.fromProtoOrNullIfDefaultInstance(proto.getRestoreInfo()))
-        .setVersionRetentionPeriod(proto.getVersionRetentionPeriod())
-        .setEarliestVersionTime(Timestamp.fromProto(proto.getEarliestVersionTime()))
-        .setEncryptionConfig(CustomerManagedEncryption.fromProtoOrNull(proto.getEncryptionConfig()))
-        .setDefaultLeader(proto.getDefaultLeader())
-        .setDialect(Dialect.fromProto(proto.getDatabaseDialect()))
-        .setProto(proto)
-        .build();
+    DatabaseInfo.Builder builder =
+        new Builder(client, DatabaseId.of(proto.getName()))
+            .setState(fromProtoState(proto.getState()))
+            .setCreateTime(Timestamp.fromProto(proto.getCreateTime()))
+            .setRestoreInfo(RestoreInfo.fromProtoOrNullIfDefaultInstance(proto.getRestoreInfo()))
+            .setVersionRetentionPeriod(proto.getVersionRetentionPeriod())
+            .setEarliestVersionTime(Timestamp.fromProto(proto.getEarliestVersionTime()))
+            .setEncryptionConfig(
+                CustomerManagedEncryption.fromProtoOrNull(proto.getEncryptionConfig()))
+            .setDefaultLeader(proto.getDefaultLeader())
+            .setDialect(Dialect.fromProto(proto.getDatabaseDialect()))
+            .setReconciling(proto.getReconciling())
+            .setProto(proto);
+    if (proto.getEnableDropProtection()) {
+      builder.enableDropProtection();
+    } else {
+      builder.disableDropProtection();
+    }
+    return builder.build();
+  }
+
+  public com.google.spanner.admin.database.v1.Database toProto() {
+    com.google.spanner.admin.database.v1.Database.Builder builder =
+        com.google.spanner.admin.database.v1.Database.newBuilder()
+            .setName(getId().getName())
+            .setState(getState().toProto())
+            .setEnableDropProtection(isDropProtectionEnabled())
+            .setReconciling(getReconciling());
+    if (getCreateTime() != null) {
+      builder.setCreateTime(getCreateTime().toProto());
+    }
+    if (getRestoreInfo() != null) {
+      builder.setRestoreInfo(getRestoreInfo().getProto());
+    }
+    if (getVersionRetentionPeriod() != null) {
+      builder.setVersionRetentionPeriod(getVersionRetentionPeriod());
+    }
+    if (getEarliestVersionTime() != null) {
+      builder.setEarliestVersionTime(getEarliestVersionTime().toProto());
+    }
+    if (getEncryptionConfig() != null) {
+      builder.setEncryptionConfig(getEncryptionConfig().toProto());
+    }
+    if (getDefaultLeader() != null) {
+      builder.setDefaultLeader(getDefaultLeader());
+    }
+    if (getDialect() != null) {
+      builder.setDatabaseDialect(getDialect().toProto());
+    }
+    return builder.build();
   }
 
   static DatabaseInfo.State fromProtoState(

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseAdminClient.java

@@ -28,6 +28,7 @@
 import com.google.spanner.admin.database.v1.CreateDatabaseRequest;
 import com.google.spanner.admin.database.v1.RestoreDatabaseMetadata;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
+import com.google.spanner.admin.database.v1.UpdateDatabaseMetadata;
 import java.util.List;
 import javax.annotation.Nullable;
 
@@ -358,6 +359,47 @@ OperationFuture<Database, RestoreDatabaseMetadata> restoreDatabase(Restore resto
    */
   Database getDatabase(String instanceId, String databaseId) throws SpannerException;
 
+  /**
+   * Updates a Cloud Spanner database. The returned {@code Operation} can be used to track the
+   * progress of the update. Throws SpannerException if the Cloud Spanner database does not exist.
+   *
+   * <p>Until completion of the returned operation:
+   *
+   * <ul>
+   *   <li>Cancelling the operation is best effort and may or may not succeed.
+   *   <li>All other attempts to modify the database are rejected.
+   *   <li>Reading the database via the API continues to give the pre-request field values.
+   * </ul>
+   *
+   * Upon completion of the returned operation:
+   *
+   * <ul>
+   *   <li>The database's new fields are readable via the API.
+   * </ul>
+   *
+   * <p>Example of updating a database.
+   *
+   * <pre>{@code
+   * String projectId = my_project_id;
+   * String instanceId = my_instance_id;
+   * String databaseId = my_database_id;
+   * Database databaseToUpdate = databaseAdminClient.newDatabaseBuilder(
+   *         DatabaseId.of(projectId, instanceId, databaseId))
+   *      .enableDropProtection().build();
+   * OperationFuture<Database, UpdateDatabaseMetadata> op = databaseAdminClient.updateDatabase(
+   *           databaseToUpdate, DatabaseField.DROP_PROTECTION);
+   * Database updateDatabase = op.get(5, TimeUnit.MINUTES);
+   * }</pre>
+   *
+   * @param database The database to update to. The current field values of the database will be
+   *     updated to the values specified in this parameter.
+   * @param fieldsToUpdate The fields that should be updated. Only these fields will have their
+   *     values updated to the values specified in {@param database}, even if there are other fields
+   *     specified in {@param database}.
+   */
+  OperationFuture<Database, UpdateDatabaseMetadata> updateDatabase(
+      Database database, DatabaseInfo.DatabaseField... fieldsToUpdate) throws SpannerException;
+
   /**
    * Gets the current state of a Cloud Spanner database backup.
    *

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseAdminClientImpl.java

@@ -23,6 +23,7 @@
 import com.google.cloud.Policy;
 import com.google.cloud.Policy.DefaultMarshaller;
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.DatabaseInfo.DatabaseField;
 import com.google.cloud.spanner.Options.ListOption;
 import com.google.cloud.spanner.SpannerImpl.PageFetcher;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
@@ -415,6 +416,27 @@ public Database getDatabase(String instanceId, String databaseId) throws Spanner
     return Database.fromProto(rpc.getDatabase(dbName), DatabaseAdminClientImpl.this);
   }
 
+  @Override
+  public OperationFuture<Database, UpdateDatabaseMetadata> updateDatabase(
+      Database database, DatabaseField... fieldsToUpdate) throws SpannerException {
+    FieldMask fieldMask = DatabaseInfo.DatabaseField.toFieldMask(fieldsToUpdate);
+    OperationFuture<com.google.spanner.admin.database.v1.Database, UpdateDatabaseMetadata>
+        rawOperationFuture = rpc.updateDatabase(database.toProto(), fieldMask);
+    return new OperationFutureImpl<>(
+        rawOperationFuture.getPollingFuture(),
+        rawOperationFuture.getInitialFuture(),
+        snapshot ->
+            Database.fromProto(
+                ProtoOperationTransformers.ResponseTransformer.create(
+                        com.google.spanner.admin.database.v1.Database.class)
+                    .apply(snapshot),
+                DatabaseAdminClientImpl.this),
+        ProtoOperationTransformers.MetadataTransformer.create(UpdateDatabaseMetadata.class),
+        e -> {
+          throw SpannerExceptionFactory.newSpannerException(e);
+        });
+  }
+
   @Override
   public OperationFuture<Void, UpdateDatabaseDdlMetadata> updateDatabaseDdl(
       final String instanceId,

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseInfo.java

@@ -16,15 +16,42 @@
 
 package com.google.cloud.spanner;
 
+import com.google.cloud.FieldSelector;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.encryption.CustomerManagedEncryption;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.FieldMask;
+import com.google.spanner.admin.database.v1.Database.State;
 import java.util.Objects;
 import javax.annotation.Nullable;
 
 /** Represents a Cloud Spanner database. */
 public class DatabaseInfo {
 
+  /** Represent an updatable field in a Cloud Spanner database. */
+  public enum DatabaseField implements FieldSelector {
+    DROP_PROTECTION("enable_drop_protection");
+
+    private final String selector;
+
+    DatabaseField(String selector) {
+      this.selector = selector;
+    }
+
+    @Override
+    public String getSelector() {
+      return selector;
+    }
+
+    static FieldMask toFieldMask(DatabaseInfo.DatabaseField... fields) {
+      FieldMask.Builder builder = FieldMask.newBuilder();
+      for (DatabaseInfo.DatabaseField field : fields) {
+        builder.addPaths(field.getSelector());
+      }
+      return builder.build();
+    }
+  }
+
   public abstract static class Builder {
     abstract Builder setState(State state);
 
@@ -58,6 +85,18 @@ public Builder setDialect(Dialect dialect) {
       throw new UnsupportedOperationException("Unimplemented");
     }
 
+    public Builder enableDropProtection() {
+      throw new UnsupportedOperationException("Unimplemented");
+    }
+
+    public Builder disableDropProtection() {
+      throw new UnsupportedOperationException("Unimplemented");
+    }
+
+    protected Builder setReconciling(boolean reconciling) {
+      throw new UnsupportedOperationException("Unimplemented");
+    }
+
     abstract Builder setProto(com.google.spanner.admin.database.v1.Database proto);
 
     /** Builds the database from this builder. */
@@ -74,6 +113,8 @@ abstract static class BuilderImpl extends Builder {
     private CustomerManagedEncryption encryptionConfig;
     private String defaultLeader;
     private Dialect dialect = Dialect.GOOGLE_STANDARD_SQL;
+    private boolean dropProtectionEnabled;
+    private boolean reconciling;
     private com.google.spanner.admin.database.v1.Database proto;
 
     BuilderImpl(DatabaseId id) {
@@ -141,6 +182,24 @@ public Builder setDialect(Dialect dialect) {
       return this;
     }
 
+    @Override
+    public Builder enableDropProtection() {
+      this.dropProtectionEnabled = true;
+      return this;
+    }
+
+    @Override
+    public Builder disableDropProtection() {
+      this.dropProtectionEnabled = false;
+      return this;
+    }
+
+    @Override
+    protected Builder setReconciling(boolean reconciling) {
+      this.reconciling = reconciling;
+      return this;
+    }
+
     @Override
     Builder setProto(@Nullable com.google.spanner.admin.database.v1.Database proto) {
       this.proto = proto;
@@ -151,13 +210,35 @@ Builder setProto(@Nullable com.google.spanner.admin.database.v1.Database proto)
   /** State of the database. */
   public enum State {
     // Not specified.
-    UNSPECIFIED,
+    UNSPECIFIED {
+      @Override
+      public com.google.spanner.admin.database.v1.Database.State toProto() {
+        return com.google.spanner.admin.database.v1.Database.State.STATE_UNSPECIFIED;
+      }
+    },
     // The database is still being created and is not ready to use.
-    CREATING,
+    CREATING {
+      @Override
+      public com.google.spanner.admin.database.v1.Database.State toProto() {
+        return com.google.spanner.admin.database.v1.Database.State.CREATING;
+      }
+    },
     // The database is fully created and ready to use.
-    READY,
+    READY {
+      @Override
+      public com.google.spanner.admin.database.v1.Database.State toProto() {
+        return com.google.spanner.admin.database.v1.Database.State.READY;
+      }
+    },
     // The database has restored and is being optimized for use.
-    READY_OPTIMIZING
+    READY_OPTIMIZING {
+      @Override
+      public com.google.spanner.admin.database.v1.Database.State toProto() {
+        return com.google.spanner.admin.database.v1.Database.State.READY_OPTIMIZING;
+      }
+    };
+
+    public abstract com.google.spanner.admin.database.v1.Database.State toProto();
   }
 
   private final DatabaseId id;
@@ -169,6 +250,8 @@ public enum State {
   private final CustomerManagedEncryption encryptionConfig;
   private final String defaultLeader;
   private final Dialect dialect;
+  private final boolean dropProtectionEnabled;
+  private final boolean reconciling;
   private final com.google.spanner.admin.database.v1.Database proto;
 
   public DatabaseInfo(DatabaseId id, State state) {
@@ -181,6 +264,8 @@ public DatabaseInfo(DatabaseId id, State state) {
     this.encryptionConfig = null;
     this.defaultLeader = null;
     this.dialect = null;
+    this.dropProtectionEnabled = false;
+    this.reconciling = false;
     this.proto = null;
   }
 
@@ -194,6 +279,8 @@ public DatabaseInfo(DatabaseId id, State state) {
     this.encryptionConfig = builder.encryptionConfig;
     this.defaultLeader = builder.defaultLeader;
     this.dialect = builder.dialect;
+    this.dropProtectionEnabled = builder.dropProtectionEnabled;
+    this.reconciling = builder.reconciling;
     this.proto = builder.proto;
   }
 
@@ -262,6 +349,14 @@ public Timestamp getEarliestVersionTime() {
     return dialect;
   }
 
+  public boolean isDropProtectionEnabled() {
+    return dropProtectionEnabled;
+  }
+
+  public boolean getReconciling() {
+    return reconciling;
+  }
+
   /** Returns the raw proto instance that was used to construct this {@link Database}. */
   public @Nullable com.google.spanner.admin.database.v1.Database getProto() {
     return proto;
@@ -284,7 +379,9 @@ public boolean equals(Object o) {
         && Objects.equals(earliestVersionTime, that.earliestVersionTime)
         && Objects.equals(encryptionConfig, that.encryptionConfig)
         && Objects.equals(defaultLeader, that.defaultLeader)
-        && Objects.equals(dialect, that.dialect);
+        && Objects.equals(dialect, that.dialect)
+        && Objects.equals(dropProtectionEnabled, that.dropProtectionEnabled)
+        && Objects.equals(reconciling, that.reconciling);
   }
 
   @Override
@@ -298,13 +395,15 @@ public int hashCode() {
         earliestVersionTime,
         encryptionConfig,
         defaultLeader,
-        dialect);
+        dialect,
+        dropProtectionEnabled,
+        reconciling);
   }
 
   @Override
   public String toString() {
     return String.format(
-        "Database[%s, %s, %s, %s, %s, %s, %s, %s, %s]",
+        "Database[%s, %s, %s, %s, %s, %s, %s, %s, %s %s %s]",
         id.getName(),
         state,
         createTime,
@@ -313,6 +412,8 @@ public String toString() {
         earliestVersionTime,
         encryptionConfig,
         defaultLeader,
-        dialect);
+        dialect,
+        dropProtectionEnabled,
+        reconciling);
   }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/encryption/CustomerManagedEncryption.java

@@ -42,6 +42,10 @@ public static CustomerManagedEncryption fromProtoOrNull(EncryptionConfig proto)
         : new CustomerManagedEncryption(proto.getKmsKeyName());
   }
 
+  public EncryptionConfig toProto() {
+    return EncryptionConfig.newBuilder().setKmsKeyName(this.getKmsKeyName()).build();
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {