Skip to content

Commit 348c064

Browse files
trollyxiatrollyxia
committed
feat: add support for Avro schema to SchemaBundle model classes
1 parent 3da3562 commit 348c064

7 files changed

Lines changed: 476 additions & 7 deletions

File tree

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

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.cloud.bigtable.admin.v2.models;
1818

1919
import com.google.api.core.InternalApi;
20+
import com.google.bigtable.admin.v2.AvroSchema;
2021
import com.google.bigtable.admin.v2.ProtoSchema;
2122
import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
2223
import com.google.common.base.Objects;
@@ -25,6 +26,8 @@
2526
import java.io.IOException;
2627
import java.nio.file.Files;
2728
import java.nio.file.Paths;
29+
import java.util.Collections;
30+
import java.util.List;
2831
import javax.annotation.Nonnull;
2932

3033
/**
@@ -70,9 +73,24 @@ public CreateSchemaBundleRequest setProtoSchemaFile(@Nonnull String protoSchemaF
7073
/** Sets the proto schema for this schema bundle. */
7174
public CreateSchemaBundleRequest setProtoSchema(@Nonnull ByteString protoSchema) {
7275
Preconditions.checkNotNull(protoSchema, "protoSchema must be set");
73-
requestBuilder.setSchemaBundle(
74-
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
75-
.setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema)));
76+
requestBuilder
77+
.getSchemaBundleBuilder()
78+
.setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema));
79+
return this;
80+
}
81+
82+
/** Sets the avro schema for this schema bundle. */
83+
public CreateSchemaBundleRequest setAvroSchema(@Nonnull String avroSchema) {
84+
Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
85+
return setAvroSchema(Collections.singletonList(avroSchema));
86+
}
87+
88+
/** Sets the avro schema for this schema bundle. */
89+
public CreateSchemaBundleRequest setAvroSchema(@Nonnull List<String> avroSchema) {
90+
Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
91+
requestBuilder
92+
.getSchemaBundleBuilder()
93+
.setAvroSchema(AvroSchema.newBuilder().addAllJsonSchemas(avroSchema));
7694
return this;
7795
}
7896

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.bigtable.admin.v2.SchemaBundleName;
2121
import com.google.common.base.Objects;
2222
import com.google.common.base.Preconditions;
23+
import java.util.List;
2324
import javax.annotation.Nonnull;
2425

2526
/**
@@ -42,7 +43,8 @@ private SchemaBundle(@Nonnull com.google.bigtable.admin.v2.SchemaBundle proto) {
4243
Preconditions.checkNotNull(proto);
4344
Preconditions.checkArgument(!proto.getName().isEmpty(), "SchemaBundle must have a name");
4445
Preconditions.checkArgument(
45-
proto.hasProtoSchema(), "Schemabundle must have a proto_schema field");
46+
proto.hasProtoSchema() || proto.hasAvroSchema(),
47+
"Schemabundle must have a proto_schema or avro_schema field");
4648
this.proto = proto;
4749
this.schemaBundleName = SchemaBundleName.parse(proto.getName());
4850
}
@@ -67,6 +69,14 @@ public com.google.protobuf.ByteString getProtoSchema() {
6769
throw new IllegalStateException("This SchemaBundle doesn't have a valid type specified");
6870
}
6971

72+
/** Gets the avro schema of this schema bundle. */
73+
public List<String> getAvroSchema() {
74+
if (proto.hasAvroSchema()) {
75+
return proto.getAvroSchema().getJsonSchemasList();
76+
}
77+
throw new IllegalStateException("This SchemaBundle doesn't have a valid type specified");
78+
}
79+
7080
/**
7181
* Creates the request protobuf. This method is considered an internal implementation detail and
7282
* not meant to be used by applications.

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

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.cloud.bigtable.admin.v2.models;
1818

1919
import com.google.api.core.InternalApi;
20+
import com.google.bigtable.admin.v2.AvroSchema;
2021
import com.google.bigtable.admin.v2.ProtoSchema;
2122
import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
2223
import com.google.common.base.Objects;
@@ -27,6 +28,8 @@
2728
import java.io.IOException;
2829
import java.nio.file.Files;
2930
import java.nio.file.Paths;
31+
import java.util.Collections;
32+
import java.util.List;
3033
import javax.annotation.Nonnull;
3134

3235
/**
@@ -90,13 +93,29 @@ public UpdateSchemaBundleRequest setProtoSchemaFile(@Nonnull String protoSchemaF
9093
public UpdateSchemaBundleRequest setProtoSchema(@Nonnull ByteString protoSchema)
9194
throws IOException {
9295
Preconditions.checkNotNull(protoSchema, "protoSchema must be set");
93-
requestBuilder.setSchemaBundle(
94-
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
95-
.setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema)));
96+
requestBuilder
97+
.getSchemaBundleBuilder()
98+
.setProtoSchema(ProtoSchema.newBuilder().setProtoDescriptors(protoSchema));
9699
updateFieldMask(com.google.bigtable.admin.v2.SchemaBundle.PROTO_SCHEMA_FIELD_NUMBER);
97100
return this;
98101
}
99102

103+
/** Sets the avro schema for this schema bundle. */
104+
public UpdateSchemaBundleRequest setAvroSchema(@Nonnull String avroSchema) {
105+
Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
106+
return setAvroSchema(Collections.singletonList(avroSchema));
107+
}
108+
109+
/** Sets the avro schema for this schema bundle. */
110+
public UpdateSchemaBundleRequest setAvroSchema(@Nonnull List<String> avroSchema) {
111+
Preconditions.checkNotNull(avroSchema, "avroSchema must be set");
112+
requestBuilder
113+
.getSchemaBundleBuilder()
114+
.setAvroSchema(AvroSchema.newBuilder().addAllJsonSchemas(avroSchema));
115+
updateFieldMask(com.google.bigtable.admin.v2.SchemaBundle.AVRO_SCHEMA_FIELD_NUMBER);
116+
return this;
117+
}
118+
100119
/**
101120
* Configures if safety warnings should be disabled. If set, then non backwards compatible changes
102121
* are allowed.

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.api.gax.rpc.UnaryCallable;
2929
import com.google.api.gax.rpc.testing.FakeOperationSnapshot;
3030
import com.google.bigtable.admin.v2.AuthorizedViewName;
31+
import com.google.bigtable.admin.v2.AvroSchema;
3132
import com.google.bigtable.admin.v2.Backup.State;
3233
import com.google.bigtable.admin.v2.BackupInfo;
3334
import com.google.bigtable.admin.v2.ChangeStreamConfig;
@@ -142,6 +143,9 @@ public class BigtableTableAdminClientTests {
142143
private static final String TEST_PROTO_SCHEMA_BUNDLE = "proto_schema_bundle.pb";
143144
// Location: `google-cloud-bigtable/src/test/resources/updated_proto_schema_bundle.pb`
144145
private static final String TEST_UPDATED_PROTO_SCHEMA_BUNDLE = "updated_proto_schema_bundle.pb";
146+
private static final String TEST_AVRO_SCHEMA = "{\"type\": \"record\", \"name\": \"User\"}";
147+
private static final String TEST_UPDATED_AVRO_SCHEMA =
148+
"{\"type\": \"record\", \"name\": \"UpdatedUser\"}";
145149

146150
private static final String INSTANCE_NAME = NameUtil.formatInstanceName(PROJECT_ID, INSTANCE_ID);
147151
private static final String TABLE_NAME =
@@ -1594,6 +1598,125 @@ public void testDeleteSchemaBundle() {
15941598
assertThat(wasCalled.get()).isTrue();
15951599
}
15961600

1601+
@Test
1602+
public void testCreateSchemaBundleWithAvroSchema() {
1603+
// Setup
1604+
Mockito.when(mockStub.createSchemaBundleOperationCallable())
1605+
.thenReturn(mockCreateSchemaBundleOperationCallable);
1606+
1607+
com.google.bigtable.admin.v2.CreateSchemaBundleRequest expectedRequest =
1608+
com.google.bigtable.admin.v2.CreateSchemaBundleRequest.newBuilder()
1609+
.setParent(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
1610+
.setSchemaBundleId(SCHEMA_BUNDLE_ID)
1611+
.setSchemaBundle(
1612+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
1613+
.setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_AVRO_SCHEMA)))
1614+
.build();
1615+
1616+
com.google.bigtable.admin.v2.SchemaBundle expectedResponse =
1617+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
1618+
.setName(
1619+
NameUtil.formatSchemaBundleName(
1620+
PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
1621+
.setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_AVRO_SCHEMA))
1622+
.build();
1623+
1624+
mockOperationResult(
1625+
mockCreateSchemaBundleOperationCallable,
1626+
expectedRequest,
1627+
expectedResponse,
1628+
CreateSchemaBundleMetadata.newBuilder()
1629+
.setName(expectedRequest.getSchemaBundle().getName())
1630+
.build());
1631+
1632+
CreateSchemaBundleRequest req =
1633+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
1634+
1635+
// Execute
1636+
SchemaBundle actualResult = adminClient.createSchemaBundle(req);
1637+
1638+
// Verify
1639+
assertThat(actualResult).isEqualTo(SchemaBundle.fromProto(expectedResponse));
1640+
assertThat(actualResult.getAvroSchema()).containsExactly(TEST_AVRO_SCHEMA);
1641+
}
1642+
1643+
@Test
1644+
public void testUpdateSchemaBundleWithAvroSchema() {
1645+
// Setup
1646+
Mockito.when(mockStub.updateSchemaBundleOperationCallable())
1647+
.thenReturn(mockUpdateSchemaBundleOperationCallable);
1648+
1649+
com.google.bigtable.admin.v2.UpdateSchemaBundleRequest expectedRequest =
1650+
com.google.bigtable.admin.v2.UpdateSchemaBundleRequest.newBuilder()
1651+
.setSchemaBundle(
1652+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
1653+
.setName(
1654+
NameUtil.formatSchemaBundleName(
1655+
PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
1656+
.setAvroSchema(
1657+
AvroSchema.newBuilder().addJsonSchemas(TEST_UPDATED_AVRO_SCHEMA)))
1658+
.setUpdateMask(FieldMask.newBuilder().addPaths("avro_schema"))
1659+
.build();
1660+
1661+
com.google.bigtable.admin.v2.SchemaBundle expectedResponse =
1662+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
1663+
.setName(
1664+
NameUtil.formatSchemaBundleName(
1665+
PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
1666+
.setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_UPDATED_AVRO_SCHEMA))
1667+
.build();
1668+
1669+
mockOperationResult(
1670+
mockUpdateSchemaBundleOperationCallable,
1671+
expectedRequest,
1672+
expectedResponse,
1673+
UpdateSchemaBundleMetadata.newBuilder()
1674+
.setName(expectedRequest.getSchemaBundle().getName())
1675+
.build());
1676+
1677+
UpdateSchemaBundleRequest req =
1678+
UpdateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
1679+
.setAvroSchema(TEST_UPDATED_AVRO_SCHEMA);
1680+
1681+
// Execute
1682+
SchemaBundle actualResult = adminClient.updateSchemaBundle(req);
1683+
1684+
// Verify
1685+
assertThat(actualResult).isEqualTo(SchemaBundle.fromProto(expectedResponse));
1686+
assertThat(actualResult.getAvroSchema()).containsExactly(TEST_UPDATED_AVRO_SCHEMA);
1687+
}
1688+
1689+
@Test
1690+
public void testGetSchemaBundleWithAvroSchema() {
1691+
// Setup
1692+
Mockito.when(mockStub.getSchemaBundleCallable()).thenReturn(mockGetSchemaBundleCallable);
1693+
1694+
com.google.bigtable.admin.v2.GetSchemaBundleRequest expectedRequest =
1695+
com.google.bigtable.admin.v2.GetSchemaBundleRequest.newBuilder()
1696+
.setName(
1697+
NameUtil.formatSchemaBundleName(
1698+
PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
1699+
.build();
1700+
1701+
com.google.bigtable.admin.v2.SchemaBundle expectedResponse =
1702+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
1703+
.setName(
1704+
NameUtil.formatSchemaBundleName(
1705+
PROJECT_ID, INSTANCE_ID, TABLE_ID, SCHEMA_BUNDLE_ID))
1706+
.setAvroSchema(AvroSchema.newBuilder().addJsonSchemas(TEST_AVRO_SCHEMA))
1707+
.build();
1708+
1709+
Mockito.when(mockGetSchemaBundleCallable.futureCall(expectedRequest))
1710+
.thenReturn(ApiFutures.immediateFuture(expectedResponse));
1711+
1712+
// Execute
1713+
SchemaBundle actualResult = adminClient.getSchemaBundle(TABLE_ID, SCHEMA_BUNDLE_ID);
1714+
1715+
// Verify
1716+
assertThat(actualResult).isEqualTo(SchemaBundle.fromProto(expectedResponse));
1717+
assertThat(actualResult.getAvroSchema()).containsExactly(TEST_AVRO_SCHEMA);
1718+
}
1719+
15971720
@Test
15981721
public void testGetBackupIamPolicy() {
15991722
// Setup

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import static com.google.common.truth.Truth.assertThat;
2020

2121
import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
22+
import com.google.common.collect.ImmutableList;
2223
import com.google.protobuf.ByteString;
2324
import java.io.IOException;
2425
import java.net.URISyntaxException;
@@ -39,6 +40,9 @@ public class CreateSchemaBundleRequestTest {
3940
private static final String TEST_PROTO_SCHEMA_BUNDLE = "proto_schema_bundle.pb";
4041
// Location: `google-cloud-bigtable/src/test/resources/updated_proto_schema_bundle.pb`
4142
private static final String TEST_UPDATED_PROTO_SCHEMA_BUNDLE = "updated_proto_schema_bundle.pb";
43+
private static final String TEST_AVRO_SCHEMA = "{\"type\": \"record\", \"name\": \"User\"}";
44+
private static final String TEST_UPDATED_AVRO_SCHEMA =
45+
"{\"type\": \"record\", \"name\": \"UpdatedUser\"}";
4246

4347
@Test
4448
public void testToProto() throws IOException, URISyntaxException {
@@ -99,6 +103,82 @@ public void testHashCode() throws IOException, URISyntaxException {
99103
.hashCode());
100104
}
101105

106+
@Test
107+
public void testToProtoWithAvroSchema() {
108+
CreateSchemaBundleRequest request =
109+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
110+
111+
com.google.bigtable.admin.v2.CreateSchemaBundleRequest requestProto =
112+
com.google.bigtable.admin.v2.CreateSchemaBundleRequest.newBuilder()
113+
.setParent(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
114+
.setSchemaBundleId(SCHEMA_BUNDLE_ID)
115+
.setSchemaBundle(
116+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
117+
.setAvroSchema(
118+
com.google.bigtable.admin.v2.AvroSchema.newBuilder()
119+
.addJsonSchemas(TEST_AVRO_SCHEMA)
120+
.build())
121+
.build())
122+
.build();
123+
assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
124+
}
125+
126+
@Test
127+
public void testToProtoWithAvroSchemaList() {
128+
CreateSchemaBundleRequest request =
129+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
130+
.setAvroSchema(ImmutableList.of(TEST_AVRO_SCHEMA, TEST_UPDATED_AVRO_SCHEMA));
131+
132+
com.google.bigtable.admin.v2.CreateSchemaBundleRequest requestProto =
133+
com.google.bigtable.admin.v2.CreateSchemaBundleRequest.newBuilder()
134+
.setParent(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID))
135+
.setSchemaBundleId(SCHEMA_BUNDLE_ID)
136+
.setSchemaBundle(
137+
com.google.bigtable.admin.v2.SchemaBundle.newBuilder()
138+
.setAvroSchema(
139+
com.google.bigtable.admin.v2.AvroSchema.newBuilder()
140+
.addJsonSchemas(TEST_AVRO_SCHEMA)
141+
.addJsonSchemas(TEST_UPDATED_AVRO_SCHEMA)
142+
.build())
143+
.build())
144+
.build();
145+
assertThat(request.toProto(PROJECT_ID, INSTANCE_ID)).isEqualTo(requestProto);
146+
}
147+
148+
@Test
149+
public void testEqualityWithAvroSchema() {
150+
CreateSchemaBundleRequest request =
151+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
152+
153+
assertThat(request)
154+
.isEqualTo(
155+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
156+
.setAvroSchema(TEST_AVRO_SCHEMA));
157+
158+
assertThat(request)
159+
.isNotEqualTo(
160+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
161+
.setAvroSchema(TEST_UPDATED_AVRO_SCHEMA));
162+
}
163+
164+
@Test
165+
public void testHashCodeWithAvroSchema() {
166+
CreateSchemaBundleRequest request =
167+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID).setAvroSchema(TEST_AVRO_SCHEMA);
168+
169+
assertThat(request.hashCode())
170+
.isEqualTo(
171+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
172+
.setAvroSchema(TEST_AVRO_SCHEMA)
173+
.hashCode());
174+
175+
assertThat(request.hashCode())
176+
.isNotEqualTo(
177+
CreateSchemaBundleRequest.of(TABLE_ID, SCHEMA_BUNDLE_ID)
178+
.setAvroSchema(TEST_UPDATED_AVRO_SCHEMA)
179+
.hashCode());
180+
}
181+
102182
private String getResourceFilePath(String filePath) throws URISyntaxException {
103183
ClassLoader cl = Thread.currentThread().getContextClassLoader();
104184
URL protoSchema = cl.getResource(filePath);

0 commit comments

Comments
 (0)