feat(bigtable): Support Avro schema bundles in Bigtable's Admin API

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/CreateSchemaBundleRequest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.bigtable.admin.v2.models;
 
 import com.google.api.core.InternalApi;
+import com.google.bigtable.admin.v2.AvroSchema;
 import com.google.bigtable.admin.v2.ProtoSchema;
 import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
 import com.google.common.base.Objects;
@@ -25,6 +26,9 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Paths;
+import java.util.Collections;
+import java.util.List;
+import java.util.logging.Logger;
 import javax.annotation.Nonnull;
 
 /**
@@ -42,6 +46,8 @@
  * @see SchemaBundle for more details.
  */
 public final class CreateSchemaBundleRequest {
+  private static final Logger LOGGER = Logger.getLogger(CreateSchemaBundleRequest.class.getName());
+
   private final String tableId;
   private final com.google.bigtable.admin.v2.CreateSchemaBundleRequest.Builder requestBuilder =
       com.google.bigtable.admin.v2.CreateSchemaBundleRequest.newBuilder();
@@ -70,9 +76,34 @@ public CreateSchemaBundleRequest setProtoSchemaFile(@Nonnull String protoSchemaF
   /** Sets the proto schema for this schema bundle. */
   public CreateSchemaBundleRequest setProtoSchema(@Nonnull ByteString protoSchema) {
     Preconditions.checkNotNull(protoSchema, "protoSchema must be set");
-    requestBuilder.setSchemaBundle(
-        com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
-            .setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema)));
+    if (requestBuilder.getSchemaBundleBuilder().hasAvroSchema()) {
+      LOGGER.warning(
+          "This schema bundle already has an Avro schema set. Setting the proto schema will"
+              + " unset the Avro schema.");
+    }
+    requestBuilder
+        .getSchemaBundleBuilder()
+        .setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema));
+    return this;
+  }
+
+  /** Sets the avro schema for this schema bundle. */
+  public CreateSchemaBundleRequest setAvroSchema(@Nonnull String avroSchema) {
+    Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
+    return setAvroSchema(Collections.singletonList(avroSchema));
+  }
+
+  /** Sets a list of avro schemas for this schema bundle. */
+  public CreateSchemaBundleRequest setAvroSchema(@Nonnull List<String> avroSchema) {
+    Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
+    if (requestBuilder.getSchemaBundleBuilder().hasProtoSchema()) {
+      LOGGER.warning(
+          "This schema bundle already has a proto schema set. Setting the Avro schema will"
+              + " unset the proto schema.");
+    }
+    requestBuilder
+        .getSchemaBundleBuilder()
+        .setAvroSchema(AvroSchema.newBuilder().addAllJsonSchemas(avroSchema));
     return this;
   }
 

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/SchemaBundle.java

@@ -20,6 +20,7 @@
 import com.google.bigtable.admin.v2.SchemaBundleName;
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
+import java.util.List;
 import javax.annotation.Nonnull;
 
 /**
@@ -41,8 +42,11 @@ public static SchemaBundle fromProto(@Nonnull com.google.bigtable.admin.v2.Schem
   private SchemaBundle(@Nonnull com.google.bigtable.admin.v2.SchemaBundle proto) {
     Preconditions.checkNotNull(proto);
     Preconditions.checkArgument(!proto.getName().isEmpty(), "SchemaBundle must have a name");
+    // proto_schema and avro_schema are defined in a protobuf oneof, so at most one
+    // can be set.
     Preconditions.checkArgument(
-        proto.hasProtoSchema(), "Schemabundle must have a proto_schema field");
+        proto.hasProtoSchema() || proto.hasAvroSchema(),
+        "Schemabundle must have a proto_schema or avro_schema field");
     this.proto = proto;
     this.schemaBundleName = SchemaBundleName.parse(proto.getName());
   }
@@ -67,6 +71,14 @@ public com.google.protobuf.ByteString getProtoSchema() {
     throw new IllegalStateException("This SchemaBundle doesn't have a valid type specified");
   }
 
+  /** Gets the avro schema of this schema bundle. */
+  public List<String> getAvroSchema() {
+    if (proto.hasAvroSchema()) {
+      return proto.getAvroSchema().getJsonSchemasList();
+    }
+    throw new IllegalStateException("This SchemaBundle does not contain an Avro schema");
+  }
+
   /**
    * Creates the request protobuf. This method is considered an internal implementation detail and
    * not meant to be used by applications.

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/UpdateSchemaBundleRequest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.bigtable.admin.v2.models;
 
 import com.google.api.core.InternalApi;
+import com.google.bigtable.admin.v2.AvroSchema;
 import com.google.bigtable.admin.v2.ProtoSchema;
 import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
 import com.google.common.base.Objects;
@@ -27,6 +28,9 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Paths;
+import java.util.Collections;
+import java.util.List;
+import java.util.logging.Logger;
 import javax.annotation.Nonnull;
 
 /**
@@ -43,6 +47,8 @@
  * @see SchemaBundle for more details.
  */
 public final class UpdateSchemaBundleRequest {
+  private static final Logger LOGGER = Logger.getLogger(UpdateSchemaBundleRequest.class.getName());
+
   private final com.google.bigtable.admin.v2.UpdateSchemaBundleRequest.Builder requestBuilder;
   private final String tableId;
   private final String schemaBundleId;
@@ -78,25 +84,51 @@ private UpdateSchemaBundleRequest(
     this.requestBuilder = requestBuilder;
   }
 
-  /** Sets the proto schema for this schema bundle. */
+  /** Updates the proto schema for this schema bundle. */
   public UpdateSchemaBundleRequest setProtoSchemaFile(@Nonnull String protoSchemaFile)
       throws IOException {
     Preconditions.checkNotNull(protoSchemaFile, "protoSchemaFile must be set");
     byte[] content = Files.readAllBytes(Paths.get(protoSchemaFile));
     return setProtoSchema(ByteString.copyFrom(content));
   }
 
-  /** Sets the proto schema for this schema bundle. */
+  /** Updates the proto schema for this schema bundle. */
   public UpdateSchemaBundleRequest setProtoSchema(@Nonnull ByteString protoSchema)
       throws IOException {
     Preconditions.checkNotNull(protoSchema, "protoSchema must be set");
-    requestBuilder.setSchemaBundle(
-        com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
-            .setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema)));
+    if (requestBuilder.getSchemaBundleBuilder().hasAvroSchema()) {
+      LOGGER.warning(
+          "This schema bundle already has an Avro schema set. Setting the proto schema will"
+              + " unset the Avro schema.");
+    }
+    requestBuilder
+        .getSchemaBundleBuilder()
+        .setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema));
     updateFieldMask(com.google.bigtable.admin.v2.SchemaBundle.PROTO_SCHEMA_FIELD_NUMBER);
     return this;
   }
 
+  /** Updates the avro schema for this schema bundle. */
+  public UpdateSchemaBundleRequest setAvroSchema(@Nonnull String avroSchema) {
+    Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
+    return setAvroSchema(Collections.singletonList(avroSchema));
+  }
+
+  /** Updates the list of avro schemas for this schema bundle. */
+  public UpdateSchemaBundleRequest setAvroSchema(@Nonnull List<String> avroSchema) {
+    Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
+    if (requestBuilder.getSchemaBundleBuilder().hasProtoSchema()) {
+      LOGGER.warning(
+          "This schema bundle already has a proto schema set. Setting the Avro schema will"
+              + " unset the proto schema.");
+    }
+    requestBuilder
+        .getSchemaBundleBuilder()
+        .setAvroSchema(AvroSchema.newBuilder().addAllJsonSchemas(avroSchema));
+    updateFieldMask(com.google.bigtable.admin.v2.SchemaBundle.AVRO_SCHEMA_FIELD_NUMBER);
+    return this;
+  }
+
   /**
    * Configures if safety warnings should be disabled. If set, then non backwards compatible changes
    * are allowed.

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientTests.java

@@ -28,6 +28,7 @@
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.api.gax.rpc.testing.FakeOperationSnapshot;
 import com.google.bigtable.admin.v2.AuthorizedViewName;
+import com.google.bigtable.admin.v2.AvroSchema;
 import com.google.bigtable.admin.v2.Backup.State;
 import com.google.bigtable.admin.v2.BackupInfo;
 import com.google.bigtable.admin.v2.ChangeStreamConfig;
@@ -142,6 +143,9 @@ public class BigtableTableAdminClientTests {
   private static final String TEST_PROTO_SCHEMA_BUNDLE = "proto_schema_bundle.pb";
   // Location: `google-cloud-bigtable/src/test/resources/updated_proto_schema_bundle.pb`
   private static final String TEST_UPDATED_PROTO_SCHEMA_BUNDLE = "updated_proto_schema_bundle.pb";
+  private static final String TEST_AVRO_SCHEMA = "{\"type\": \"record\", \"name\": \"User\"}";
+  private static final String TEST_UPDATED_AVRO_SCHEMA =
+      "{\"type\": \"record\", \"name\": \"UpdatedUser\"}";
 
   private static final String INSTANCE_NAME = NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID);
   private static final String TABLE_NAME =
@@ -1594,6 +1598,125 @@ public void testDeleteSchemaBundle() {
     assertThat(wasCalled.get()).isTrue();
   }
 
+  @Test
+  public void testCreateSchemaBundleWithAvroSchema() {
+    // Setup
+    Mockito.when(mockStub.createSchemaBundleOperationCallable())
+        .thenReturn(mockCreateSchemaBundleOperationCallable);
+
+    com.google.bigtable.admin.v2.CreateSchemaBundleRequest expectedRequest =
+        com.google.bigtable.admin.v2.CreateSchemaBundleRequest.newBuilder()
+            .setParent(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
+            .setSchemaBundleId(SCHEMA_BUNDLE_ID)
+            .setSchemaBundle(
+                com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
+                    .setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_AVRO_SCHEMA)))
+            .build();
+
+    com.google.bigtable.admin.v2.SchemaBundle expectedResponse =
+        com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
+            .setName(
+                NameUtil.formatSchemaBundleName(
+                    PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
+            .setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_AVRO_SCHEMA))
+            .build();
+
+    mockOperationResult(
+        mockCreateSchemaBundleOperationCallable,
+        expectedRequest,
+        expectedResponse,
+        CreateSchemaBundleMetadata.newBuilder()
+            .setName(expectedRequest.getSchemaBundle().getName())
+            .build());
+
+    CreateSchemaBundleRequest req =
+        CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
+
+    // Execute
+    SchemaBundle actualResult = adminClient.createSchemaBundle(req);
+
+    // Verify
+    assertThat(actualResult).isEqualTo(SchemaBundle.fromProto(expectedResponse));
+    assertThat(actualResult.getAvroSchema()).containsExactly(TEST_AVRO_SCHEMA);
+  }
+
+  @Test
+  public void testUpdateSchemaBundleWithAvroSchema() {
+    // Setup
+    Mockito.when(mockStub.updateSchemaBundleOperationCallable())
+        .thenReturn(mockUpdateSchemaBundleOperationCallable);
+
+    com.google.bigtable.admin.v2.UpdateSchemaBundleRequest expectedRequest =
+        com.google.bigtable.admin.v2.UpdateSchemaBundleRequest.newBuilder()
+            .setSchemaBundle(
+                com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
+                    .setName(
+                        NameUtil.formatSchemaBundleName(
+                            PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
+                    .setAvroSchema(
+                        AvroSchema.newBuilder().addJsonSchemas(TEST_UPDATED_AVRO_SCHEMA)))
+            .setUpdateMask(FieldMask.newBuilder().addPaths("avro_schema"))
+            .build();
+
+    com.google.bigtable.admin.v2.SchemaBundle expectedResponse =
+        com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
+            .setName(
+                NameUtil.formatSchemaBundleName(
+                    PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
+            .setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_UPDATED_AVRO_SCHEMA))
+            .build();
+
+    mockOperationResult(
+        mockUpdateSchemaBundleOperationCallable,
+        expectedRequest,
+        expectedResponse,
+        UpdateSchemaBundleMetadata.newBuilder()
+            .setName(expectedRequest.getSchemaBundle().getName())
+            .build());
+
+    UpdateSchemaBundleRequest req =
+        UpdateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
+            .setAvroSchema(TEST_UPDATED_AVRO_SCHEMA);
+
+    // Execute
+    SchemaBundle actualResult = adminClient.updateSchemaBundle(req);
+
+    // Verify
+    assertThat(actualResult).isEqualTo(SchemaBundle.fromProto(expectedResponse));
+    assertThat(actualResult.getAvroSchema()).containsExactly(TEST_UPDATED_AVRO_SCHEMA);
+  }
+
+  @Test
+  public void testGetSchemaBundleWithAvroSchema() {
+    // Setup
+    Mockito.when(mockStub.getSchemaBundleCallable()).thenReturn(mockGetSchemaBundleCallable);
+
+    com.google.bigtable.admin.v2.GetSchemaBundleRequest expectedRequest =
+        com.google.bigtable.admin.v2.GetSchemaBundleRequest.newBuilder()
+            .setName(
+                NameUtil.formatSchemaBundleName(
+                    PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
+            .build();
+
+    com.google.bigtable.admin.v2.SchemaBundle expectedResponse =
+        com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
+            .setName(
+                NameUtil.formatSchemaBundleName(
+                    PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
+            .setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_AVRO_SCHEMA))
+            .build();
+
+    Mockito.when(mockGetSchemaBundleCallable.futureCall(expectedRequest))
+        .thenReturn(ApiFutures.immediateFuture(expectedResponse));
+
+    // Execute
+    SchemaBundle actualResult = adminClient.getSchemaBundle(TABLE_ID, SCHEMA_BUNDLE_ID);
+
+    // Verify
+    assertThat(actualResult).isEqualTo(SchemaBundle.fromProto(expectedResponse));
+    assertThat(actualResult.getAvroSchema()).containsExactly(TEST_AVRO_SCHEMA);
+  }
+
   @Test
   public void testGetBackupIamPolicy() {
     // Setup