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,8 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Paths;
+import java.util.Collections;
+import java.util.List;
 import javax.annotation.Nonnull;
 
 /**
@@ -70,9 +73,24 @@ 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)));
+    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 the avro schema for this schema bundle. */
+  public CreateSchemaBundleRequest setAvroSchema(@Nonnull List<String> avroSchema) {
+    Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
+    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;
 
 /**
@@ -42,7 +43,8 @@ private SchemaBundle(@Nonnull com.google.bigtable.admin.v2.SchemaBundle proto) {
     Preconditions.checkNotNull(proto);
     Preconditions.checkArgument(!proto.getName().isEmpty(), "SchemaBundle must have a name");
     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 +69,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,8 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Paths;
+import java.util.Collections;
+import java.util.List;
 import javax.annotation.Nonnull;
 
 /**
@@ -90,13 +93,29 @@ public UpdateSchemaBundleRequest setProtoSchemaFile(@Nonnull String protoSchemaF
   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)));
+    requestBuilder
+        .getSchemaBundleBuilder()
+        .setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema));
     updateFieldMask(com.google.bigtable.admin.v2.SchemaBundle.PROTO_SCHEMA_FIELD_NUMBER);
     return this;
   }
 
+  /** Sets 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));
+  }
+
+  /** Sets the avro schema for this schema bundle. */
+  public UpdateSchemaBundleRequest setAvroSchema(@Nonnull List<String> avroSchema) {
+    Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
+    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

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

@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.net.URISyntaxException;
@@ -39,6 +40,9 @@ public class CreateSchemaBundleRequestTest {
   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\"}";
 
   @Test
   public void testToProto() throws IOException, URISyntaxException {
@@ -99,6 +103,82 @@ public void testHashCode() throws IOException, URISyntaxException {
                 .hashCode());
   }
 
+  @Test
+  public void testToProtoWithAvroSchema() {
+    CreateSchemaBundleRequest request =
+        CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
+
+    com.google.bigtable.admin.v2.CreateSchemaBundleRequest requestProto =
+        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(
+                        com.google.bigtable.admin.v2.AvroSchema.newBuilder()
+                            .addJsonSchemas(TEST_AVRO_SCHEMA)
+                            .build())
+                    .build())
+            .build();
+    assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
+  }
+
+  @Test
+  public void testToProtoWithAvroSchemaList() {
+    CreateSchemaBundleRequest request =
+        CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
+            .setAvroSchema(ImmutableList.of(TEST_AVRO_SCHEMA, TEST_UPDATED_AVRO_SCHEMA));
+
+    com.google.bigtable.admin.v2.CreateSchemaBundleRequest requestProto =
+        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(
+                        com.google.bigtable.admin.v2.AvroSchema.newBuilder()
+                            .addJsonSchemas(TEST_AVRO_SCHEMA)
+                            .addJsonSchemas(TEST_UPDATED_AVRO_SCHEMA)
+                            .build())
+                    .build())
+            .build();
+    assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
+  }
+
+  @Test
+  public void testEqualityWithAvroSchema() {
+    CreateSchemaBundleRequest request =
+        CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
+
+    assertThat(request)
+        .isEqualTo(
+            CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
+                .setAvroSchema(TEST_AVRO_SCHEMA));
+
+    assertThat(request)
+        .isNotEqualTo(
+            CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
+                .setAvroSchema(TEST_UPDATED_AVRO_SCHEMA));
+  }
+
+  @Test
+  public void testHashCodeWithAvroSchema() {
+    CreateSchemaBundleRequest request =
+        CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
+
+    assertThat(request.hashCode())
+        .isEqualTo(
+            CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
+                .setAvroSchema(TEST_AVRO_SCHEMA)
+                .hashCode());
+
+    assertThat(request.hashCode())
+        .isNotEqualTo(
+            CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
+                .setAvroSchema(TEST_UPDATED_AVRO_SCHEMA)
+                .hashCode());
+  }
+
   private String getResourceFilePath(String filePath) throws URISyntaxException {
     ClassLoader cl = Thread.currentThread().getContextClassLoader();
     URL protoSchema = cl.getResource(filePath);