diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BQTableSchemaToProtoDescriptor.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BQTableSchemaToProtoDescriptor.java new file mode 100644 index 0000000000..cfb34a3f8e --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BQTableSchemaToProtoDescriptor.java @@ -0,0 +1,152 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.DescriptorProtos.DescriptorProto; +import com.google.protobuf.DescriptorProtos.FieldDescriptorProto; +import com.google.protobuf.DescriptorProtos.FileDescriptorProto; +import com.google.protobuf.Descriptors; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Descriptors.FileDescriptor; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; + +/** + * Converts a BQ table schema to protobuf descriptor. All field names will be converted to lowercase + * when constructing the protobuf descriptor. The mapping between field types and field modes are + * shown in the ImmutableMaps below. + */ +public class BQTableSchemaToProtoDescriptor { + private static ImmutableMap + BQTableSchemaModeMap = + ImmutableMap.of( + TableFieldSchema.Mode.NULLABLE, FieldDescriptorProto.Label.LABEL_OPTIONAL, + TableFieldSchema.Mode.REPEATED, FieldDescriptorProto.Label.LABEL_REPEATED, + TableFieldSchema.Mode.REQUIRED, FieldDescriptorProto.Label.LABEL_REQUIRED); + + private static ImmutableMap + BQTableSchemaTypeMap = + new ImmutableMap.Builder() + .put(TableFieldSchema.Type.BOOL, FieldDescriptorProto.Type.TYPE_BOOL) + .put(TableFieldSchema.Type.BYTES, FieldDescriptorProto.Type.TYPE_BYTES) + .put(TableFieldSchema.Type.DATE, FieldDescriptorProto.Type.TYPE_INT32) + .put(TableFieldSchema.Type.DATETIME, FieldDescriptorProto.Type.TYPE_INT64) + .put(TableFieldSchema.Type.DOUBLE, FieldDescriptorProto.Type.TYPE_DOUBLE) + .put(TableFieldSchema.Type.GEOGRAPHY, FieldDescriptorProto.Type.TYPE_STRING) + .put(TableFieldSchema.Type.INT64, FieldDescriptorProto.Type.TYPE_INT64) + .put(TableFieldSchema.Type.NUMERIC, FieldDescriptorProto.Type.TYPE_BYTES) + .put(TableFieldSchema.Type.STRING, FieldDescriptorProto.Type.TYPE_STRING) + .put(TableFieldSchema.Type.STRUCT, FieldDescriptorProto.Type.TYPE_MESSAGE) + .put(TableFieldSchema.Type.TIME, FieldDescriptorProto.Type.TYPE_INT64) + .put(TableFieldSchema.Type.TIMESTAMP, FieldDescriptorProto.Type.TYPE_INT64) + .build(); + + /** + * Converts TableFieldSchema to a Descriptors.Descriptor object. + * + * @param BQTableSchema + * @throws Descriptors.DescriptorValidationException + */ + public static Descriptor convertBQTableSchemaToProtoDescriptor(TableSchema BQTableSchema) + throws Descriptors.DescriptorValidationException { + Preconditions.checkNotNull(BQTableSchema, "BQTableSchema is null."); + return convertBQTableSchemaToProtoDescriptorImpl( + BQTableSchema, "root", new HashMap, Descriptor>()); + } + + /** + * Converts a TableFieldSchema to a Descriptors.Descriptor object. + * + * @param BQTableSchema + * @param scope Keeps track of current scope to prevent repeated naming while constructing + * descriptor. + * @param dependencyMap Stores already constructed descriptors to prevent reconstruction + * @throws Descriptors.DescriptorValidationException + */ + private static Descriptor convertBQTableSchemaToProtoDescriptorImpl( + TableSchema BQTableSchema, + String scope, + HashMap, Descriptor> dependencyMap) + throws Descriptors.DescriptorValidationException { + List dependenciesList = new ArrayList(); + List fields = new ArrayList(); + int index = 1; + for (TableFieldSchema BQTableField : BQTableSchema.getFieldsList()) { + String currentScope = scope + "__" + BQTableField.getName(); + if (BQTableField.getType() == TableFieldSchema.Type.STRUCT) { + ImmutableList fieldList = + ImmutableList.copyOf(BQTableField.getFieldsList()); + if (dependencyMap.containsKey(fieldList)) { + Descriptor descriptor = dependencyMap.get(fieldList); + dependenciesList.add(descriptor.getFile()); + fields.add(convertBQTableFieldToProtoField(BQTableField, index++, descriptor.getName())); + } else { + Descriptor descriptor = + convertBQTableSchemaToProtoDescriptorImpl( + TableSchema.newBuilder().addAllFields(fieldList).build(), + currentScope, + dependencyMap); + dependenciesList.add(descriptor.getFile()); + dependencyMap.put(fieldList, descriptor); + fields.add(convertBQTableFieldToProtoField(BQTableField, index++, currentScope)); + } + } else { + fields.add(convertBQTableFieldToProtoField(BQTableField, index++, currentScope)); + } + } + FileDescriptor[] dependenciesArray = new FileDescriptor[dependenciesList.size()]; + dependenciesArray = dependenciesList.toArray(dependenciesArray); + DescriptorProto descriptorProto = + DescriptorProto.newBuilder().setName(scope).addAllField(fields).build(); + FileDescriptorProto fileDescriptorProto = + FileDescriptorProto.newBuilder().addMessageType(descriptorProto).build(); + FileDescriptor fileDescriptor = + FileDescriptor.buildFrom(fileDescriptorProto, dependenciesArray); + Descriptor descriptor = fileDescriptor.findMessageTypeByName(scope); + return descriptor; + } + + /** + * Converts a BQTableField to ProtoField + * + * @param BQTableField BQ Field used to construct a FieldDescriptorProto + * @param index Index for protobuf fields. + * @param scope used to name descriptors + */ + private static FieldDescriptorProto convertBQTableFieldToProtoField( + TableFieldSchema BQTableField, int index, String scope) { + TableFieldSchema.Mode mode = BQTableField.getMode(); + String fieldName = BQTableField.getName().toLowerCase(); + if (BQTableField.getType() == TableFieldSchema.Type.STRUCT) { + return FieldDescriptorProto.newBuilder() + .setName(fieldName) + .setTypeName(scope) + .setLabel((FieldDescriptorProto.Label) BQTableSchemaModeMap.get(mode)) + .setNumber(index) + .build(); + } + return FieldDescriptorProto.newBuilder() + .setName(fieldName) + .setType((FieldDescriptorProto.Type) BQTableSchemaTypeMap.get(BQTableField.getType())) + .setLabel((FieldDescriptorProto.Label) BQTableSchemaModeMap.get(mode)) + .setNumber(index) + .build(); + } +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigDecimalByteStringEncoder.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigDecimalByteStringEncoder.java new file mode 100644 index 0000000000..b5c5992437 --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/BigDecimalByteStringEncoder.java @@ -0,0 +1,86 @@ +/* + * Copyright 2021 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * This code was ported from ZetaSQL and can be found here: + * https://github.com/google/zetasql/blob/c55f967a5ae35b476437210c529691d8a73f5507/java/com/google/zetasql/Value.java + */ + +package com.google.cloud.bigquery.storage.v1; + +import com.google.common.primitives.Bytes; +import com.google.protobuf.ByteString; +import java.math.BigDecimal; +import java.math.BigInteger; + +public class BigDecimalByteStringEncoder { + private static int NumericScale = 9; + private static final BigDecimal MAX_NUMERIC_VALUE = + new BigDecimal("99999999999999999999999999999.999999999"); + private static final BigDecimal MIN_NUMERIC_VALUE = + new BigDecimal("-99999999999999999999999999999.999999999"); + + public static ByteString encodeToNumericByteString(BigDecimal bigDecimal) { + ByteString byteString = + serializeBigDecimal( + bigDecimal, NumericScale, MAX_NUMERIC_VALUE, MIN_NUMERIC_VALUE, "ByteString"); + return byteString; + } + + public static BigDecimal decodeNumericByteString(ByteString byteString) { + BigDecimal bigDecimal = + deserializeBigDecimal( + byteString, NumericScale, MAX_NUMERIC_VALUE, MIN_NUMERIC_VALUE, "BigDecimal"); + return bigDecimal; + } + // Make these private and make public wrapper that internalizes these min/max/scale/type + private static BigDecimal deserializeBigDecimal( + ByteString serializedValue, + int scale, + BigDecimal maxValue, + BigDecimal minValue, + String typeName) { + byte[] bytes = serializedValue.toByteArray(); + // NUMERIC/BIGNUMERIC values are serialized as scaled integers in two's complement form in + // little endian order. BigInteger requires the same encoding but in big endian order, + // therefore we must reverse the bytes that come from the proto. + Bytes.reverse(bytes); + BigInteger scaledValue = new BigInteger(bytes); + BigDecimal decimalValue = new BigDecimal(scaledValue, scale); + if (decimalValue.compareTo(maxValue) > 0 || decimalValue.compareTo(minValue) < 0) { + throw new IllegalArgumentException(typeName + " overflow: " + decimalValue.toPlainString()); + } + return decimalValue; + } + /** Returns a numeric Value that equals to {@code v}. */ + private static ByteString serializeBigDecimal( + BigDecimal v, int scale, BigDecimal maxValue, BigDecimal minValue, String typeName) { + if (v.scale() > scale) { + throw new IllegalArgumentException( + typeName + " scale cannot exceed " + scale + ": " + v.toPlainString()); + } + if (v.compareTo(maxValue) > 0 || v.compareTo(minValue) < 0) { + throw new IllegalArgumentException(typeName + " overflow: " + v.toPlainString()); + } + byte[] bytes = v.setScale(scale).unscaledValue().toByteArray(); + // NUMERIC/BIGNUMERIC values are serialized as scaled integers in two's complement form in + // little endian + // order. BigInteger requires the same encoding but in big endian order, therefore we must + // reverse the bytes that come from the proto. + Bytes.reverse(bytes); + return ByteString.copyFrom(bytes); + } +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/CivilTimeEncoder.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/CivilTimeEncoder.java new file mode 100644 index 0000000000..be2d1fd929 --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/CivilTimeEncoder.java @@ -0,0 +1,314 @@ +/* + * Copyright 2021 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import static com.google.common.base.Preconditions.checkArgument; + +import org.threeten.bp.DateTimeException; +import org.threeten.bp.LocalDateTime; +import org.threeten.bp.LocalTime; +import org.threeten.bp.temporal.ChronoUnit; + +/** + * Ported from ZetaSQL CivilTimeEncoder Original code can be found at: + * https://github.com/google/zetasql/blob/master/java/com/google/zetasql/CivilTimeEncoder.java + * Encoder for TIME and DATETIME values, according to civil_time encoding. + * + *

The valid range and number of bits required by each date/time field is as the following: + * + * + * + * + * + * + * + * + * + * + * + *
Field Range #Bits
Year [1, 9999] 14
Month [1, 12] 4
Day [1, 31] 5
Hour [0, 23] 5
Minute [0, 59] 6
Second [0, 59]* 6
Micros [0, 999999] 20
Nanos [0, 999999999] 30
+ * + *

* Leap second is not supported. + * + *

When encoding the TIME or DATETIME into a bit field, larger date/time field is on the more + * significant side. + */ +public final class CivilTimeEncoder { + private static final int NANO_LENGTH = 30; + private static final int MICRO_LENGTH = 20; + + private static final int NANO_SHIFT = 0; + private static final int MICRO_SHIFT = 0; + private static final int SECOND_SHIFT = 0; + private static final int MINUTE_SHIFT = 6; + private static final int HOUR_SHIFT = 12; + private static final int DAY_SHIFT = 17; + private static final int MONTH_SHIFT = 22; + private static final int YEAR_SHIFT = 26; + + private static final long NANO_MASK = 0x3FFFFFFFL; + private static final long MICRO_MASK = 0xFFFFFL; + private static final long SECOND_MASK = 0x3FL; + private static final long MINUTE_MASK = 0xFC0L; + private static final long HOUR_MASK = 0x1F000L; + private static final long DAY_MASK = 0x3E0000L; + private static final long MONTH_MASK = 0x3C00000L; + private static final long YEAR_MASK = 0xFFFC000000L; + + private static final long TIME_SECONDS_MASK = 0x1FFFFL; + private static final long TIME_MICROS_MASK = 0x1FFFFFFFFFL; + private static final long TIME_NANOS_MASK = 0x7FFFFFFFFFFFL; + private static final long DATETIME_SECONDS_MASK = 0xFFFFFFFFFFL; + private static final long DATETIME_MICROS_MASK = 0xFFFFFFFFFFFFFFFL; + + /** + * Encodes {@code time} as a 4-byte integer with seconds precision. + * + *

Encoding is as the following: + * + * 

+   *      3         2         1
+   * MSB 10987654321098765432109876543210 LSB
+   *                    | H ||  M ||  S |
+   *
+ * + * @see #decodePacked32TimeSeconds(int) + */ + @SuppressWarnings("GoodTime-ApiWithNumericTimeUnit") + private static int encodePacked32TimeSeconds(LocalTime time) { + checkValidTimeSeconds(time); + int bitFieldTimeSeconds = 0x0; + bitFieldTimeSeconds |= time.getHour() << HOUR_SHIFT; + bitFieldTimeSeconds |= time.getMinute() << MINUTE_SHIFT; + bitFieldTimeSeconds |= time.getSecond() << SECOND_SHIFT; + return bitFieldTimeSeconds; + } + + /** + * Decodes {@code bitFieldTimeSeconds} as a {@link LocalTime} with seconds precision. + * + *

Encoding is as the following: + * + * 

+   *      3         2         1
+   * MSB 10987654321098765432109876543210 LSB
+   *                    | H ||  M ||  S |
+   *
+ * + * @see #encodePacked32TimeSeconds(LocalTime) + */ + @SuppressWarnings("GoodTime-ApiWithNumericTimeUnit") + private static LocalTime decodePacked32TimeSeconds(int bitFieldTimeSeconds) { + checkValidBitField(bitFieldTimeSeconds, TIME_SECONDS_MASK); + int hourOfDay = getFieldFromBitField(bitFieldTimeSeconds, HOUR_MASK, HOUR_SHIFT); + int minuteOfHour = getFieldFromBitField(bitFieldTimeSeconds, MINUTE_MASK, MINUTE_SHIFT); + int secondOfMinute = getFieldFromBitField(bitFieldTimeSeconds, SECOND_MASK, SECOND_SHIFT); + // LocalTime validates the input parameters. + try { + return LocalTime.of(hourOfDay, minuteOfHour, secondOfMinute); + } catch (DateTimeException e) { + throw new IllegalArgumentException(e.getMessage(), e); + } + } + + /** + * Encodes {@code time} as a 8-byte integer with microseconds precision. + * + *

Encoding is as the following: + * + * 

+   *        6         5         4         3         2         1
+   * MSB 3210987654321098765432109876543210987654321098765432109876543210 LSB
+   *                                | H ||  M ||  S ||-------micros-----|
+   *
+ * + * @see #decodePacked64TimeMicros(long) + * @see #encodePacked64TimeMicros(LocalTime) + */ + @SuppressWarnings("GoodTime") + public static long encodePacked64TimeMicros(LocalTime time) { + checkValidTimeMicros(time); + return (((long) encodePacked32TimeSeconds(time)) << MICRO_LENGTH) | (time.getNano() / 1_000L); + } + + /** + * Decodes {@code bitFieldTimeMicros} as a {@link LocalTime} with microseconds precision. + * + *

Encoding is as the following: + * + * 

+   *        6         5         4         3         2         1
+   * MSB 3210987654321098765432109876543210987654321098765432109876543210 LSB
+   *                                | H ||  M ||  S ||-------micros-----|
+   *
+ * + * @see #encodePacked64TimeMicros(LocalTime) + */ + @SuppressWarnings("GoodTime-ApiWithNumericTimeUnit") + public static LocalTime decodePacked64TimeMicros(long bitFieldTimeMicros) { + checkValidBitField(bitFieldTimeMicros, TIME_MICROS_MASK); + int bitFieldTimeSeconds = (int) (bitFieldTimeMicros >> MICRO_LENGTH); + LocalTime timeSeconds = decodePacked32TimeSeconds(bitFieldTimeSeconds); + int microOfSecond = getFieldFromBitField(bitFieldTimeMicros, MICRO_MASK, MICRO_SHIFT); + checkValidMicroOfSecond(microOfSecond); + LocalTime time = timeSeconds.withNano(microOfSecond * 1000); + checkValidTimeMicros(time); + return time; + } + + /** + * Encodes {@code dateTime} as a 8-byte integer with seconds precision. + * + *

Encoding is as the following: + * + * 

+   *        6         5         4         3         2         1
+   * MSB 3210987654321098765432109876543210987654321098765432109876543210 LSB
+   *                             |--- year ---||m || D || H ||  M ||  S |
+   *
+ * + * @see #decodePacked64DatetimeSeconds(long) + */ + @SuppressWarnings("GoodTime-ApiWithNumericTimeUnit") + private static long encodePacked64DatetimeSeconds(LocalDateTime dateTime) { + checkValidDateTimeSeconds(dateTime); + long bitFieldDatetimeSeconds = 0x0L; + bitFieldDatetimeSeconds |= (long) dateTime.getYear() << YEAR_SHIFT; + bitFieldDatetimeSeconds |= (long) dateTime.getMonthValue() << MONTH_SHIFT; + bitFieldDatetimeSeconds |= (long) dateTime.getDayOfMonth() << DAY_SHIFT; + bitFieldDatetimeSeconds |= (long) encodePacked32TimeSeconds(dateTime.toLocalTime()); + return bitFieldDatetimeSeconds; + } + + /** + * Decodes {@code bitFieldDatetimeSeconds} as a {@link LocalDateTime} with seconds precision. + * + *

Encoding is as the following: + * + * 

+   *        6         5         4         3         2         1
+   * MSB 3210987654321098765432109876543210987654321098765432109876543210 LSBa
+   *                             |--- year ---||m || D || H ||  M ||  S |
+   *
+ * + * @see #encodePacked64DatetimeSeconds(LocalDateTime) + */ + @SuppressWarnings("GoodTime-ApiWithNumericTimeUnit") + private static LocalDateTime decodePacked64DatetimeSeconds(long bitFieldDatetimeSeconds) { + checkValidBitField(bitFieldDatetimeSeconds, DATETIME_SECONDS_MASK); + int bitFieldTimeSeconds = (int) (bitFieldDatetimeSeconds & TIME_SECONDS_MASK); + LocalTime timeSeconds = decodePacked32TimeSeconds(bitFieldTimeSeconds); + int year = getFieldFromBitField(bitFieldDatetimeSeconds, YEAR_MASK, YEAR_SHIFT); + int monthOfYear = getFieldFromBitField(bitFieldDatetimeSeconds, MONTH_MASK, MONTH_SHIFT); + int dayOfMonth = getFieldFromBitField(bitFieldDatetimeSeconds, DAY_MASK, DAY_SHIFT); + try { + LocalDateTime dateTime = + LocalDateTime.of( + year, + monthOfYear, + dayOfMonth, + timeSeconds.getHour(), + timeSeconds.getMinute(), + timeSeconds.getSecond()); + checkValidDateTimeSeconds(dateTime); + return dateTime; + } catch (DateTimeException e) { + throw new IllegalArgumentException(e.getMessage(), e); + } + } + + /** + * Encodes {@code dateTime} as a 8-byte integer with microseconds precision. + * + *

Encoding is as the following: + * + * 

+   *        6         5         4         3         2         1
+   * MSB 3210987654321098765432109876543210987654321098765432109876543210 LSB
+   *         |--- year ---||m || D || H ||  M ||  S ||-------micros-----|
+   *
+ * + * @see #decodePacked64DatetimeMicros(long) + */ + @SuppressWarnings({"GoodTime-ApiWithNumericTimeUnit", "JavaLocalDateTimeGetNano"}) + public static long encodePacked64DatetimeMicros(LocalDateTime dateTime) { + checkValidDateTimeMicros(dateTime); + return (encodePacked64DatetimeSeconds(dateTime) << MICRO_LENGTH) + | (dateTime.getNano() / 1_000L); + } + + /** + * Decodes {@code bitFieldDatetimeMicros} as a {@link LocalDateTime} with microseconds precision. + * + *

Encoding is as the following: + * + * 

+   *        6         5         4         3         2         1
+   * MSB 3210987654321098765432109876543210987654321098765432109876543210 LSB
+   *         |--- year ---||m || D || H ||  M ||  S ||-------micros-----|
+   *
+ * + * @see #encodePacked64DatetimeMicros(LocalDateTime) + */ + @SuppressWarnings("GoodTime-ApiWithNumericTimeUnit") + public static LocalDateTime decodePacked64DatetimeMicros(long bitFieldDatetimeMicros) { + checkValidBitField(bitFieldDatetimeMicros, DATETIME_MICROS_MASK); + long bitFieldDatetimeSeconds = bitFieldDatetimeMicros >> MICRO_LENGTH; + LocalDateTime dateTimeSeconds = decodePacked64DatetimeSeconds(bitFieldDatetimeSeconds); + int microOfSecond = getFieldFromBitField(bitFieldDatetimeMicros, MICRO_MASK, MICRO_SHIFT); + checkValidMicroOfSecond(microOfSecond); + LocalDateTime dateTime = dateTimeSeconds.withNano(microOfSecond * 1_000); + checkValidDateTimeMicros(dateTime); + return dateTime; + } + + private static int getFieldFromBitField(long bitField, long mask, int shift) { + return (int) ((bitField & mask) >> shift); + } + + private static void checkValidTimeSeconds(LocalTime time) { + checkArgument(time.getHour() >= 0 && time.getHour() <= 23); + checkArgument(time.getMinute() >= 0 && time.getMinute() <= 59); + checkArgument(time.getSecond() >= 0 && time.getSecond() <= 59); + } + + private static void checkValidDateTimeSeconds(LocalDateTime dateTime) { + checkArgument(dateTime.getYear() >= 1 && dateTime.getYear() <= 9999); + checkArgument(dateTime.getMonthValue() >= 1 && dateTime.getMonthValue() <= 12); + checkArgument(dateTime.getDayOfMonth() >= 1 && dateTime.getDayOfMonth() <= 31); + checkValidTimeSeconds(dateTime.toLocalTime()); + } + + private static void checkValidTimeMicros(LocalTime time) { + checkValidTimeSeconds(time); + checkArgument(time.equals(time.truncatedTo(ChronoUnit.MICROS))); + } + + private static void checkValidDateTimeMicros(LocalDateTime dateTime) { + checkValidDateTimeSeconds(dateTime); + checkArgument(dateTime.equals(dateTime.truncatedTo(ChronoUnit.MICROS))); + } + + private static void checkValidMicroOfSecond(int microOfSecond) { + checkArgument(microOfSecond >= 0 && microOfSecond <= 999999); + } + + private static void checkValidBitField(long bitField, long mask) { + checkArgument((bitField & ~mask) == 0x0L); + } + + private CivilTimeEncoder() {} +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java new file mode 100644 index 0000000000..4a13b4dfcd --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java @@ -0,0 +1,334 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.core.ApiFuture; +import com.google.api.gax.batching.FlowControlSettings; +import com.google.api.gax.core.CredentialsProvider; +import com.google.api.gax.rpc.TransportChannelProvider; +import com.google.common.base.Preconditions; +import com.google.protobuf.Descriptors; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Message; +import java.io.IOException; +import java.util.logging.Logger; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import javax.annotation.Nullable; +import org.json.JSONArray; +import org.json.JSONObject; + +/** + * A StreamWriter that can write JSON data (JSONObjects) to BigQuery tables. The JsonStreamWriter is + * built on top of a StreamWriter, and it simply converts all JSON data to protobuf messages then + * calls StreamWriter's append() method to write to BigQuery tables. + */ +public class JsonStreamWriter implements AutoCloseable { + private static String streamPatternString = + "projects/[^/]+/datasets/[^/]+/tables/[^/]+/streams/[^/]+"; + private static Pattern streamPattern = Pattern.compile(streamPatternString); + private static final Logger LOG = Logger.getLogger(JsonStreamWriter.class.getName()); + + private BigQueryWriteClient client; + private String streamName; + private StreamWriter streamWriter; + private StreamWriter.Builder streamWriterBuilder; + private Descriptor descriptor; + private TableSchema tableSchema; + + /** + * Constructs the JsonStreamWriter + * + * @param builder The Builder object for the JsonStreamWriter + */ + private JsonStreamWriter(Builder builder) + throws Descriptors.DescriptorValidationException, IllegalArgumentException, IOException, + InterruptedException { + this.client = builder.client; + this.descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(builder.tableSchema); + + if (this.client == null) { + streamWriterBuilder = StreamWriter.newBuilder(builder.streamName); + } else { + streamWriterBuilder = StreamWriter.newBuilder(builder.streamName, builder.client); + } + streamWriterBuilder.setWriterSchema(ProtoSchemaConverter.convert(this.descriptor)); + setStreamWriterSettings( + builder.channelProvider, + builder.credentialsProvider, + builder.endpoint, + builder.flowControlSettings, + builder.traceId); + this.streamWriter = streamWriterBuilder.build(); + this.streamName = builder.streamName; + } + + /** + * Writes a JSONArray that contains JSONObjects to the BigQuery table by first converting the JSON + * data to protobuf messages, then using StreamWriter's append() to write the data. + * + * @param jsonArr The JSON array that contains JSONObjects to be written + * @return ApiFuture returns an AppendRowsResponse message wrapped in an + * ApiFuture + */ + public ApiFuture append(JSONArray jsonArr) { + return append(jsonArr, -1); + } + + /** + * Writes a JSONArray that contains JSONObjects to the BigQuery table by first converting the JSON + * data to protobuf messages, then using StreamWriter's append() to write the data. + * + * @param jsonArr The JSON array that contains JSONObjects to be written + * @param offset Offset for deduplication + * @return ApiFuture returns an AppendRowsResponse message wrapped in an + * ApiFuture + */ + public ApiFuture append(JSONArray jsonArr, long offset) { + ProtoRows.Builder rowsBuilder = ProtoRows.newBuilder(); + // Any error in convertJsonToProtoMessage will throw an + // IllegalArgumentException/IllegalStateException/NullPointerException and will halt processing + // of JSON data. + for (int i = 0; i < jsonArr.length(); i++) { + JSONObject json = jsonArr.getJSONObject(i); + Message protoMessage = JsonToProtoMessage.convertJsonToProtoMessage(this.descriptor, json); + rowsBuilder.addSerializedRows(protoMessage.toByteString()); + } + // Need to make sure refreshAppendAndSetDescriptor finish first before this can run + synchronized (this) { + final ApiFuture appendResponseFuture = + this.streamWriter.append(rowsBuilder.build(), offset); + return appendResponseFuture; + } + } + + /** + * Gets streamName + * + * @return String + */ + public String getStreamName() { + return this.streamName; + } + + /** + * Gets current descriptor + * + * @return Descriptor + */ + public Descriptor getDescriptor() { + return this.descriptor; + } + + /** Sets all StreamWriter settings. */ + private void setStreamWriterSettings( + @Nullable TransportChannelProvider channelProvider, + @Nullable CredentialsProvider credentialsProvider, + @Nullable String endpoint, + @Nullable FlowControlSettings flowControlSettings, + @Nullable String traceId) { + if (channelProvider != null) { + streamWriterBuilder.setChannelProvider(channelProvider); + } + if (credentialsProvider != null) { + streamWriterBuilder.setCredentialsProvider(credentialsProvider); + } + if (endpoint != null) { + streamWriterBuilder.setEndpoint(endpoint); + } + if (traceId != null) { + streamWriterBuilder.setTraceId("JsonWriterBeta_" + traceId); + } else { + streamWriterBuilder.setTraceId("JsonWriterBeta:null"); + } + if (flowControlSettings != null) { + if (flowControlSettings.getMaxOutstandingRequestBytes() != null) { + streamWriterBuilder.setMaxInflightBytes( + flowControlSettings.getMaxOutstandingRequestBytes()); + } + if (flowControlSettings.getMaxOutstandingElementCount() != null) { + streamWriterBuilder.setMaxInflightRequests( + flowControlSettings.getMaxOutstandingElementCount()); + } + } + } + + /** + * newBuilder that constructs a JsonStreamWriter builder with BigQuery client being initialized by + * StreamWriter by default. + * + * @param streamOrTableName name of the stream that must follow + * "projects/[^/]+/datasets/[^/]+/tables/[^/]+/streams/[^/]+" or table name + * "projects/[^/]+/datasets/[^/]+/tables/[^/]+" + * @param tableSchema The schema of the table when the stream was created, which is passed back + * through {@code WriteStream} + * @return Builder + */ + public static Builder newBuilder(String streamOrTableName, TableSchema tableSchema) { + Preconditions.checkNotNull(streamOrTableName, "StreamOrTableName is null."); + Preconditions.checkNotNull(tableSchema, "TableSchema is null."); + return new Builder(streamOrTableName, tableSchema, null); + } + + /** + * newBuilder that constructs a JsonStreamWriter builder. + * + * @param streamOrTableName name of the stream that must follow + * "projects/[^/]+/datasets/[^/]+/tables/[^/]+/streams/[^/]+" + * @param tableSchema The schema of the table when the stream was created, which is passed back + * through {@code WriteStream} + * @param client + * @return Builder + */ + public static Builder newBuilder( + String streamOrTableName, TableSchema tableSchema, BigQueryWriteClient client) { + Preconditions.checkNotNull(streamOrTableName, "StreamName is null."); + Preconditions.checkNotNull(tableSchema, "TableSchema is null."); + Preconditions.checkNotNull(client, "BigQuery client is null."); + return new Builder(streamOrTableName, tableSchema, client); + } + + /** Closes the underlying StreamWriter. */ + @Override + public void close() { + this.streamWriter.close(); + } + + public static final class Builder { + private String streamName; + private BigQueryWriteClient client; + private TableSchema tableSchema; + + private TransportChannelProvider channelProvider; + private CredentialsProvider credentialsProvider; + private FlowControlSettings flowControlSettings; + private String endpoint; + private boolean createDefaultStream = false; + private String traceId; + + private static String streamPatternString = + "(projects/[^/]+/datasets/[^/]+/tables/[^/]+)/streams/[^/]+"; + private static String tablePatternString = "(projects/[^/]+/datasets/[^/]+/tables/[^/]+)"; + + private static Pattern streamPattern = Pattern.compile(streamPatternString); + private static Pattern tablePattern = Pattern.compile(tablePatternString); + + /** + * Constructor for JsonStreamWriter's Builder + * + * @param streamOrTableName name of the stream that must follow + * "projects/[^/]+/datasets/[^/]+/tables/[^/]+/streams/[^/]+" or + * "projects/[^/]+/datasets/[^/]+/tables/[^/]+" + * @param tableSchema schema used to convert Json to proto messages. + * @param client + */ + private Builder(String streamOrTableName, TableSchema tableSchema, BigQueryWriteClient client) { + Matcher streamMatcher = streamPattern.matcher(streamOrTableName); + if (!streamMatcher.matches()) { + Matcher tableMatcher = tablePattern.matcher(streamOrTableName); + if (!tableMatcher.matches()) { + throw new IllegalArgumentException("Invalid name: " + streamOrTableName); + } else { + this.streamName = streamOrTableName + "/_default"; + } + } else { + this.streamName = streamOrTableName; + } + this.tableSchema = tableSchema; + this.client = client; + } + + /** + * Setter for the underlying StreamWriter's TransportChannelProvider. + * + * @param channelProvider + * @return Builder + */ + public Builder setChannelProvider(TransportChannelProvider channelProvider) { + this.channelProvider = + Preconditions.checkNotNull(channelProvider, "ChannelProvider is null."); + return this; + } + + /** + * Setter for the underlying StreamWriter's CredentialsProvider. + * + * @param credentialsProvider + * @return Builder + */ + public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) { + this.credentialsProvider = + Preconditions.checkNotNull(credentialsProvider, "CredentialsProvider is null."); + return this; + } + + /** + * Setter for the underlying StreamWriter's FlowControlSettings. + * + * @param flowControlSettings + * @return Builder + */ + public Builder setFlowControlSettings(FlowControlSettings flowControlSettings) { + Preconditions.checkNotNull(flowControlSettings, "FlowControlSettings is null."); + this.flowControlSettings = + Preconditions.checkNotNull(flowControlSettings, "FlowControlSettings is null."); + return this; + } + + /** + * Stream name on the builder. + * + * @return Builder + */ + public String getStreamName() { + return streamName; + } + + /** + * Setter for the underlying StreamWriter's Endpoint. + * + * @param endpoint + * @return Builder + */ + public Builder setEndpoint(String endpoint) { + this.endpoint = Preconditions.checkNotNull(endpoint, "Endpoint is null."); + return this; + } + + /** + * Setter for a traceId to help identify traffic origin. + * + * @param traceId + * @return Builder + */ + public Builder setTraceId(String traceId) { + this.traceId = Preconditions.checkNotNull(traceId, "TraceId is null."); + return this; + } + + /** + * Builds JsonStreamWriter + * + * @return JsonStreamWriter + */ + public JsonStreamWriter build() + throws Descriptors.DescriptorValidationException, IllegalArgumentException, IOException, + InterruptedException { + return new JsonStreamWriter(this); + } + } +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessage.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessage.java new file mode 100644 index 0000000000..d2cdbab982 --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessage.java @@ -0,0 +1,349 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Descriptors.FieldDescriptor; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Message; +import com.google.protobuf.UninitializedMessageException; +import java.util.logging.Logger; +import org.json.JSONArray; +import org.json.JSONException; +import org.json.JSONObject; + +/** + * Converts Json data to protocol buffer messages given the protocol buffer descriptor. The protobuf + * descriptor must have all fields lowercased. + */ +public class JsonToProtoMessage { + private static final Logger LOG = Logger.getLogger(JsonToProtoMessage.class.getName()); + private static ImmutableMap FieldTypeToDebugMessage = + new ImmutableMap.Builder() + .put(FieldDescriptor.Type.BOOL, "boolean") + .put(FieldDescriptor.Type.BYTES, "bytes") + .put(FieldDescriptor.Type.INT32, "int32") + .put(FieldDescriptor.Type.DOUBLE, "double") + .put(FieldDescriptor.Type.INT64, "int64") + .put(FieldDescriptor.Type.STRING, "string") + .put(FieldDescriptor.Type.MESSAGE, "object") + .build(); + + /** + * Converts Json data to protocol buffer messages given the protocol buffer descriptor. + * + * @param protoSchema + * @param json + * @throws IllegalArgumentException when JSON data is not compatible with proto descriptor. + */ + public static DynamicMessage convertJsonToProtoMessage(Descriptor protoSchema, JSONObject json) + throws IllegalArgumentException { + Preconditions.checkNotNull(json, "JSONObject is null."); + Preconditions.checkNotNull(protoSchema, "Protobuf descriptor is null."); + Preconditions.checkState(json.length() != 0, "JSONObject is empty."); + + return convertJsonToProtoMessageImpl(protoSchema, json, "root", /*topLevel=*/ true); + } + + /** + * Converts Json data to protocol buffer messages given the protocol buffer descriptor. + * + * @param protoSchema + * @param json + * @param jsonScope Debugging purposes + * @param topLevel checks if root level has any matching fields. + * @throws IllegalArgumentException when JSON data is not compatible with proto descriptor. + */ + private static DynamicMessage convertJsonToProtoMessageImpl( + Descriptor protoSchema, JSONObject json, String jsonScope, boolean topLevel) + throws IllegalArgumentException { + + DynamicMessage.Builder protoMsg = DynamicMessage.newBuilder(protoSchema); + String[] jsonNames = JSONObject.getNames(json); + if (jsonNames == null) { + return protoMsg.build(); + } + for (int i = 0; i < jsonNames.length; i++) { + String jsonName = jsonNames[i]; + // We want lowercase here to support case-insensitive data writes. + // The protobuf descriptor that is used is assumed to have all lowercased fields + String jsonLowercaseName = jsonName.toLowerCase(); + String currentScope = jsonScope + "." + jsonName; + FieldDescriptor field = protoSchema.findFieldByName(jsonLowercaseName); + if (field == null) { + throw new IllegalArgumentException( + String.format("JSONObject has fields unknown to BigQuery: %s.", currentScope)); + } + if (!field.isRepeated()) { + fillField(protoMsg, field, json, jsonName, currentScope); + } else { + fillRepeatedField(protoMsg, field, json, jsonName, currentScope); + } + } + + DynamicMessage msg; + try { + msg = protoMsg.build(); + } catch (UninitializedMessageException e) { + String errorMsg = e.getMessage(); + int idxOfColon = errorMsg.indexOf(":"); + String missingFieldName = errorMsg.substring(idxOfColon + 2); + throw new IllegalArgumentException( + String.format( + "JSONObject does not have the required field %s.%s.", jsonScope, missingFieldName)); + } + if (topLevel && msg.getSerializedSize() == 0) { + throw new IllegalArgumentException("The created protobuf message is empty."); + } + return msg; + } + + /** + * Fills a non-repetaed protoField with the json data. + * + * @param protoMsg The protocol buffer message being constructed + * @param fieldDescriptor + * @param json + * @param exactJsonKeyName Exact key name in JSONObject instead of lowercased version + * @param currentScope Debugging purposes + * @throws IllegalArgumentException when JSON data is not compatible with proto descriptor. + */ + private static void fillField( + DynamicMessage.Builder protoMsg, + FieldDescriptor fieldDescriptor, + JSONObject json, + String exactJsonKeyName, + String currentScope) + throws IllegalArgumentException { + + java.lang.Object val = json.get(exactJsonKeyName); + if (val == JSONObject.NULL) { + return; + } + switch (fieldDescriptor.getType()) { + case BOOL: + if (val instanceof Boolean) { + protoMsg.setField(fieldDescriptor, (Boolean) val); + return; + } + break; + case BYTES: + if (val instanceof ByteString) { + protoMsg.setField(fieldDescriptor, ((ByteString) val).toByteArray()); + return; + } else if (val instanceof JSONArray) { + try { + byte[] bytes = new byte[((JSONArray) val).length()]; + for (int j = 0; j < ((JSONArray) val).length(); j++) { + bytes[j] = (byte) ((JSONArray) val).getInt(j); + if (bytes[j] != ((JSONArray) val).getInt(j)) { + throw new IllegalArgumentException( + String.format( + "Error: " + + currentScope + + "[" + + j + + "] could not be converted to byte[].")); + } + } + protoMsg.setField(fieldDescriptor, bytes); + } catch (JSONException e) { + throw new IllegalArgumentException( + String.format("Error: " + currentScope + "could not be converted to byte[].")); + } + } + break; + case INT64: + if (val instanceof Integer) { + protoMsg.setField(fieldDescriptor, new Long((Integer) val)); + return; + } else if (val instanceof Long) { + protoMsg.setField(fieldDescriptor, (Long) val); + return; + } + break; + case INT32: + if (val instanceof Integer) { + protoMsg.setField(fieldDescriptor, (Integer) val); + return; + } + break; + case STRING: + if (val instanceof String) { + protoMsg.setField(fieldDescriptor, (String) val); + return; + } + break; + case DOUBLE: + if (val instanceof Double) { + protoMsg.setField(fieldDescriptor, (Double) val); + return; + } else if (val instanceof Float) { + protoMsg.setField(fieldDescriptor, new Double((Float) val)); + return; + } + break; + case MESSAGE: + if (val instanceof JSONObject) { + Message.Builder message = protoMsg.newBuilderForField(fieldDescriptor); + protoMsg.setField( + fieldDescriptor, + convertJsonToProtoMessageImpl( + fieldDescriptor.getMessageType(), + json.getJSONObject(exactJsonKeyName), + currentScope, + /*topLevel =*/ false)); + return; + } + break; + } + throw new IllegalArgumentException( + String.format( + "JSONObject does not have a %s field at %s.", + FieldTypeToDebugMessage.get(fieldDescriptor.getType()), currentScope)); + } + + /** + * Fills a repeated protoField with the json data. + * + * @param protoMsg The protocol buffer message being constructed + * @param fieldDescriptor + * @param json If root level has no matching fields, throws exception. + * @param exactJsonKeyName Exact key name in JSONObject instead of lowercased version + * @param currentScope Debugging purposes + * @throws IllegalArgumentException when JSON data is not compatible with proto descriptor. + */ + private static void fillRepeatedField( + DynamicMessage.Builder protoMsg, + FieldDescriptor fieldDescriptor, + JSONObject json, + String exactJsonKeyName, + String currentScope) + throws IllegalArgumentException { + + JSONArray jsonArray; + try { + jsonArray = json.getJSONArray(exactJsonKeyName); + } catch (JSONException e) { + throw new IllegalArgumentException( + "JSONObject does not have a array field at " + currentScope + "."); + } + java.lang.Object val; + int index; + boolean fail = false; + for (int i = 0; i < jsonArray.length(); i++) { + val = jsonArray.get(i); + index = i; + switch (fieldDescriptor.getType()) { + case BOOL: + if (val instanceof Boolean) { + protoMsg.addRepeatedField(fieldDescriptor, (Boolean) val); + } else { + fail = true; + } + break; + case BYTES: + if (val instanceof JSONArray) { + try { + byte[] bytes = new byte[((JSONArray) val).length()]; + for (int j = 0; j < ((JSONArray) val).length(); j++) { + bytes[j] = (byte) ((JSONArray) val).getInt(j); + if (bytes[j] != ((JSONArray) val).getInt(j)) { + throw new IllegalArgumentException( + String.format( + "Error: " + + currentScope + + "[" + + index + + "] could not be converted to byte[].")); + } + } + protoMsg.addRepeatedField(fieldDescriptor, bytes); + } catch (JSONException e) { + throw new IllegalArgumentException( + String.format( + "Error: " + + currentScope + + "[" + + index + + "] could not be converted to byte[].")); + } + } else if (val instanceof ByteString) { + protoMsg.addRepeatedField(fieldDescriptor, ((ByteString) val).toByteArray()); + return; + } else { + fail = true; + } + break; + case INT64: + if (val instanceof Integer) { + protoMsg.addRepeatedField(fieldDescriptor, new Long((Integer) val)); + } else if (val instanceof Long) { + protoMsg.addRepeatedField(fieldDescriptor, (Long) val); + } else { + fail = true; + } + break; + case INT32: + if (val instanceof Integer) { + protoMsg.addRepeatedField(fieldDescriptor, (Integer) val); + } else { + fail = true; + } + break; + case STRING: + if (val instanceof String) { + protoMsg.addRepeatedField(fieldDescriptor, (String) val); + } else { + fail = true; + } + break; + case DOUBLE: + if (val instanceof Double) { + protoMsg.addRepeatedField(fieldDescriptor, (Double) val); + } else if (val instanceof Float) { + protoMsg.addRepeatedField(fieldDescriptor, new Double((float) val)); + } else { + fail = true; + } + break; + case MESSAGE: + if (val instanceof JSONObject) { + Message.Builder message = protoMsg.newBuilderForField(fieldDescriptor); + protoMsg.addRepeatedField( + fieldDescriptor, + convertJsonToProtoMessageImpl( + fieldDescriptor.getMessageType(), + jsonArray.getJSONObject(i), + currentScope, + /*topLevel =*/ false)); + } else { + fail = true; + } + break; + } + if (fail) { + throw new IllegalArgumentException( + String.format( + "JSONObject does not have a %s field at %s[%d].", + FieldTypeToDebugMessage.get(fieldDescriptor.getType()), currentScope, index)); + } + } + } +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ProtoSchemaConverter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ProtoSchemaConverter.java new file mode 100644 index 0000000000..c55d4181a5 --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ProtoSchemaConverter.java @@ -0,0 +1,118 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.InvalidArgumentException; +import com.google.protobuf.DescriptorProtos.DescriptorProto; +import com.google.protobuf.DescriptorProtos.EnumDescriptorProto; +import com.google.protobuf.DescriptorProtos.FieldDescriptorProto; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Descriptors.FieldDescriptor; +import io.grpc.Status; +import java.util.HashSet; +import java.util.Set; + +// A Converter class that turns a native protobuf::DescriptorProto to a self contained +// protobuf::DescriptorProto +// that can be reconstructed by the backend. +public class ProtoSchemaConverter { + private static String getNameFromFullName(String fullName) { + return fullName.replace('.', '_'); + } + + private static ProtoSchema convertInternal( + Descriptor input, + Set visitedTypes, + Set enumTypes, + Set structTypes, + DescriptorProto.Builder rootProtoSchema) { + DescriptorProto.Builder resultProto = DescriptorProto.newBuilder(); + if (rootProtoSchema == null) { + rootProtoSchema = resultProto; + } + String protoFullName = input.getFullName(); + String protoName = getNameFromFullName(protoFullName); + resultProto.setName(protoName); + Set localEnumTypes = new HashSet(); + visitedTypes.add(input.getFullName()); + for (int i = 0; i < input.getFields().size(); i++) { + FieldDescriptor inputField = input.getFields().get(i); + FieldDescriptorProto.Builder resultField = inputField.toProto().toBuilder(); + if (inputField.getType() == FieldDescriptor.Type.GROUP + || inputField.getType() == FieldDescriptor.Type.MESSAGE) { + String msgFullName = inputField.getMessageType().getFullName(); + String msgName = getNameFromFullName(msgFullName); + if (structTypes.contains(msgFullName)) { + resultField.setTypeName(msgName); + } else { + if (visitedTypes.contains(msgFullName)) { + throw new InvalidArgumentException( + "Recursive type is not supported:" + inputField.getMessageType().getFullName(), + null, + GrpcStatusCode.of(Status.Code.INVALID_ARGUMENT), + false); + } + visitedTypes.add(msgFullName); + rootProtoSchema.addNestedType( + convertInternal( + inputField.getMessageType(), + visitedTypes, + enumTypes, + structTypes, + rootProtoSchema) + .getProtoDescriptor()); + visitedTypes.remove(msgFullName); + resultField.setTypeName( + rootProtoSchema.getNestedType(rootProtoSchema.getNestedTypeCount() - 1).getName()); + } + } + + if (inputField.getType() == FieldDescriptor.Type.ENUM) { + // For enums, in order to avoid value conflict, we will always define + // a enclosing struct called enum_full_name_E that includes the actual + // enum. + String enumFullName = inputField.getEnumType().getFullName(); + String enclosingTypeName = getNameFromFullName(enumFullName) + "_E"; + String enumName = inputField.getEnumType().getName(); + String actualEnumFullName = enclosingTypeName + "." + enumName; + if (enumTypes.contains(enumFullName)) { + resultField.setTypeName(actualEnumFullName); + } else { + EnumDescriptorProto enumType = inputField.getEnumType().toProto(); + resultProto.addNestedType( + DescriptorProto.newBuilder() + .setName(enclosingTypeName) + .addEnumType(enumType.toBuilder().setName(enumName)) + .build()); + resultField.setTypeName(actualEnumFullName); + enumTypes.add(enumFullName); + } + } + resultProto.addField(resultField); + } + structTypes.add(protoFullName); + + return ProtoSchema.newBuilder().setProtoDescriptor(resultProto.build()).build(); + } + + public static ProtoSchema convert(Descriptor descriptor) { + Set visitedTypes = new HashSet(); + Set enumTypes = new HashSet(); + Set structTypes = new HashSet(); + return convertInternal(descriptor, visitedTypes, enumTypes, structTypes, null); + } +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamConnection.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamConnection.java new file mode 100644 index 0000000000..20c5c9397d --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamConnection.java @@ -0,0 +1,104 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.gax.rpc.BidiStreamingCallable; +import com.google.api.gax.rpc.ClientStream; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.StreamController; +import io.grpc.Status; +import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; + +/** + * StreamConnection is responsible for writing requests to a GRPC bidirecional connection. + * + *

StreamWriter creates a connection. Two callback functions are necessary: request_callback and + * done_callback. Request callback is used for every request, and done callback is used to notify + * the user that the connection is closed and no more callbacks will be received from this + * connection. + * + *

The stream writer will accept all the requests without flow control, and makes the callbacks + * in receiving order. + * + *

It's user's responsibility to do the flow control and maintain the lifetime of the requests. + */ +public class StreamConnection { + private BidiStreamingCallable bidiStreamingCallable; + private ClientStream clientStream; + + private RequestCallback requestCallback; + private DoneCallback doneCallback; + + public StreamConnection( + BigQueryWriteClient client, RequestCallback requestCallback, DoneCallback doneCallback) { + this.requestCallback = requestCallback; + this.doneCallback = doneCallback; + + bidiStreamingCallable = client.appendRowsCallable(); + clientStream = + bidiStreamingCallable.splitCall( + new ResponseObserver() { + + @Override + public void onStart(StreamController controller) { + // no-op + } + + @Override + public void onResponse(AppendRowsResponse response) { + StreamConnection.this.requestCallback.run(response); + } + + @Override + public void onError(Throwable t) { + StreamConnection.this.doneCallback.run(t); + } + + @Override + public void onComplete() { + StreamConnection.this.doneCallback.run( + new StatusRuntimeException( + Status.fromCode(Code.CANCELLED) + .withDescription("Stream is closed by user."))); + } + }); + } + + /** + * Sends a request to the bi-directional stream connection. + * + * @param request request to send. + */ + public void send(AppendRowsRequest request) { + clientStream.send(request); + } + + /** Close the bi-directional stream connection. */ + public void close() { + clientStream.closeSend(); + } + + /** Invoked when a response is received from the server. */ + public static interface RequestCallback { + public void run(AppendRowsResponse response); + } + + /** Invoked when server closes the connection. */ + public static interface DoneCallback { + public void run(Throwable finalStatus); + } +} diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java new file mode 100644 index 0000000000..80a935ee93 --- /dev/null +++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java @@ -0,0 +1,621 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.core.CredentialsProvider; +import com.google.api.gax.rpc.FixedHeaderProvider; +import com.google.api.gax.rpc.TransportChannelProvider; +import com.google.cloud.bigquery.storage.v1.AppendRowsRequest.ProtoData; +import com.google.cloud.bigquery.storage.v1.StreamConnection.DoneCallback; +import com.google.cloud.bigquery.storage.v1.StreamConnection.RequestCallback; +import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.Uninterruptibles; +import com.google.protobuf.Int64Value; +import io.grpc.Status; +import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; +import java.io.IOException; +import java.util.Deque; +import java.util.LinkedList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; +import java.util.logging.Logger; +import javax.annotation.concurrent.GuardedBy; + +/** + * A BigQuery Stream Writer that can be used to write data into BigQuery Table. + * + *

TODO: Support batching. + * + *

TODO: Support schema change. + */ +public class StreamWriter implements AutoCloseable { + private static final Logger log = Logger.getLogger(StreamWriter.class.getName()); + + private Lock lock; + private Condition hasMessageInWaitingQueue; + private Condition inflightReduced; + + /* + * The identifier of stream to write to. + */ + private final String streamName; + + /* + * The proto schema of rows to write. + */ + private final ProtoSchema writerSchema; + + /* + * Max allowed inflight requests in the stream. Method append is blocked at this. + */ + private final long maxInflightRequests; + + /* + * Max allowed inflight bytes in the stream. Method append is blocked at this. + */ + private final long maxInflightBytes; + + /* + * TraceId for debugging purpose. + */ + private final String traceId; + + /* + * Tracks current inflight requests in the stream. + */ + @GuardedBy("lock") + private long inflightRequests = 0; + + /* + * Tracks current inflight bytes in the stream. + */ + @GuardedBy("lock") + private long inflightBytes = 0; + + /* + * Indicates whether user has called Close() or not. + */ + @GuardedBy("lock") + private boolean userClosed = false; + + /* + * The final status of connection. Set to nonnull when connection is permanently closed. + */ + @GuardedBy("lock") + private Throwable connectionFinalStatus = null; + + /* + * Contains requests buffered in the client and not yet sent to server. + */ + @GuardedBy("lock") + private final Deque waitingRequestQueue; + + /* + * Contains sent append requests waiting for response from server. + */ + @GuardedBy("lock") + private final Deque inflightRequestQueue; + + /* + * A client used to interact with BigQuery. + */ + private BigQueryWriteClient client; + + /* + * If true, the client above is created by this writer and should be closed. + */ + private boolean ownsBigQueryWriteClient = false; + + /* + * Wraps the underlying bi-directional stream connection with server. + */ + private StreamConnection streamConnection; + + /* + * A separate thread to handle actual communication with server. + */ + private Thread appendThread; + + /** The maximum size of one request. Defined by the API. */ + public static long getApiMaxRequestBytes() { + return 10L * 1000L * 1000L; // 10 megabytes (https://en.wikipedia.org/wiki/Megabyte) + } + + private StreamWriter(Builder builder) throws IOException { + this.lock = new ReentrantLock(); + this.hasMessageInWaitingQueue = lock.newCondition(); + this.inflightReduced = lock.newCondition(); + this.streamName = builder.streamName; + if (builder.writerSchema == null) { + throw new StatusRuntimeException( + Status.fromCode(Code.INVALID_ARGUMENT) + .withDescription("Writer schema must be provided when building this writer.")); + } + this.writerSchema = builder.writerSchema; + this.maxInflightRequests = builder.maxInflightRequest; + this.maxInflightBytes = builder.maxInflightBytes; + this.traceId = builder.traceId; + this.waitingRequestQueue = new LinkedList(); + this.inflightRequestQueue = new LinkedList(); + if (builder.client == null) { + BigQueryWriteSettings stubSettings = + BigQueryWriteSettings.newBuilder() + .setCredentialsProvider(builder.credentialsProvider) + .setTransportChannelProvider(builder.channelProvider) + .setEndpoint(builder.endpoint) + // (b/185842996): Temporily fix this by explicitly providing the header. + .setHeaderProvider( + FixedHeaderProvider.create( + "x-goog-request-params", "write_stream=" + this.streamName)) + .build(); + this.client = BigQueryWriteClient.create(stubSettings); + this.ownsBigQueryWriteClient = true; + } else { + this.client = builder.client; + this.ownsBigQueryWriteClient = false; + } + + this.streamConnection = + new StreamConnection( + this.client, + new RequestCallback() { + @Override + public void run(AppendRowsResponse response) { + requestCallback(response); + } + }, + new DoneCallback() { + @Override + public void run(Throwable finalStatus) { + doneCallback(finalStatus); + } + }); + this.appendThread = + new Thread( + new Runnable() { + @Override + public void run() { + appendLoop(); + } + }); + this.appendThread.start(); + } + + /** + * Schedules the writing of a message. + * + *

Example of writing a message. + * + * 

{@code
+   * AppendRowsRequest message;
+   * ApiFuture messageIdFuture = writer.append(message);
+   * ApiFutures.addCallback(messageIdFuture, new ApiFutureCallback() {
+   *   public void onSuccess(AppendRowsResponse response) {
+   *     if (!response.hasError()) {
+   *       System.out.println("written with offset: " + response.getAppendResult().getOffset());
+   *     } else {
+   *       System.out.println("received an in stream error: " + response.getError().toString());
+   *     }
+   *   }
+   *
+   *   public void onFailure(Throwable t) {
+   *     System.out.println("failed to write: " + t);
+   *   }
+   * }, MoreExecutors.directExecutor());
+   * }
+ * + * @param rows the rows in serialized format to write to BigQuery. + * @param offset the offset of the first row. + * @return the append response wrapped in a future. + */ + public ApiFuture append(ProtoRows rows, long offset) { + AppendRowsRequest.Builder requestBuilder = AppendRowsRequest.newBuilder(); + requestBuilder.setProtoRows(ProtoData.newBuilder().setRows(rows).build()); + if (offset >= 0) { + requestBuilder.setOffset(Int64Value.of(offset)); + } + return appendInternal(requestBuilder.build()); + } + + private ApiFuture appendInternal(AppendRowsRequest message) { + AppendRequestAndResponse requestWrapper = new AppendRequestAndResponse(message); + if (requestWrapper.messageSize > getApiMaxRequestBytes()) { + requestWrapper.appendResult.setException( + new StatusRuntimeException( + Status.fromCode(Code.INVALID_ARGUMENT) + .withDescription( + "MessageSize is too large. Max allow: " + + getApiMaxRequestBytes() + + " Actual: " + + requestWrapper.messageSize))); + return requestWrapper.appendResult; + } + this.lock.lock(); + try { + if (userClosed) { + requestWrapper.appendResult.setException( + new StatusRuntimeException( + Status.fromCode(Status.Code.FAILED_PRECONDITION) + .withDescription("Stream is already closed"))); + return requestWrapper.appendResult; + } + if (connectionFinalStatus != null) { + requestWrapper.appendResult.setException( + new StatusRuntimeException( + Status.fromCode(Status.Code.FAILED_PRECONDITION) + .withDescription( + "Stream is closed due to " + connectionFinalStatus.toString()))); + return requestWrapper.appendResult; + } + + ++this.inflightRequests; + this.inflightBytes += requestWrapper.messageSize; + waitingRequestQueue.addLast(requestWrapper); + hasMessageInWaitingQueue.signal(); + maybeWaitForInflightQuota(); + return requestWrapper.appendResult; + } finally { + this.lock.unlock(); + } + } + + @GuardedBy("lock") + private void maybeWaitForInflightQuota() { + while (this.inflightRequests >= this.maxInflightRequests + || this.inflightBytes >= this.maxInflightBytes) { + try { + inflightReduced.await(100, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + log.warning( + "Interrupted while waiting for inflight quota. Stream: " + + streamName + + " Error: " + + e.toString()); + throw new StatusRuntimeException( + Status.fromCode(Code.CANCELLED) + .withCause(e) + .withDescription("Interrupted while waiting for quota.")); + } + } + } + + /** Close the stream writer. Shut down all resources. */ + @Override + public void close() { + log.info("User closing stream: " + streamName); + this.lock.lock(); + try { + this.userClosed = true; + } finally { + this.lock.unlock(); + } + log.info("Waiting for append thread to finish. Stream: " + streamName); + try { + appendThread.join(); + log.info("User close complete. Stream: " + streamName); + } catch (InterruptedException e) { + // Unexpected. Just swallow the exception with logging. + log.warning( + "Append handler join is interrupted. Stream: " + streamName + " Error: " + e.toString()); + } + if (this.ownsBigQueryWriteClient) { + this.client.close(); + try { + this.client.awaitTermination(1, TimeUnit.MINUTES); + } catch (InterruptedException ignored) { + } + } + } + + /* + * This loop is executed in a separate thread. + * + * It takes requests from waiting queue and sends them to server. + */ + private void appendLoop() { + boolean isFirstRequestInConnection = true; + Deque localQueue = new LinkedList(); + while (!waitingQueueDrained()) { + this.lock.lock(); + try { + hasMessageInWaitingQueue.await(100, TimeUnit.MILLISECONDS); + while (!this.waitingRequestQueue.isEmpty()) { + AppendRequestAndResponse requestWrapper = this.waitingRequestQueue.pollFirst(); + this.inflightRequestQueue.addLast(requestWrapper); + localQueue.addLast(requestWrapper); + } + } catch (InterruptedException e) { + log.warning( + "Interrupted while waiting for message. Stream: " + + streamName + + " Error: " + + e.toString()); + } finally { + this.lock.unlock(); + } + + if (localQueue.isEmpty()) { + continue; + } + + // TODO: Add reconnection here. + while (!localQueue.isEmpty()) { + AppendRowsRequest preparedRequest = + prepareRequestBasedOnPosition( + localQueue.pollFirst().message, isFirstRequestInConnection); + this.streamConnection.send(preparedRequest); + isFirstRequestInConnection = false; + } + } + + log.info("Cleanup starts. Stream: " + streamName); + // At this point, the waiting queue is drained, so no more requests. + // We can close the stream connection and handle the remaining inflight requests. + this.streamConnection.close(); + waitForDoneCallback(); + + // At this point, there cannot be more callback. It is safe to clean up all inflight requests. + log.info( + "Stream connection is fully closed. Cleaning up inflight requests. Stream: " + streamName); + cleanupInflightRequests(); + log.info("Append thread is done. Stream: " + streamName); + } + + /* + * Returns true if waiting queue is drain, a.k.a. no more requests in the waiting queue. + * + * It serves as a signal to append thread that there cannot be any more requests in the waiting + * queue and it can prepare to stop. + */ + private boolean waitingQueueDrained() { + this.lock.lock(); + try { + return (this.userClosed || this.connectionFinalStatus != null) + && this.waitingRequestQueue.isEmpty(); + } finally { + this.lock.unlock(); + } + } + + private void waitForDoneCallback() { + log.info("Waiting for done callback from stream connection. Stream: " + streamName); + while (true) { + this.lock.lock(); + try { + if (connectionFinalStatus != null) { + // Done callback is received, return. + return; + } + } finally { + this.lock.unlock(); + } + Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS); + } + } + + private AppendRowsRequest prepareRequestBasedOnPosition( + AppendRowsRequest original, boolean isFirstRequest) { + AppendRowsRequest.Builder requestBuilder = original.toBuilder(); + if (isFirstRequest) { + if (this.writerSchema != null) { + requestBuilder.getProtoRowsBuilder().setWriterSchema(this.writerSchema); + } + requestBuilder.setWriteStream(this.streamName); + if (this.traceId != null) { + requestBuilder.setTraceId(this.traceId); + } + } else { + requestBuilder.clearWriteStream(); + requestBuilder.getProtoRowsBuilder().clearWriterSchema(); + } + return requestBuilder.build(); + } + + private void cleanupInflightRequests() { + Throwable finalStatus; + Deque localQueue = new LinkedList(); + this.lock.lock(); + try { + finalStatus = this.connectionFinalStatus; + while (!this.inflightRequestQueue.isEmpty()) { + localQueue.addLast(pollInflightRequestQueue()); + } + } finally { + this.lock.unlock(); + } + log.info( + "Cleaning " + + localQueue.size() + + " inflight requests with error: " + + finalStatus.toString()); + while (!localQueue.isEmpty()) { + localQueue.pollFirst().appendResult.setException(finalStatus); + } + } + + private void requestCallback(AppendRowsResponse response) { + AppendRequestAndResponse requestWrapper; + this.lock.lock(); + try { + requestWrapper = pollInflightRequestQueue(); + } finally { + this.lock.unlock(); + } + if (response.hasError()) { + StatusRuntimeException exception = + new StatusRuntimeException( + Status.fromCodeValue(response.getError().getCode()) + .withDescription(response.getError().getMessage())); + requestWrapper.appendResult.setException(exception); + } else { + requestWrapper.appendResult.set(response); + } + } + + private void doneCallback(Throwable finalStatus) { + log.info( + "Received done callback. Stream: " + + streamName + + " Final status: " + + finalStatus.toString()); + this.lock.lock(); + try { + this.connectionFinalStatus = finalStatus; + } finally { + this.lock.unlock(); + } + } + + @GuardedBy("lock") + private AppendRequestAndResponse pollInflightRequestQueue() { + AppendRequestAndResponse requestWrapper = this.inflightRequestQueue.pollFirst(); + --this.inflightRequests; + this.inflightBytes -= requestWrapper.messageSize; + this.inflightReduced.signal(); + return requestWrapper; + } + + /** + * Constructs a new {@link StreamWriterV2.Builder} using the given stream and client. AppendRows + * needs special headers to be added to client, so a passed in client will not work. This should + * be used by test only. + */ + public static StreamWriter.Builder newBuilder(String streamName, BigQueryWriteClient client) { + return new StreamWriter.Builder(streamName, client); + } + + /** Constructs a new {@link StreamWriterV2.Builder} using the given stream. */ + public static StreamWriter.Builder newBuilder(String streamName) { + return new StreamWriter.Builder(streamName); + } + + /** A builder of {@link StreamWriterV2}s. */ + public static final class Builder { + + private static final long DEFAULT_MAX_INFLIGHT_REQUESTS = 1000L; + + private static final long DEFAULT_MAX_INFLIGHT_BYTES = 100 * 1024 * 1024; // 100Mb. + + private String streamName; + + private BigQueryWriteClient client; + + private ProtoSchema writerSchema = null; + + private long maxInflightRequest = DEFAULT_MAX_INFLIGHT_REQUESTS; + + private long maxInflightBytes = DEFAULT_MAX_INFLIGHT_BYTES; + + private String endpoint = BigQueryWriteSettings.getDefaultEndpoint(); + + private TransportChannelProvider channelProvider = + BigQueryWriteSettings.defaultGrpcTransportProviderBuilder().setChannelsPerCpu(1).build(); + + private CredentialsProvider credentialsProvider = + BigQueryWriteSettings.defaultCredentialsProviderBuilder().build(); + + private String traceId = null; + + private Builder(String streamName) { + this.streamName = Preconditions.checkNotNull(streamName); + this.client = null; + } + + private Builder(String streamName, BigQueryWriteClient client) { + this.streamName = Preconditions.checkNotNull(streamName); + this.client = Preconditions.checkNotNull(client); + } + + /** Sets the proto schema of the rows. */ + public Builder setWriterSchema(ProtoSchema writerSchema) { + this.writerSchema = writerSchema; + return this; + } + + public Builder setMaxInflightRequests(long value) { + this.maxInflightRequest = value; + return this; + } + + public Builder setMaxInflightBytes(long value) { + this.maxInflightBytes = value; + return this; + } + + /** Gives the ability to override the gRPC endpoint. */ + public Builder setEndpoint(String endpoint) { + this.endpoint = Preconditions.checkNotNull(endpoint, "Endpoint is null."); + return this; + } + + /** + * {@code ChannelProvider} to use to create Channels, which must point at Cloud BigQuery Storage + * API endpoint. + * + *

For performance, this client benefits from having multiple underlying connections. See + * {@link com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.Builder#setPoolSize(int)}. + */ + public Builder setChannelProvider(TransportChannelProvider channelProvider) { + this.channelProvider = + Preconditions.checkNotNull(channelProvider, "ChannelProvider is null."); + return this; + } + + /** {@code CredentialsProvider} to use to create Credentials to authenticate calls. */ + public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) { + this.credentialsProvider = + Preconditions.checkNotNull(credentialsProvider, "CredentialsProvider is null."); + return this; + } + + /** + * Sets traceId for debuging purpose. TraceId must follow the format of + * CustomerDomain:DebugString, e.g. DATAFLOW:job_id_x. + */ + public Builder setTraceId(String traceId) { + int colonIndex = traceId.indexOf(':'); + if (colonIndex == -1 || colonIndex == 0 || colonIndex == traceId.length() - 1) { + throw new IllegalArgumentException( + "TraceId must follow the format of A:B. Actual:" + traceId); + } + this.traceId = traceId; + return this; + } + + /** Builds the {@code StreamWriterV2}. */ + public StreamWriter build() throws IOException { + return new StreamWriter(this); + } + } + + // Class that wraps AppendRowsRequest and its corresponding Response future. + private static final class AppendRequestAndResponse { + final SettableApiFuture appendResult; + final AppendRowsRequest message; + final long messageSize; + + AppendRequestAndResponse(AppendRowsRequest message) { + this.appendResult = SettableApiFuture.create(); + this.message = message; + this.messageSize = message.getProtoRows().getSerializedSize(); + } + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/BQTableSchemaToProtoDescriptorTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/BQTableSchemaToProtoDescriptorTest.java new file mode 100644 index 0000000000..07cb1c8657 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/BQTableSchemaToProtoDescriptorTest.java @@ -0,0 +1,400 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import static org.junit.Assert.*; + +import com.google.cloud.bigquery.storage.test.JsonTest.*; +import com.google.cloud.bigquery.storage.test.SchemaTest.*; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Descriptors.FieldDescriptor; +import java.util.HashMap; +import java.util.Map; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class BQTableSchemaToProtoDescriptorTest { + // This is a map between the TableFieldSchema.Type and the descriptor it is supposed to + // produce. The produced descriptor will be used to check against the entry values here. + private static ImmutableMap + BQTableTypeToCorrectProtoDescriptorTest = + new ImmutableMap.Builder() + .put(TableFieldSchema.Type.BOOL, BoolType.getDescriptor()) + .put(TableFieldSchema.Type.BYTES, BytesType.getDescriptor()) + .put(TableFieldSchema.Type.DATE, Int32Type.getDescriptor()) + .put(TableFieldSchema.Type.DATETIME, Int64Type.getDescriptor()) + .put(TableFieldSchema.Type.DOUBLE, DoubleType.getDescriptor()) + .put(TableFieldSchema.Type.GEOGRAPHY, StringType.getDescriptor()) + .put(TableFieldSchema.Type.INT64, Int64Type.getDescriptor()) + .put(TableFieldSchema.Type.NUMERIC, BytesType.getDescriptor()) + .put(TableFieldSchema.Type.STRING, StringType.getDescriptor()) + .put(TableFieldSchema.Type.TIME, Int64Type.getDescriptor()) + .put(TableFieldSchema.Type.TIMESTAMP, Int64Type.getDescriptor()) + .build(); + + // Creates mapping from descriptor to how many times it was reused. + private void mapDescriptorToCount(Descriptor descriptor, HashMap map) { + for (FieldDescriptor field : descriptor.getFields()) { + if (field.getType() == FieldDescriptor.Type.MESSAGE) { + Descriptor subDescriptor = field.getMessageType(); + String messageName = subDescriptor.getName(); + if (map.containsKey(messageName)) { + map.put(messageName, map.get(messageName) + 1); + } else { + map.put(messageName, 1); + } + mapDescriptorToCount(subDescriptor, map); + } + } + } + + private void isDescriptorEqual(Descriptor convertedProto, Descriptor originalProto) { + // Check same number of fields + assertEquals(convertedProto.getFields().size(), originalProto.getFields().size()); + for (FieldDescriptor convertedField : convertedProto.getFields()) { + // Check field name + FieldDescriptor originalField = originalProto.findFieldByName(convertedField.getName()); + assertNotNull(originalField); + // Check type + FieldDescriptor.Type convertedType = convertedField.getType(); + FieldDescriptor.Type originalType = originalField.getType(); + assertEquals(convertedField.getName(), convertedType, originalType); + // Check mode + assertTrue( + (originalField.isRepeated() == convertedField.isRepeated()) + && (originalField.isRequired() == convertedField.isRequired()) + && (originalField.isOptional() == convertedField.isOptional())); + // Recursively check nested messages + if (convertedType == FieldDescriptor.Type.MESSAGE) { + isDescriptorEqual(convertedField.getMessageType(), originalField.getMessageType()); + } + } + } + + @Test + public void testSimpleTypes() throws Exception { + for (Map.Entry entry : + BQTableTypeToCorrectProtoDescriptorTest.entrySet()) { + final TableFieldSchema tableFieldSchema = + TableFieldSchema.newBuilder() + .setType(entry.getKey()) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_field_type") + .build(); + final TableSchema tableSchema = + TableSchema.newBuilder().addFields(0, tableFieldSchema).build(); + final Descriptor descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(tableSchema); + isDescriptorEqual(descriptor, entry.getValue()); + } + } + + @Test + public void testStructSimple() throws Exception { + final TableFieldSchema StringType = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_field_type") + .build(); + final TableFieldSchema tableFieldSchema = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_field_type") + .addFields(0, StringType) + .build(); + final TableSchema tableSchema = TableSchema.newBuilder().addFields(0, tableFieldSchema).build(); + final Descriptor descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(tableSchema); + isDescriptorEqual(descriptor, MessageType.getDescriptor()); + } + + @Test + public void testStructComplex() throws Exception { + final TableFieldSchema test_int = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_int") + .build(); + final TableFieldSchema test_string = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_string") + .build(); + final TableFieldSchema test_bytes = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.BYTES) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("test_bytes") + .build(); + final TableFieldSchema test_bool = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.BOOL) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_bool") + .build(); + final TableFieldSchema test_double = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DOUBLE) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_double") + .build(); + final TableFieldSchema test_date = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DATE) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("test_date") + .build(); + final TableFieldSchema ComplexLvl2 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.REQUIRED) + .addFields(0, test_int) + .setName("complex_lvl2") + .build(); + final TableFieldSchema ComplexLvl1 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.REQUIRED) + .addFields(0, test_int) + .addFields(1, ComplexLvl2) + .setName("complex_lvl1") + .build(); + final TableFieldSchema TEST_NUMERIC = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_numeric") + .build(); + final TableFieldSchema TEST_GEO = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.GEOGRAPHY) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_geo") + .build(); + final TableFieldSchema TEST_TIMESTAMP = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.TIMESTAMP) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_timestamp") + .build(); + final TableFieldSchema TEST_TIME = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.TIME) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_time") + .build(); + final TableFieldSchema TEST_NUMERIC_REPEATED = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_numeric_repeated") + .build(); + final TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, test_int) + .addFields(1, test_string) + .addFields(2, test_bytes) + .addFields(3, test_bool) + .addFields(4, test_double) + .addFields(5, test_date) + .addFields(6, ComplexLvl1) + .addFields(7, ComplexLvl2) + .addFields(8, TEST_NUMERIC) + .addFields(9, TEST_GEO) + .addFields(10, TEST_TIMESTAMP) + .addFields(11, TEST_TIME) + .addFields(12, TEST_NUMERIC_REPEATED) + .build(); + final Descriptor descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(tableSchema); + isDescriptorEqual(descriptor, ComplexRoot.getDescriptor()); + } + + @Test + public void testCasingComplexStruct() throws Exception { + final TableFieldSchema required = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("tEsT_ReQuIrEd") + .build(); + final TableFieldSchema repeated = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("tESt_repEATed") + .build(); + final TableFieldSchema optional = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_opTIONal") + .build(); + final TableFieldSchema test_int = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("TEST_INT") + .build(); + final TableFieldSchema test_string = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("TEST_STRING") + .build(); + final TableFieldSchema test_bytes = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.BYTES) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("TEST_BYTES") + .build(); + final TableFieldSchema test_bool = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.BOOL) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("TEST_BOOL") + .build(); + final TableFieldSchema test_double = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DOUBLE) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("TEST_DOUBLE") + .build(); + final TableFieldSchema test_date = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DATE) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("TEST_DATE") + .build(); + final TableFieldSchema option_test = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.REQUIRED) + .addFields(0, required) + .addFields(1, repeated) + .addFields(2, optional) + .setName("option_test") + .build(); + final TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, test_int) + .addFields(1, test_string) + .addFields(2, test_bytes) + .addFields(3, test_bool) + .addFields(4, test_double) + .addFields(5, test_date) + .addFields(6, option_test) + .build(); + final Descriptor descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(tableSchema); + isDescriptorEqual(descriptor, CasingComplex.getDescriptor()); + } + + @Test + public void testOptions() throws Exception { + final TableFieldSchema required = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("test_required") + .build(); + final TableFieldSchema repeated = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_repeated") + .build(); + final TableFieldSchema optional = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_optional") + .build(); + final TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, required) + .addFields(1, repeated) + .addFields(2, optional) + .build(); + final Descriptor descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(tableSchema); + isDescriptorEqual(descriptor, OptionTest.getDescriptor()); + } + + @Test + public void testDescriptorReuseDuringCreation() throws Exception { + final TableFieldSchema test_int = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_int") + .build(); + final TableFieldSchema reuse_lvl2 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("reuse_lvl2") + .addFields(0, test_int) + .build(); + final TableFieldSchema reuse_lvl1 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("reuse_lvl1") + .addFields(0, test_int) + .addFields(0, reuse_lvl2) + .build(); + final TableFieldSchema reuse_lvl1_1 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("reuse_lvl1_1") + .addFields(0, test_int) + .addFields(0, reuse_lvl2) + .build(); + final TableFieldSchema reuse_lvl1_2 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("reuse_lvl1_2") + .addFields(0, test_int) + .addFields(0, reuse_lvl2) + .build(); + final TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, reuse_lvl1) + .addFields(1, reuse_lvl1_1) + .addFields(2, reuse_lvl1_2) + .build(); + final Descriptor descriptor = + BQTableSchemaToProtoDescriptor.convertBQTableSchemaToProtoDescriptor(tableSchema); + HashMap descriptorToCount = new HashMap(); + mapDescriptorToCount(descriptor, descriptorToCount); + assertEquals(descriptorToCount.size(), 2); + assertTrue(descriptorToCount.containsKey("root__reuse_lvl1")); + assertEquals(descriptorToCount.get("root__reuse_lvl1").intValue(), 3); + assertTrue(descriptorToCount.containsKey("root__reuse_lvl1__reuse_lvl2")); + assertEquals(descriptorToCount.get("root__reuse_lvl1__reuse_lvl2").intValue(), 3); + isDescriptorEqual(descriptor, ReuseRoot.getDescriptor()); + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeBigQueryWrite.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeBigQueryWrite.java new file mode 100644 index 0000000000..8a6a8d3d98 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeBigQueryWrite.java @@ -0,0 +1,93 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.gax.grpc.testing.MockGrpcService; +import com.google.protobuf.AbstractMessage; +import io.grpc.ServerServiceDefinition; +import java.util.LinkedList; +import java.util.List; +import java.util.concurrent.ScheduledExecutorService; +import org.threeten.bp.Duration; + +/** + * A fake implementation of {@link MockGrpcService}, that can be used to test clients of a + * StreamWriter. It forwards calls to the real implementation (@link FakeBigQueryWriteImpl}. + */ +public class FakeBigQueryWrite implements MockGrpcService { + private final FakeBigQueryWriteImpl serviceImpl; + + public FakeBigQueryWrite() { + serviceImpl = new FakeBigQueryWriteImpl(); + } + + @Override + public List getRequests() { + return new LinkedList(serviceImpl.getCapturedRequests()); + } + + public void waitForResponseScheduled() throws InterruptedException { + serviceImpl.waitForResponseScheduled(); + } + + public List getAppendRequests() { + return serviceImpl.getCapturedRequests(); + } + + public List getWriteStreamRequests() { + return serviceImpl.getCapturedWriteRequests(); + } + + @Override + public void addResponse(AbstractMessage response) { + if (response instanceof AppendRowsResponse) { + serviceImpl.addResponse((AppendRowsResponse) response); + } else if (response instanceof WriteStream) { + serviceImpl.addWriteStreamResponse((WriteStream) response); + } else if (response instanceof FlushRowsResponse) { + serviceImpl.addFlushRowsResponse((FlushRowsResponse) response); + } else { + throw new IllegalStateException("Unsupported service"); + } + } + + @Override + public void addException(Exception exception) { + serviceImpl.addConnectionError(exception); + } + + @Override + public ServerServiceDefinition getServiceDefinition() { + return serviceImpl.bindService(); + } + + @Override + public void reset() { + serviceImpl.reset(); + } + + public void setResponseDelay(Duration delay) { + serviceImpl.setResponseDelay(delay); + } + + public void setResponseSleep(Duration sleep) { + serviceImpl.setResponseSleep(sleep); + } + + public void setExecutor(ScheduledExecutorService executor) { + serviceImpl.setExecutor(executor); + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeBigQueryWriteImpl.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeBigQueryWriteImpl.java new file mode 100644 index 0000000000..ce59e02663 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeBigQueryWriteImpl.java @@ -0,0 +1,234 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.common.base.Optional; +import com.google.common.util.concurrent.Uninterruptibles; +import io.grpc.stub.StreamObserver; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.logging.Logger; +import org.threeten.bp.Duration; + +/** + * A fake implementation of {@link BigQueryWriteImplBase} that can acts like server in StreamWriter + * unit testing. + */ +class FakeBigQueryWriteImpl extends BigQueryWriteGrpc.BigQueryWriteImplBase { + private static final Logger LOG = Logger.getLogger(FakeBigQueryWriteImpl.class.getName()); + + private final LinkedBlockingQueue requests = new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue writeRequests = + new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue flushRequests = new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue responses = new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue writeResponses = new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue flushResponses = new LinkedBlockingQueue<>(); + private final AtomicInteger nextMessageId = new AtomicInteger(1); + private boolean autoPublishResponse; + private ScheduledExecutorService executor = null; + private Duration responseDelay = Duration.ZERO; + + private Duration responseSleep = Duration.ZERO; + private Semaphore responseSemaphore = new Semaphore(0, true); + + /** Class used to save the state of a possible response. */ + private static class Response { + Optional appendResponse; + Optional error; + + public Response(AppendRowsResponse appendResponse) { + this.appendResponse = Optional.of(appendResponse); + this.error = Optional.absent(); + } + + public Response(Throwable exception) { + this.appendResponse = Optional.absent(); + this.error = Optional.of(exception); + } + + public AppendRowsResponse getResponse() { + return appendResponse.get(); + } + + public Throwable getError() { + return error.get(); + } + + boolean isError() { + return error.isPresent(); + } + + @Override + public String toString() { + if (isError()) { + return error.get().toString(); + } + return appendResponse.get().toString(); + } + } + + @Override + public void getWriteStream( + GetWriteStreamRequest request, StreamObserver responseObserver) { + Object response = writeResponses.remove(); + if (response instanceof WriteStream) { + writeRequests.add(request); + responseObserver.onNext((WriteStream) response); + responseObserver.onCompleted(); + } else if (response instanceof Exception) { + responseObserver.onError((Exception) response); + } else { + responseObserver.onError(new IllegalArgumentException("Unrecognized response type")); + } + } + + @Override + public void flushRows( + FlushRowsRequest request, StreamObserver responseObserver) { + Object response = writeResponses.remove(); + if (response instanceof FlushRowsResponse) { + flushRequests.add(request); + responseObserver.onNext((FlushRowsResponse) response); + responseObserver.onCompleted(); + } else if (response instanceof Exception) { + responseObserver.onError((Exception) response); + } else { + responseObserver.onError(new IllegalArgumentException("Unrecognized response type")); + } + } + + public void waitForResponseScheduled() throws InterruptedException { + responseSemaphore.acquire(); + } + + @Override + public StreamObserver appendRows( + final StreamObserver responseObserver) { + StreamObserver requestObserver = + new StreamObserver() { + @Override + public void onNext(AppendRowsRequest value) { + LOG.fine("Get request:" + value.toString()); + final Response response = responses.remove(); + requests.add(value); + if (responseSleep.compareTo(Duration.ZERO) > 0) { + LOG.info("Sleeping before response for " + responseSleep.toString()); + Uninterruptibles.sleepUninterruptibly( + responseSleep.toMillis(), TimeUnit.MILLISECONDS); + } + if (responseDelay == Duration.ZERO) { + sendResponse(response, responseObserver); + } else { + final Response responseToSend = response; + // TODO(yirutang): This is very wrong because it messes up response/complete ordering. + LOG.fine("Schedule a response to be sent at delay"); + executor.schedule( + new Runnable() { + @Override + public void run() { + sendResponse(responseToSend, responseObserver); + } + }, + responseDelay.toMillis(), + TimeUnit.MILLISECONDS); + } + responseSemaphore.release(); + } + + @Override + public void onError(Throwable t) { + responseObserver.onError(t); + } + + @Override + public void onCompleted() { + responseObserver.onCompleted(); + } + }; + return requestObserver; + } + + private void sendResponse( + Response response, StreamObserver responseObserver) { + LOG.fine("Sending response: " + response.toString()); + if (response.isError()) { + responseObserver.onError(response.getError()); + } else { + responseObserver.onNext(response.getResponse()); + } + } + + /** Set an executor to use to delay publish responses. */ + public FakeBigQueryWriteImpl setExecutor(ScheduledExecutorService executor) { + this.executor = executor; + return this; + } + + /** Set an amount of time by which to delay publish responses. */ + public FakeBigQueryWriteImpl setResponseDelay(Duration responseDelay) { + this.responseDelay = responseDelay; + return this; + } + + /** Set an amount of time by which to sleep before publishing responses. */ + public FakeBigQueryWriteImpl setResponseSleep(Duration responseSleep) { + this.responseSleep = responseSleep; + return this; + } + + public FakeBigQueryWriteImpl addResponse(AppendRowsResponse appendRowsResponse) { + responses.add(new Response(appendRowsResponse)); + return this; + } + + public FakeBigQueryWriteImpl addResponse(AppendRowsResponse.Builder appendResponseBuilder) { + return addResponse(appendResponseBuilder.build()); + } + + public FakeBigQueryWriteImpl addWriteStreamResponse(WriteStream response) { + writeResponses.add(response); + return this; + } + + public FakeBigQueryWriteImpl addFlushRowsResponse(FlushRowsResponse response) { + flushResponses.add(response); + return this; + } + + public FakeBigQueryWriteImpl addConnectionError(Throwable error) { + responses.add(new Response(error)); + return this; + } + + public List getCapturedRequests() { + return new ArrayList(requests); + } + + public List getCapturedWriteRequests() { + return new ArrayList(writeRequests); + } + + public void reset() { + requests.clear(); + responses.clear(); + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeClock.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeClock.java new file mode 100644 index 0000000000..6a83c820c7 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeClock.java @@ -0,0 +1,41 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.core.ApiClock; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +/** A Clock to help with testing time-based logic. */ +public class FakeClock implements ApiClock { + + private final AtomicLong millis = new AtomicLong(); + + // Advances the clock value by {@code time} in {@code timeUnit}. + public void advance(long time, TimeUnit timeUnit) { + millis.addAndGet(timeUnit.toMillis(time)); + } + + @Override + public long nanoTime() { + return millisTime() * 1000_000L; + } + + @Override + public long millisTime() { + return millis.get(); + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeScheduledExecutorService.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeScheduledExecutorService.java new file mode 100644 index 0000000000..0869fdc788 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/FakeScheduledExecutorService.java @@ -0,0 +1,346 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.core.ApiClock; +import com.google.common.primitives.Ints; +import com.google.common.util.concurrent.SettableFuture; +import java.util.ArrayList; +import java.util.Deque; +import java.util.LinkedList; +import java.util.List; +import java.util.PriorityQueue; +import java.util.concurrent.AbstractExecutorService; +import java.util.concurrent.Callable; +import java.util.concurrent.Delayed; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.logging.Logger; +import org.threeten.bp.Duration; +import org.threeten.bp.Instant; + +/** + * Fake implementation of {@link ScheduledExecutorService} that allows tests control the reference + * time of the executor and decide when to execute any outstanding task. + */ +public class FakeScheduledExecutorService extends AbstractExecutorService + implements ScheduledExecutorService { + private static final Logger LOG = Logger.getLogger(FakeScheduledExecutorService.class.getName()); + + private final AtomicBoolean shutdown = new AtomicBoolean(false); + private final PriorityQueue> pendingCallables = new PriorityQueue<>(); + private final FakeClock clock = new FakeClock(); + private final Deque expectedWorkQueue = new LinkedList<>(); + + public ApiClock getClock() { + return clock; + } + + @Override + public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { + return schedulePendingCallable( + new PendingCallable<>( + Duration.ofMillis(unit.toMillis(delay)), command, PendingCallableType.NORMAL)); + } + + @Override + public ScheduledFuture schedule(Callable callable, long delay, TimeUnit unit) { + return schedulePendingCallable( + new PendingCallable<>( + Duration.ofMillis(unit.toMillis(delay)), callable, PendingCallableType.NORMAL)); + } + + @Override + public ScheduledFuture scheduleAtFixedRate( + Runnable command, long initialDelay, long period, TimeUnit unit) { + return schedulePendingCallable( + new PendingCallable<>( + Duration.ofMillis(unit.toMillis(initialDelay)), + command, + PendingCallableType.FIXED_RATE)); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay( + Runnable command, long initialDelay, long delay, TimeUnit unit) { + return schedulePendingCallable( + new PendingCallable<>( + Duration.ofMillis(unit.toMillis(initialDelay)), + command, + PendingCallableType.FIXED_DELAY)); + } + + /** + * This will advance the reference time of the executor and execute (in the same thread) any + * outstanding callable which execution time has passed. + */ + public void advanceTime(Duration toAdvance) { + LOG.info( + "Advance to time to:" + + Instant.ofEpochMilli(clock.millisTime() + toAdvance.toMillis()).toString()); + clock.advance(toAdvance.toMillis(), TimeUnit.MILLISECONDS); + work(); + } + + private void work() { + for (; ; ) { + PendingCallable callable = null; + Instant cmpTime = Instant.ofEpochMilli(clock.millisTime()); + if (!pendingCallables.isEmpty()) { + LOG.info( + "Going to call: Current time: " + + cmpTime.toString() + + " Scheduled time: " + + pendingCallables.peek().getScheduledTime().toString() + + " Creation time:" + + pendingCallables.peek().getCreationTime().toString()); + } + synchronized (pendingCallables) { + if (pendingCallables.isEmpty() + || pendingCallables.peek().getScheduledTime().isAfter(cmpTime)) { + break; + } + callable = pendingCallables.poll(); + } + if (callable != null) { + try { + callable.call(); + } catch (Exception e) { + // We ignore any callable exception, which should be set to the future but not relevant to + // advanceTime. + } + } + } + + synchronized (pendingCallables) { + if (shutdown.get() && pendingCallables.isEmpty()) { + pendingCallables.notifyAll(); + } + } + } + + @Override + public void shutdown() { + if (shutdown.getAndSet(true)) { + throw new IllegalStateException("This executor has been shutdown already"); + } + } + + @Override + public List shutdownNow() { + if (shutdown.getAndSet(true)) { + throw new IllegalStateException("This executor has been shutdown already"); + } + List pending = new ArrayList<>(); + for (final PendingCallable pendingCallable : pendingCallables) { + pending.add( + new Runnable() { + @Override + public void run() { + pendingCallable.call(); + } + }); + } + synchronized (pendingCallables) { + pendingCallables.notifyAll(); + pendingCallables.clear(); + } + return pending; + } + + @Override + public boolean isShutdown() { + return shutdown.get(); + } + + @Override + public boolean isTerminated() { + return pendingCallables.isEmpty(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + synchronized (pendingCallables) { + if (pendingCallables.isEmpty()) { + return true; + } + LOG.info("Wating on pending callables" + pendingCallables.size()); + pendingCallables.wait(unit.toMillis(timeout)); + return pendingCallables.isEmpty(); + } + } + + @Override + public void execute(Runnable command) { + if (shutdown.get()) { + throw new IllegalStateException("This executor has been shutdown"); + } + command.run(); + } + + ScheduledFuture schedulePendingCallable(PendingCallable callable) { + LOG.info( + "Schedule pending callable called " + callable.delay + " " + callable.getScheduledTime()); + if (shutdown.get()) { + throw new IllegalStateException("This executor has been shutdown"); + } + synchronized (pendingCallables) { + pendingCallables.add(callable); + } + work(); + synchronized (expectedWorkQueue) { + // We compare by the callable delay in order decide when to remove expectations from the + // expected work queue, i.e. only the expected work that matches the delay of the scheduled + // callable is removed from the queue. + if (!expectedWorkQueue.isEmpty() && expectedWorkQueue.peek().equals(callable.delay)) { + expectedWorkQueue.poll(); + } + expectedWorkQueue.notifyAll(); + } + + return callable.getScheduledFuture(); + } + + enum PendingCallableType { + NORMAL, + FIXED_RATE, + FIXED_DELAY + } + + /** Class that saves the state of an scheduled pending callable. */ + class PendingCallable implements Comparable> { + Instant creationTime = Instant.ofEpochMilli(clock.millisTime()); + Duration delay; + Callable pendingCallable; + SettableFuture future = SettableFuture.create(); + AtomicBoolean cancelled = new AtomicBoolean(false); + AtomicBoolean done = new AtomicBoolean(false); + PendingCallableType type; + + PendingCallable(Duration delay, final Runnable runnable, PendingCallableType type) { + pendingCallable = + new Callable() { + @Override + public T call() { + runnable.run(); + return null; + } + }; + this.type = type; + this.delay = delay; + } + + PendingCallable(Duration delay, Callable callable, PendingCallableType type) { + pendingCallable = callable; + this.type = type; + this.delay = delay; + } + + private Instant getScheduledTime() { + return creationTime.plus(delay); + } + + private Instant getCreationTime() { + return creationTime; + } + + ScheduledFuture getScheduledFuture() { + return new ScheduledFuture() { + @Override + public long getDelay(TimeUnit unit) { + return unit.convert( + getScheduledTime().toEpochMilli() - clock.millisTime(), TimeUnit.MILLISECONDS); + } + + @Override + public int compareTo(Delayed o) { + return Ints.saturatedCast( + getDelay(TimeUnit.MILLISECONDS) - o.getDelay(TimeUnit.MILLISECONDS)); + } + + @Override + public boolean cancel(boolean mayInterruptIfRunning) { + synchronized (this) { + cancelled.set(true); + return !done.get(); + } + } + + @Override + public boolean isCancelled() { + return cancelled.get(); + } + + @Override + public boolean isDone() { + return done.get(); + } + + @Override + public T get() throws InterruptedException, ExecutionException { + return future.get(); + } + + @Override + public T get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + return future.get(timeout, unit); + } + }; + } + + T call() { + T result = null; + synchronized (this) { + if (cancelled.get()) { + return null; + } + try { + result = pendingCallable.call(); + future.set(result); + } catch (Exception e) { + future.setException(e); + } finally { + switch (type) { + case NORMAL: + done.set(true); + break; + case FIXED_DELAY: + this.creationTime = Instant.ofEpochMilli(clock.millisTime()); + schedulePendingCallable(this); + break; + case FIXED_RATE: + this.creationTime = this.creationTime.plus(delay); + schedulePendingCallable(this); + break; + default: + // Nothing to do + } + } + } + return result; + } + + @Override + public int compareTo(PendingCallable other) { + return getScheduledTime().compareTo(other.getScheduledTime()); + } + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java new file mode 100644 index 0000000000..ea9e7e6f4c --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java @@ -0,0 +1,538 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import static org.junit.Assert.assertEquals; + +import com.google.api.core.ApiFuture; +import com.google.api.gax.core.ExecutorProvider; +import com.google.api.gax.core.InstantiatingExecutorProvider; +import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.api.gax.grpc.testing.MockGrpcService; +import com.google.api.gax.grpc.testing.MockServiceHelper; +import com.google.cloud.bigquery.storage.test.JsonTest.ComplexRoot; +import com.google.cloud.bigquery.storage.test.Test.FooType; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors.DescriptorValidationException; +import com.google.protobuf.Int64Value; +import com.google.protobuf.Timestamp; +import java.io.IOException; +import java.math.BigDecimal; +import java.util.Arrays; +import java.util.UUID; +import java.util.concurrent.ExecutionException; +import java.util.logging.Logger; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.threeten.bp.Instant; +import org.threeten.bp.LocalTime; + +@RunWith(JUnit4.class) +public class JsonStreamWriterTest { + private static final Logger LOG = Logger.getLogger(JsonStreamWriterTest.class.getName()); + private static final String TEST_STREAM = "projects/p/datasets/d/tables/t/streams/s"; + private static final String TEST_TABLE = "projects/p/datasets/d/tables/t"; + private static final ExecutorProvider SINGLE_THREAD_EXECUTOR = + InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).build(); + private static LocalChannelProvider channelProvider; + private FakeScheduledExecutorService fakeExecutor; + private FakeBigQueryWrite testBigQueryWrite; + private static MockServiceHelper serviceHelper; + + private final TableFieldSchema FOO = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("foo") + .build(); + private final TableSchema TABLE_SCHEMA = TableSchema.newBuilder().addFields(0, FOO).build(); + + private final TableFieldSchema BAR = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("bar") + .build(); + private final TableFieldSchema BAZ = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("baz") + .build(); + private final TableSchema UPDATED_TABLE_SCHEMA = + TableSchema.newBuilder().addFields(0, FOO).addFields(1, BAR).build(); + private final TableSchema UPDATED_TABLE_SCHEMA_2 = + TableSchema.newBuilder().addFields(0, FOO).addFields(1, BAR).addFields(2, BAZ).build(); + + private final TableFieldSchema TEST_INT = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.INT64) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_int") + .build(); + private final TableFieldSchema TEST_STRING = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_string") + .build(); + private final TableFieldSchema TEST_BYTES = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.BYTES) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("test_bytes") + .build(); + private final TableFieldSchema TEST_BOOL = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.BOOL) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_bool") + .build(); + private final TableFieldSchema TEST_DOUBLE = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DOUBLE) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_double") + .build(); + private final TableFieldSchema TEST_DATE = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DATE) + .setMode(TableFieldSchema.Mode.REQUIRED) + .setName("test_date") + .build(); + private final TableFieldSchema COMPLEXLVL2 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.REQUIRED) + .addFields(0, TEST_INT) + .setName("complex_lvl2") + .build(); + private final TableFieldSchema COMPLEXLVL1 = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRUCT) + .setMode(TableFieldSchema.Mode.REQUIRED) + .addFields(0, TEST_INT) + .addFields(1, COMPLEXLVL2) + .setName("complex_lvl1") + .build(); + private final TableFieldSchema TEST_NUMERIC = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_numeric") + .build(); + private final TableFieldSchema TEST_NUMERIC_REPEATED = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_numeric_repeated") + .build(); + private final TableFieldSchema TEST_GEO = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.GEOGRAPHY) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_geo") + .build(); + private final TableFieldSchema TEST_TIMESTAMP = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.TIMESTAMP) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_timestamp") + .build(); + private final TableFieldSchema TEST_TIME = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.TIME) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_time") + .build(); + private final TableSchema COMPLEX_TABLE_SCHEMA = + TableSchema.newBuilder() + .addFields(0, TEST_INT) + .addFields(1, TEST_STRING) + .addFields(2, TEST_BYTES) + .addFields(3, TEST_BOOL) + .addFields(4, TEST_DOUBLE) + .addFields(5, TEST_DATE) + .addFields(6, COMPLEXLVL1) + .addFields(7, COMPLEXLVL2) + .addFields(8, TEST_NUMERIC) + .addFields(9, TEST_GEO) + .addFields(10, TEST_TIMESTAMP) + .addFields(11, TEST_TIME) + .addFields(12, TEST_NUMERIC_REPEATED) + .build(); + + @Before + public void setUp() throws Exception { + testBigQueryWrite = new FakeBigQueryWrite(); + serviceHelper = + new MockServiceHelper( + UUID.randomUUID().toString(), Arrays.asList(testBigQueryWrite)); + serviceHelper.start(); + channelProvider = serviceHelper.createChannelProvider(); + fakeExecutor = new FakeScheduledExecutorService(); + testBigQueryWrite.setExecutor(fakeExecutor); + Instant time = Instant.now(); + Timestamp timestamp = + Timestamp.newBuilder().setSeconds(time.getEpochSecond()).setNanos(time.getNano()).build(); + // Add enough GetWriteStream response. + for (int i = 0; i < 4; i++) { + testBigQueryWrite.addResponse( + WriteStream.newBuilder().setName(TEST_STREAM).setCreateTime(timestamp).build()); + } + } + + @After + public void tearDown() throws Exception { + serviceHelper.stop(); + } + + private JsonStreamWriter.Builder getTestJsonStreamWriterBuilder( + String testStream, TableSchema BQTableSchema) { + return JsonStreamWriter.newBuilder(testStream, BQTableSchema) + .setChannelProvider(channelProvider) + .setCredentialsProvider(NoCredentialsProvider.create()); + } + + @Test + public void testTwoParamNewBuilder_nullSchema() { + try { + getTestJsonStreamWriterBuilder(null, TABLE_SCHEMA); + Assert.fail("expected NullPointerException"); + } catch (NullPointerException e) { + assertEquals(e.getMessage(), "StreamOrTableName is null."); + } + } + + @Test + public void testTwoParamNewBuilder_nullStream() { + try { + getTestJsonStreamWriterBuilder(TEST_STREAM, null); + Assert.fail("expected NullPointerException"); + } catch (NullPointerException e) { + assertEquals(e.getMessage(), "TableSchema is null."); + } + } + + @Test + public void testTwoParamNewBuilder() + throws DescriptorValidationException, IOException, InterruptedException { + JsonStreamWriter writer = getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build(); + assertEquals(TEST_STREAM, writer.getStreamName()); + } + + @Test + public void testSingleAppendSimpleJson() throws Exception { + FooType expectedProto = FooType.newBuilder().setFoo("allen").build(); + JSONObject foo = new JSONObject(); + foo.put("foo", "allen"); + JSONArray jsonArr = new JSONArray(); + jsonArr.put(foo); + + try (JsonStreamWriter writer = + getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA) + .setTraceId("test:empty") + .build()) { + + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build()) + .build()); + + ApiFuture appendFuture = writer.append(jsonArr); + assertEquals(0L, appendFuture.get().getAppendResult().getOffset().getValue()); + appendFuture.get(); + assertEquals( + 1, + testBigQueryWrite + .getAppendRequests() + .get(0) + .getProtoRows() + .getRows() + .getSerializedRowsCount()); + assertEquals( + testBigQueryWrite + .getAppendRequests() + .get(0) + .getProtoRows() + .getRows() + .getSerializedRows(0), + expectedProto.toByteString()); + assertEquals( + testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriterBeta_test:empty"); + } + } + + @Test + public void testSingleAppendMultipleSimpleJson() throws Exception { + FooType expectedProto = FooType.newBuilder().setFoo("allen").build(); + JSONObject foo = new JSONObject(); + foo.put("foo", "allen"); + JSONObject foo1 = new JSONObject(); + foo1.put("foo", "allen"); + JSONObject foo2 = new JSONObject(); + foo2.put("foo", "allen"); + JSONObject foo3 = new JSONObject(); + foo3.put("foo", "allen"); + JSONArray jsonArr = new JSONArray(); + jsonArr.put(foo); + jsonArr.put(foo1); + jsonArr.put(foo2); + jsonArr.put(foo3); + + try (JsonStreamWriter writer = + getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build()) { + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build()) + .build()); + + ApiFuture appendFuture = writer.append(jsonArr); + + assertEquals(0L, appendFuture.get().getAppendResult().getOffset().getValue()); + appendFuture.get(); + assertEquals( + 4, + testBigQueryWrite + .getAppendRequests() + .get(0) + .getProtoRows() + .getRows() + .getSerializedRowsCount()); + assertEquals( + testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriterBeta:null"); + for (int i = 0; i < 4; i++) { + assertEquals( + testBigQueryWrite + .getAppendRequests() + .get(0) + .getProtoRows() + .getRows() + .getSerializedRows(i), + expectedProto.toByteString()); + } + } + } + + @Test + public void testMultipleAppendSimpleJson() throws Exception { + FooType expectedProto = FooType.newBuilder().setFoo("allen").build(); + JSONObject foo = new JSONObject(); + foo.put("foo", "allen"); + JSONArray jsonArr = new JSONArray(); + jsonArr.put(foo); + + try (JsonStreamWriter writer = + getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build()) { + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build()) + .build()); + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(1)).build()) + .build()); + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(2)).build()) + .build()); + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(3)).build()) + .build()); + ApiFuture appendFuture; + for (int i = 0; i < 4; i++) { + appendFuture = writer.append(jsonArr); + assertEquals((long) i, appendFuture.get().getAppendResult().getOffset().getValue()); + appendFuture.get(); + assertEquals( + 1, + testBigQueryWrite + .getAppendRequests() + .get(i) + .getProtoRows() + .getRows() + .getSerializedRowsCount()); + assertEquals( + testBigQueryWrite + .getAppendRequests() + .get(i) + .getProtoRows() + .getRows() + .getSerializedRows(0), + expectedProto.toByteString()); + } + } + } + + @Test + public void testSingleAppendComplexJson() throws Exception { + ComplexRoot expectedProto = + ComplexRoot.newBuilder() + .setTestInt(1) + .addTestString("a") + .addTestString("b") + .addTestString("c") + .setTestBytes(ByteString.copyFrom("hello".getBytes())) + .setTestBool(true) + .addTestDouble(1.1) + .addTestDouble(2.2) + .addTestDouble(3.3) + .addTestDouble(4.4) + .setTestDate(1) + .setComplexLvl1( + com.google.cloud.bigquery.storage.test.JsonTest.ComplexLvl1.newBuilder() + .setTestInt(2) + .setComplexLvl2( + com.google.cloud.bigquery.storage.test.JsonTest.ComplexLvl2.newBuilder() + .setTestInt(3) + .build()) + .build()) + .setComplexLvl2( + com.google.cloud.bigquery.storage.test.JsonTest.ComplexLvl2.newBuilder() + .setTestInt(3) + .build()) + .setTestNumeric( + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("1.23456"))) + .setTestGeo("POINT(1,1)") + .setTestTimestamp(12345678) + .setTestTime(CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(1, 0, 1))) + .addTestNumericRepeated( + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("0"))) + .addTestNumericRepeated( + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("99999999999999999999999999999.999999999"))) + .addTestNumericRepeated( + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("-99999999999999999999999999999.999999999"))) + .build(); + JSONObject complex_lvl2 = new JSONObject(); + complex_lvl2.put("test_int", 3); + + JSONObject complex_lvl1 = new JSONObject(); + complex_lvl1.put("test_int", 2); + complex_lvl1.put("complex_lvl2", complex_lvl2); + + JSONObject json = new JSONObject(); + json.put("test_int", 1); + json.put("test_string", new JSONArray(new String[] {"a", "b", "c"})); + json.put("test_bytes", ByteString.copyFrom("hello".getBytes())); + json.put("test_bool", true); + json.put("test_DOUBLe", new JSONArray(new Double[] {1.1, 2.2, 3.3, 4.4})); + json.put("test_date", 1); + json.put("complex_lvl1", complex_lvl1); + json.put("complex_lvl2", complex_lvl2); + json.put( + "test_numeric", + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("1.23456"))); + json.put( + "test_numeric_repeated", + new JSONArray( + new byte[][] { + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("0")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("99999999999999999999999999999.999999999")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("-99999999999999999999999999999.999999999")) + .toByteArray(), + })); + json.put("test_geo", "POINT(1,1)"); + json.put("test_timestamp", 12345678); + json.put("test_time", CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(1, 0, 1))); + JSONArray jsonArr = new JSONArray(); + jsonArr.put(json); + + try (JsonStreamWriter writer = + getTestJsonStreamWriterBuilder(TEST_STREAM, COMPLEX_TABLE_SCHEMA).build()) { + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(0)).build()) + .build()); + + ApiFuture appendFuture = writer.append(jsonArr); + + assertEquals(0L, appendFuture.get().getAppendResult().getOffset().getValue()); + appendFuture.get(); + assertEquals( + 1, + testBigQueryWrite + .getAppendRequests() + .get(0) + .getProtoRows() + .getRows() + .getSerializedRowsCount()); + assertEquals( + testBigQueryWrite + .getAppendRequests() + .get(0) + .getProtoRows() + .getRows() + .getSerializedRows(0), + expectedProto.toByteString()); + } + } + + @Test + public void testAppendOutOfRangeException() throws Exception { + try (JsonStreamWriter writer = + getTestJsonStreamWriterBuilder(TEST_STREAM, TABLE_SCHEMA).build()) { + testBigQueryWrite.addResponse( + AppendRowsResponse.newBuilder() + .setError(com.google.rpc.Status.newBuilder().setCode(11).build()) + .build()); + JSONObject foo = new JSONObject(); + foo.put("foo", "allen"); + JSONArray jsonArr = new JSONArray(); + jsonArr.put(foo); + ApiFuture appendFuture = writer.append(jsonArr); + try { + appendFuture.get(); + Assert.fail("expected ExecutionException"); + } catch (ExecutionException ex) { + assertEquals(ex.getCause().getMessage(), "OUT_OF_RANGE: "); + } + } + } + + @Test + public void testCreateDefaultStream() throws Exception { + TableSchema tableSchema = + TableSchema.newBuilder().addFields(0, TEST_INT).addFields(1, TEST_STRING).build(); + try (JsonStreamWriter writer = + JsonStreamWriter.newBuilder(TEST_TABLE, tableSchema) + .setChannelProvider(channelProvider) + .setCredentialsProvider(NoCredentialsProvider.create()) + .build()) { + assertEquals("projects/p/datasets/d/tables/t/_default", writer.getStreamName()); + } + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessageTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessageTest.java new file mode 100644 index 0000000000..b8eba3c893 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessageTest.java @@ -0,0 +1,754 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import com.google.cloud.bigquery.storage.test.JsonTest.*; +import com.google.cloud.bigquery.storage.test.SchemaTest.*; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Message; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Map; +import java.util.logging.Logger; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.Assert; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class JsonToProtoMessageTest { + private static final Logger LOG = Logger.getLogger(JsonToProtoMessageTest.class.getName()); + private static ImmutableMap AllTypesToDebugMessageTest = + new ImmutableMap.Builder() + .put(BoolType.getDescriptor(), "boolean") + .put(BytesType.getDescriptor(), "bytes") + .put(Int64Type.getDescriptor(), "int64") + .put(Int32Type.getDescriptor(), "int32") + .put(DoubleType.getDescriptor(), "double") + .put(StringType.getDescriptor(), "string") + .put(RepeatedType.getDescriptor(), "array") + .put(ObjectType.getDescriptor(), "object") + .build(); + + private static ImmutableMap AllTypesToCorrectProto = + new ImmutableMap.Builder() + .put( + BoolType.getDescriptor(), + new Message[] {BoolType.newBuilder().setTestFieldType(true).build()}) + .put( + BytesType.getDescriptor(), + new Message[] { + BytesType.newBuilder().setTestFieldType(ByteString.copyFromUtf8("test")).build() + }) + .put( + Int64Type.getDescriptor(), + new Message[] { + Int64Type.newBuilder().setTestFieldType(Long.MAX_VALUE).build(), + Int64Type.newBuilder().setTestFieldType(new Long(Integer.MAX_VALUE)).build() + }) + .put( + Int32Type.getDescriptor(), + new Message[] {Int32Type.newBuilder().setTestFieldType(Integer.MAX_VALUE).build()}) + .put( + DoubleType.getDescriptor(), + new Message[] {DoubleType.newBuilder().setTestFieldType(1.23).build()}) + .put( + StringType.getDescriptor(), + new Message[] {StringType.newBuilder().setTestFieldType("test").build()}) + .put( + RepeatedType.getDescriptor(), + new Message[] { + RepeatedType.newBuilder() + .addAllTestFieldType( + new ArrayList() { + { + add(1L); + add(2L); + add(3L); + } + }) + .build() + }) + .put( + ObjectType.getDescriptor(), + new Message[] { + ObjectType.newBuilder() + .setTestFieldType(ComplexLvl2.newBuilder().setTestInt(1).build()) + .build() + }) + .build(); + + private static ImmutableMap AllRepeatedTypesToDebugMessageTest = + new ImmutableMap.Builder() + .put(RepeatedBool.getDescriptor(), "boolean") + .put(RepeatedBytes.getDescriptor(), "bytes") + .put(RepeatedInt64.getDescriptor(), "int64") + .put(RepeatedInt32.getDescriptor(), "int32") + .put(RepeatedDouble.getDescriptor(), "double") + .put(RepeatedString.getDescriptor(), "string") + .put(RepeatedObject.getDescriptor(), "object") + .build(); + + private static ImmutableMap AllRepeatedTypesToCorrectProto = + new ImmutableMap.Builder() + .put( + RepeatedBool.getDescriptor(), + new Message[] { + RepeatedBool.newBuilder().addTestRepeated(true).addTestRepeated(false).build() + }) + .put( + RepeatedBytes.getDescriptor(), + new Message[] { + RepeatedBytes.newBuilder() + .addTestRepeated(ByteString.copyFrom(new byte[] {0})) + .addTestRepeated(ByteString.copyFrom(new byte[] {0, -116, -122, 71})) + .build(), + RepeatedBytes.newBuilder() + .addTestRepeated( + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("0"))) + .addTestRepeated( + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("1.2"))) + .build() + }) + .put( + RepeatedString.getDescriptor(), + new Message[] { + RepeatedString.newBuilder().addTestRepeated("hello").addTestRepeated("test").build() + }) + .put( + RepeatedInt64.getDescriptor(), + new Message[] { + RepeatedInt64.newBuilder() + .addTestRepeated(Long.MAX_VALUE) + .addTestRepeated(Long.MIN_VALUE) + .addTestRepeated(Integer.MAX_VALUE) + .addTestRepeated(Integer.MIN_VALUE) + .addTestRepeated(Short.MAX_VALUE) + .addTestRepeated(Short.MIN_VALUE) + .addTestRepeated(Byte.MAX_VALUE) + .addTestRepeated(Byte.MIN_VALUE) + .addTestRepeated(0) + .build(), + RepeatedInt64.newBuilder() + .addTestRepeated(Integer.MAX_VALUE) + .addTestRepeated(Integer.MIN_VALUE) + .addTestRepeated(Short.MAX_VALUE) + .addTestRepeated(Short.MIN_VALUE) + .addTestRepeated(Byte.MAX_VALUE) + .addTestRepeated(Byte.MIN_VALUE) + .addTestRepeated(0) + .build() + }) + .put( + RepeatedInt32.getDescriptor(), + new Message[] { + RepeatedInt32.newBuilder() + .addTestRepeated(Integer.MAX_VALUE) + .addTestRepeated(Integer.MIN_VALUE) + .addTestRepeated(Short.MAX_VALUE) + .addTestRepeated(Short.MIN_VALUE) + .addTestRepeated(Byte.MAX_VALUE) + .addTestRepeated(Byte.MIN_VALUE) + .addTestRepeated(0) + .build() + }) + .put( + RepeatedDouble.getDescriptor(), + new Message[] { + RepeatedDouble.newBuilder() + .addTestRepeated(Double.MAX_VALUE) + .addTestRepeated(Double.MIN_VALUE) + .addTestRepeated(Float.MAX_VALUE) + .addTestRepeated(Float.MIN_VALUE) + .build(), + RepeatedDouble.newBuilder() + .addTestRepeated(Float.MAX_VALUE) + .addTestRepeated(Float.MIN_VALUE) + .build() + }) + .put( + RepeatedObject.getDescriptor(), + new Message[] { + RepeatedObject.newBuilder() + .addTestRepeated(ComplexLvl2.newBuilder().setTestInt(1).build()) + .addTestRepeated(ComplexLvl2.newBuilder().setTestInt(2).build()) + .addTestRepeated(ComplexLvl2.newBuilder().setTestInt(3).build()) + .build() + }) + .build(); + + private static JSONObject[] simpleJSONObjects = { + new JSONObject().put("test_field_type", Long.MAX_VALUE), + new JSONObject().put("test_field_type", Integer.MAX_VALUE), + new JSONObject().put("test_field_type", 1.23), + new JSONObject().put("test_field_type", true), + new JSONObject().put("test_field_type", ByteString.copyFromUtf8("test")), + new JSONObject().put("test_field_type", new JSONArray("[1, 2, 3]")), + new JSONObject().put("test_field_type", new JSONObject().put("test_int", 1)), + new JSONObject().put("test_field_type", "test") + }; + + private static JSONObject[] simpleJSONArrays = { + new JSONObject() + .put( + "test_repeated", + new JSONArray( + new Long[] { + Long.MAX_VALUE, + Long.MIN_VALUE, + (long) Integer.MAX_VALUE, + (long) Integer.MIN_VALUE, + (long) Short.MAX_VALUE, + (long) Short.MIN_VALUE, + (long) Byte.MAX_VALUE, + (long) Byte.MIN_VALUE, + 0L + })), + new JSONObject() + .put( + "test_repeated", + new JSONArray( + new Integer[] { + Integer.MAX_VALUE, + Integer.MIN_VALUE, + (int) Short.MAX_VALUE, + (int) Short.MIN_VALUE, + (int) Byte.MAX_VALUE, + (int) Byte.MIN_VALUE, + 0 + })), + new JSONObject() + .put( + "test_repeated", + new JSONArray( + new Double[] { + Double.MAX_VALUE, + Double.MIN_VALUE, + (double) Float.MAX_VALUE, + (double) Float.MIN_VALUE + })), + new JSONObject() + .put("test_repeated", new JSONArray(new Float[] {Float.MAX_VALUE, Float.MIN_VALUE})), + new JSONObject().put("test_repeated", new JSONArray(new Boolean[] {true, false})), + new JSONObject().put("test_repeated", new JSONArray(new String[] {"hello", "test"})), + new JSONObject() + .put( + "test_repeated", + new JSONArray( + new byte[][] { + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("0")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("1.2")) + .toByteArray() + })), + new JSONObject().put("test_repeated", new JSONArray(new int[][] {{11111, 22222}})), + new JSONObject().put("test_repeated", new JSONArray(new char[][] {{'a', 'b'}, {'c'}})), + new JSONObject().put("test_repeated", new JSONArray(new String[][] {{"hello"}, {"test"}})), + new JSONObject() + .put( + "test_repeated", + new JSONArray( + new JSONObject[] { + new JSONObject().put("test_int", 1), + new JSONObject().put("test_int", 2), + new JSONObject().put("test_int", 3) + })) + }; + + @Test + public void testDifferentNameCasing() throws Exception { + TestInt64 expectedProto = + TestInt64.newBuilder().setByte(1).setShort(1).setInt(1).setLong(1).build(); + + JSONObject json = new JSONObject(); + json.put("bYtE", (byte) 1); + json.put("SHORT", (short) 1); + json.put("inT", 1); + json.put("lONg", 1L); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestInt64.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testInt64() throws Exception { + TestInt64 expectedProto = + TestInt64.newBuilder().setByte(1).setShort(1).setInt(1).setLong(1).build(); + JSONObject json = new JSONObject(); + json.put("byte", (byte) 1); + json.put("short", (short) 1); + json.put("int", 1); + json.put("long", 1L); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestInt64.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testInt32() throws Exception { + TestInt32 expectedProto = TestInt32.newBuilder().setByte(1).setShort(1).setInt(1).build(); + JSONObject json = new JSONObject(); + json.put("byte", (byte) 1); + json.put("short", (short) 1); + json.put("int", 1); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestInt32.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testInt32NotMatchInt64() throws Exception { + JSONObject json = new JSONObject(); + json.put("byte", (byte) 1); + json.put("short", (short) 1); + json.put("int", 1L); + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestInt32.getDescriptor(), json); + Assert.fail("should fail"); + } catch (IllegalArgumentException e) { + assertEquals("JSONObject does not have a int32 field at root.int.", e.getMessage()); + } + } + + @Test + public void testDouble() throws Exception { + TestDouble expectedProto = TestDouble.newBuilder().setDouble(1.2).setFloat(3.4f).build(); + JSONObject json = new JSONObject(); + json.put("double", 1.2); + json.put("float", 3.4f); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestDouble.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testAllTypes() throws Exception { + for (Map.Entry entry : AllTypesToDebugMessageTest.entrySet()) { + int success = 0; + for (JSONObject json : simpleJSONObjects) { + try { + LOG.info("Testing " + json + " over " + entry.getKey().getFullName()); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(entry.getKey(), json); + LOG.info("Convert Success!"); + assertEquals(protoMsg, AllTypesToCorrectProto.get(entry.getKey())[success]); + success += 1; + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject does not have a " + entry.getValue() + " field at root.test_field_type.", + e.getMessage()); + } + } + if (entry.getKey() == Int64Type.getDescriptor()) { + assertEquals(entry.getKey().getFullName(), 2, success); + } else { + assertEquals(entry.getKey().getFullName(), 1, success); + } + } + } + + @Test + public void testAllRepeatedTypesWithLimits() throws Exception { + for (Map.Entry entry : AllRepeatedTypesToDebugMessageTest.entrySet()) { + int success = 0; + for (JSONObject json : simpleJSONArrays) { + try { + LOG.info("Testing " + json + " over " + entry.getKey().getFullName()); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(entry.getKey(), json); + LOG.info("Convert Success!"); + assertEquals( + protoMsg.toString(), + protoMsg, + AllRepeatedTypesToCorrectProto.get(entry.getKey())[success]); + success += 1; + } catch (IllegalArgumentException e) { + LOG.info(e.getMessage()); + assertTrue( + e.getMessage() + .equals( + "JSONObject does not have a " + + entry.getValue() + + " field at root.test_repeated[0].") + || e.getMessage() + .equals("Error: root.test_repeated[0] could not be converted to byte[].")); + } + } + if (entry.getKey() == RepeatedInt64.getDescriptor() + || entry.getKey() == RepeatedDouble.getDescriptor()) { + assertEquals(entry.getKey().getFullName(), 2, success); + } else { + assertEquals(entry.getKey().getFullName(), 1, success); + } + } + } + + @Test + public void testOptional() throws Exception { + TestInt64 expectedProto = TestInt64.newBuilder().setByte(1).build(); + JSONObject json = new JSONObject(); + json.put("byte", 1); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestInt64.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testRepeatedIsOptional() throws Exception { + TestRepeatedIsOptional expectedProto = + TestRepeatedIsOptional.newBuilder().setRequiredDouble(1.1).build(); + JSONObject json = new JSONObject(); + json.put("required_double", 1.1); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestRepeatedIsOptional.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testRequired() throws Exception { + JSONObject json = new JSONObject(); + json.put("optional_double", 1.1); + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestRequired.getDescriptor(), json); + Assert.fail("should fail"); + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject does not have the required field root.required_double.", e.getMessage()); + } + } + + @Test + public void testStructSimple() throws Exception { + MessageType expectedProto = + MessageType.newBuilder() + .setTestFieldType(StringType.newBuilder().setTestFieldType("test").build()) + .build(); + JSONObject stringType = new JSONObject(); + stringType.put("test_field_type", "test"); + JSONObject json = new JSONObject(); + json.put("test_field_type", stringType); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(MessageType.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testStructSimpleFail() throws Exception { + JSONObject stringType = new JSONObject(); + stringType.put("test_field_type", 1); + JSONObject json = new JSONObject(); + json.put("test_field_type", stringType); + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(MessageType.getDescriptor(), json); + Assert.fail("should fail"); + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject does not have a string field at root.test_field_type.test_field_type.", + e.getMessage()); + } + } + + @Test + public void testStructComplex() throws Exception { + ComplexRoot expectedProto = + ComplexRoot.newBuilder() + .setTestInt(1) + .addTestString("a") + .addTestString("b") + .addTestString("c") + .setTestBytes(ByteString.copyFrom("hello".getBytes())) + .setTestBool(true) + .addTestDouble(1.1) + .addTestDouble(2.2) + .addTestDouble(3.3) + .addTestDouble(4.4) + .setTestDate(1) + .setComplexLvl1( + ComplexLvl1.newBuilder() + .setTestInt(2) + .setComplexLvl2(ComplexLvl2.newBuilder().setTestInt(3).build()) + .build()) + .setComplexLvl2(ComplexLvl2.newBuilder().setTestInt(3).build()) + .build(); + JSONObject complex_lvl2 = new JSONObject(); + complex_lvl2.put("test_int", 3); + + JSONObject complex_lvl1 = new JSONObject(); + complex_lvl1.put("test_int", 2); + complex_lvl1.put("complex_lvl2", complex_lvl2); + + JSONObject json = new JSONObject(); + json.put("test_int", 1); + json.put("test_string", new JSONArray(new String[] {"a", "b", "c"})); + json.put("test_bytes", ByteString.copyFromUtf8("hello")); + json.put("test_bool", true); + json.put("test_DOUBLe", new JSONArray(new Double[] {1.1, 2.2, 3.3, 4.4})); + json.put("test_date", 1); + json.put("complex_lvl1", complex_lvl1); + json.put("complex_lvl2", complex_lvl2); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(ComplexRoot.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testStructComplexFail() throws Exception { + JSONObject complex_lvl2 = new JSONObject(); + complex_lvl2.put("test_int", 3); + + JSONObject complex_lvl1 = new JSONObject(); + complex_lvl1.put("test_int", "not_int"); + complex_lvl1.put("complex_lvl2", complex_lvl2); + + JSONObject json = new JSONObject(); + json.put("test_int", 1); + json.put("test_string", new JSONArray(new String[] {"a", "b", "c"})); + json.put("test_bytes", ByteString.copyFromUtf8("hello")); + json.put("test_bool", true); + json.put("test_double", new JSONArray(new Double[] {1.1, 2.2, 3.3, 4.4})); + json.put("test_date", 1); + json.put("complex_lvl1", complex_lvl1); + json.put("complex_lvl2", complex_lvl2); + + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(ComplexRoot.getDescriptor(), json); + Assert.fail("should fail"); + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject does not have a int64 field at root.complex_lvl1.test_int.", e.getMessage()); + } + } + + @Test + public void testRepeatedWithMixedTypes() throws Exception { + JSONObject json = new JSONObject(); + json.put("test_repeated", new JSONArray("[1.1, 2.2, true]")); + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(RepeatedDouble.getDescriptor(), json); + Assert.fail("should fail"); + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject does not have a double field at root.test_repeated[2].", e.getMessage()); + } + } + + @Test + public void testNestedRepeatedComplex() throws Exception { + NestedRepeated expectedProto = + NestedRepeated.newBuilder() + .addDouble(1.1) + .addDouble(2.2) + .addDouble(3.3) + .addDouble(4.4) + .addDouble(5.5) + .addInt(1) + .addInt(2) + .addInt(3) + .addInt(4) + .addInt(5) + .setRepeatedString( + RepeatedString.newBuilder() + .addTestRepeated("hello") + .addTestRepeated("this") + .addTestRepeated("is") + .addTestRepeated("a") + .addTestRepeated("test") + .build()) + .build(); + double[] doubleArr = {1.1, 2.2, 3.3, 4.4, 5.5}; + String[] stringArr = {"hello", "this", "is", "a", "test"}; + int[] intArr = {1, 2, 3, 4, 5}; + + JSONObject json = new JSONObject(); + json.put("double", new JSONArray(doubleArr)); + json.put("int", new JSONArray(intArr)); + JSONObject jsonRepeatedString = new JSONObject(); + jsonRepeatedString.put("test_repeated", new JSONArray(stringArr)); + json.put("repeated_string", jsonRepeatedString); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(NestedRepeated.getDescriptor(), json); + assertEquals(protoMsg, expectedProto); + } + + @Test + public void testNestedRepeatedComplexFail() throws Exception { + double[] doubleArr = {1.1, 2.2, 3.3, 4.4, 5.5}; + Boolean[] fakeStringArr = {true, false}; + int[] intArr = {1, 2, 3, 4, 5}; + + JSONObject json = new JSONObject(); + json.put("double", new JSONArray(doubleArr)); + json.put("int", new JSONArray(intArr)); + JSONObject jsonRepeatedString = new JSONObject(); + jsonRepeatedString.put("test_repeated", new JSONArray(fakeStringArr)); + json.put("repeated_string", jsonRepeatedString); + + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(NestedRepeated.getDescriptor(), json); + Assert.fail("should fail"); + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject does not have a string field at root.repeated_string.test_repeated[0].", + e.getMessage()); + } + } + + @Test + public void testEmptySecondLevelObject() throws Exception { + ComplexLvl1 expectedProto = + ComplexLvl1.newBuilder() + .setTestInt(1) + .setComplexLvl2(ComplexLvl2.newBuilder().build()) + .build(); + JSONObject complexLvl2 = new JSONObject(); + JSONObject json = new JSONObject(); + json.put("test_int", 1); + json.put("complex_lvl2", complexLvl2); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(ComplexLvl1.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testAllowUnknownFieldsError() throws Exception { + JSONObject json = new JSONObject(); + json.put("test_repeated", new JSONArray(new int[] {1, 2, 3, 4, 5})); + json.put("string", "hello"); + + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(RepeatedInt64.getDescriptor(), json); + Assert.fail("Should fail"); + } catch (IllegalArgumentException e) { + assertEquals("JSONObject has fields unknown to BigQuery: root.string.", e.getMessage()); + } + } + + @Test + public void testEmptyProtoMessage() throws Exception { + JSONObject json = new JSONObject(); + json.put("test_repeated", new JSONArray(new int[0])); + + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(RepeatedInt64.getDescriptor(), json); + Assert.fail("Should fail"); + } catch (IllegalArgumentException e) { + assertEquals("The created protobuf message is empty.", e.getMessage()); + } + } + + @Test + public void testEmptyJSONObject() throws Exception { + JSONObject json = new JSONObject(); + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(Int64Type.getDescriptor(), json); + Assert.fail("Should fail"); + } catch (IllegalStateException e) { + assertEquals("JSONObject is empty.", e.getMessage()); + } + } + + @Test + public void testNullJson() throws Exception { + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(Int64Type.getDescriptor(), null); + Assert.fail("Should fail"); + } catch (NullPointerException e) { + assertEquals("JSONObject is null.", e.getMessage()); + } + } + + @Test + public void testNullDescriptor() throws Exception { + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(null, new JSONObject()); + Assert.fail("Should fail"); + } catch (NullPointerException e) { + assertEquals("Protobuf descriptor is null.", e.getMessage()); + } + } + + @Test + public void testAllowUnknownFieldsSecondLevel() throws Exception { + JSONObject complex_lvl2 = new JSONObject(); + complex_lvl2.put("no_match", 1); + JSONObject json = new JSONObject(); + json.put("test_int", 1); + json.put("complex_lvl2", complex_lvl2); + + try { + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(ComplexLvl1.getDescriptor(), json); + Assert.fail("Should fail"); + } catch (IllegalArgumentException e) { + assertEquals( + "JSONObject has fields unknown to BigQuery: root.complex_lvl2.no_match.", e.getMessage()); + } + } + + @Test + public void testTopLevelMatchSecondLevelMismatch() throws Exception { + ComplexLvl1 expectedProto = + ComplexLvl1.newBuilder() + .setTestInt(1) + .setComplexLvl2(ComplexLvl2.newBuilder().build()) + .build(); + JSONObject complex_lvl2 = new JSONObject(); + JSONObject json = new JSONObject(); + json.put("test_int", 1); + json.put("complex_lvl2", complex_lvl2); + + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(ComplexLvl1.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } + + @Test + public void testJsonNullValue() throws Exception { + TestInt64 expectedProto = TestInt64.newBuilder().setInt(1).build(); + JSONObject json = new JSONObject(); + json.put("long", JSONObject.NULL); + json.put("int", 1); + DynamicMessage protoMsg = + JsonToProtoMessage.convertJsonToProtoMessage(TestInt64.getDescriptor(), json); + assertEquals(expectedProto, protoMsg); + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ProtoSchemaConverterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ProtoSchemaConverterTest.java new file mode 100644 index 0000000000..8a2b8dc97b --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ProtoSchemaConverterTest.java @@ -0,0 +1,192 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import com.google.api.gax.rpc.InvalidArgumentException; +import com.google.cloud.bigquery.storage.test.Test.*; +import com.google.protobuf.DescriptorProtos.FileDescriptorProto; +import com.google.protobuf.Descriptors; +import org.junit.*; + +public class ProtoSchemaConverterTest { + @Test + public void convertSimple() { + AllSupportedTypes testProto = AllSupportedTypes.newBuilder().setStringValue("abc").build(); + ProtoSchema protoSchema = ProtoSchemaConverter.convert(testProto.getDescriptorForType()); + Assert.assertEquals( + "name: \"com_google_cloud_bigquery_storage_test_AllSupportedTypes\"\n" + + "field {\n" + + " name: \"int32_value\"\n" + + " number: 1\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_INT32\n" + + "}\n" + + "field {\n" + + " name: \"int64_value\"\n" + + " number: 2\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_INT64\n" + + "}\n" + + "field {\n" + + " name: \"uint32_value\"\n" + + " number: 3\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_UINT32\n" + + "}\n" + + "field {\n" + + " name: \"uint64_value\"\n" + + " number: 4\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_UINT64\n" + + "}\n" + + "field {\n" + + " name: \"float_value\"\n" + + " number: 5\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_FLOAT\n" + + "}\n" + + "field {\n" + + " name: \"double_value\"\n" + + " number: 6\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_DOUBLE\n" + + "}\n" + + "field {\n" + + " name: \"bool_value\"\n" + + " number: 7\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_BOOL\n" + + "}\n" + + "field {\n" + + " name: \"enum_value\"\n" + + " number: 8\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_ENUM\n" + + " type_name: \"com_google_cloud_bigquery_storage_test_TestEnum_E.TestEnum\"\n" + + "}\n" + + "field {\n" + + " name: \"string_value\"\n" + + " number: 9\n" + + " label: LABEL_REQUIRED\n" + + " type: TYPE_STRING\n" + + "}\n" + + "nested_type {\n" + + " name: \"com_google_cloud_bigquery_storage_test_TestEnum_E\"\n" + + " enum_type {\n" + + " name: \"TestEnum\"\n" + + " value {\n" + + " name: \"TestEnum0\"\n" + + " number: 0\n" + + " }\n" + + " value {\n" + + " name: \"TestEnum1\"\n" + + " number: 1\n" + + " }\n" + + " }\n" + + "}\n", + protoSchema.getProtoDescriptor().toString()); + } + + @Test + public void convertNested() { + ComplicateType testProto = ComplicateType.newBuilder().build(); + ProtoSchema protoSchema = ProtoSchemaConverter.convert(testProto.getDescriptorForType()); + Assert.assertEquals( + "name: \"com_google_cloud_bigquery_storage_test_ComplicateType\"\n" + + "field {\n" + + " name: \"nested_repeated_type\"\n" + + " number: 1\n" + + " label: LABEL_REPEATED\n" + + " type: TYPE_MESSAGE\n" + + " type_name: \"com_google_cloud_bigquery_storage_test_NestedType\"\n" + + "}\n" + + "field {\n" + + " name: \"inner_type\"\n" + + " number: 2\n" + + " label: LABEL_OPTIONAL\n" + + " type: TYPE_MESSAGE\n" + + " type_name: \"com_google_cloud_bigquery_storage_test_InnerType\"\n" + + "}\n" + + "nested_type {\n" + + " name: \"com_google_cloud_bigquery_storage_test_InnerType\"\n" + + " field {\n" + + " name: \"value\"\n" + + " number: 1\n" + + " label: LABEL_REPEATED\n" + + " type: TYPE_STRING\n" + + " }\n" + + "}\n" + + "nested_type {\n" + + " name: \"com_google_cloud_bigquery_storage_test_NestedType\"\n" + + " field {\n" + + " name: \"inner_type\"\n" + + " number: 1\n" + + " label: LABEL_REPEATED\n" + + " type: TYPE_MESSAGE\n" + + " type_name: \"com_google_cloud_bigquery_storage_test_InnerType\"\n" + + " }\n" + + "}\n", + protoSchema.getProtoDescriptor().toString()); + } + + @Test + public void convertRecursive() { + try { + RecursiveType testProto = RecursiveType.newBuilder().build(); + ProtoSchema protoSchema = ProtoSchemaConverter.convert(testProto.getDescriptorForType()); + Assert.fail("No exception raised"); + } catch (InvalidArgumentException e) { + Assert.assertEquals( + "Recursive type is not supported:com.google.cloud.bigquery.storage.test.RecursiveType", + e.getMessage()); + } + } + + @Test + public void convertRecursiveTopMessage() { + try { + RecursiveTypeTopMessage testProto = RecursiveTypeTopMessage.newBuilder().build(); + ProtoSchema protoSchema = ProtoSchemaConverter.convert(testProto.getDescriptorForType()); + Assert.fail("No exception raised"); + } catch (InvalidArgumentException e) { + Assert.assertEquals( + "Recursive type is not supported:com.google.cloud.bigquery.storage.test.RecursiveTypeTopMessage", + e.getMessage()); + } + } + + @Test + public void convertDuplicateType() { + DuplicateType testProto = DuplicateType.newBuilder().build(); + ProtoSchema protoSchema = ProtoSchemaConverter.convert(testProto.getDescriptorForType()); + + FileDescriptorProto fileDescriptorProto = + FileDescriptorProto.newBuilder() + .setName("foo.proto") + .addMessageType(protoSchema.getProtoDescriptor()) + .build(); + try { + Descriptors.FileDescriptor fs = + Descriptors.FileDescriptor.buildFrom( + fileDescriptorProto, new Descriptors.FileDescriptor[0]); + Descriptors.Descriptor type = + fs.findMessageTypeByName(protoSchema.getProtoDescriptor().getName()); + Assert.assertEquals(4, type.getFields().size()); + } catch (Descriptors.DescriptorValidationException ex) { + Assert.fail("Got unexpected exception: " + ex.getMessage()); + } + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java new file mode 100644 index 0000000000..a07616dbdb --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java @@ -0,0 +1,507 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery.storage.v1; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.api.core.ApiFuture; +import com.google.api.gax.core.NoCredentialsProvider; +import com.google.api.gax.grpc.testing.MockGrpcService; +import com.google.api.gax.grpc.testing.MockServiceHelper; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.StatusCode.Code; +import com.google.cloud.bigquery.storage.test.Test.FooType; +import com.google.common.base.Strings; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Int64Value; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.logging.Logger; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.function.ThrowingRunnable; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.threeten.bp.Duration; + +@RunWith(JUnit4.class) +public class StreamWriterTest { + private static final Logger log = Logger.getLogger(StreamWriterTest.class.getName()); + private static final String TEST_STREAM = "projects/p/datasets/d/tables/t/streams/s"; + private static final String TEST_TRACE_ID = "DATAFLOW:job_id"; + private FakeScheduledExecutorService fakeExecutor; + private FakeBigQueryWrite testBigQueryWrite; + private static MockServiceHelper serviceHelper; + private BigQueryWriteClient client; + + @Before + public void setUp() throws Exception { + testBigQueryWrite = new FakeBigQueryWrite(); + serviceHelper = + new MockServiceHelper( + UUID.randomUUID().toString(), Arrays.asList(testBigQueryWrite)); + serviceHelper.start(); + fakeExecutor = new FakeScheduledExecutorService(); + testBigQueryWrite.setExecutor(fakeExecutor); + client = + BigQueryWriteClient.create( + BigQueryWriteSettings.newBuilder() + .setCredentialsProvider(NoCredentialsProvider.create()) + .setTransportChannelProvider(serviceHelper.createChannelProvider()) + .build()); + } + + @After + public void tearDown() throws Exception { + log.info("tearDown called"); + client.close(); + serviceHelper.stop(); + } + + private StreamWriter getTestStreamWriter() throws IOException { + return StreamWriter.newBuilder(TEST_STREAM, client) + .setWriterSchema(createProtoSchema()) + .setTraceId(TEST_TRACE_ID) + .build(); + } + + private ProtoSchema createProtoSchema() { + return ProtoSchema.newBuilder() + .setProtoDescriptor( + DescriptorProtos.DescriptorProto.newBuilder() + .setName("Message") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("foo") + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .setNumber(1) + .build()) + .build()) + .build(); + } + + private ProtoRows createProtoRows(String[] messages) { + ProtoRows.Builder rowsBuilder = ProtoRows.newBuilder(); + for (String message : messages) { + FooType foo = FooType.newBuilder().setFoo(message).build(); + rowsBuilder.addSerializedRows(foo.toByteString()); + } + return rowsBuilder.build(); + } + + private AppendRowsResponse createAppendResponse(long offset) { + return AppendRowsResponse.newBuilder() + .setAppendResult( + AppendRowsResponse.AppendResult.newBuilder().setOffset(Int64Value.of(offset)).build()) + .build(); + } + + private AppendRowsResponse createAppendResponseWithError(Status.Code code, String message) { + return AppendRowsResponse.newBuilder() + .setError(com.google.rpc.Status.newBuilder().setCode(code.value()).setMessage(message)) + .build(); + } + + private ApiFuture sendTestMessage(StreamWriter writer, String[] messages) { + return writer.append(createProtoRows(messages), -1); + } + + private static T assertFutureException( + Class expectedThrowable, final Future future) { + return assertThrows( + expectedThrowable, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + try { + future.get(); + } catch (ExecutionException ex) { + // Future wraps exception with ExecutionException. So unwrapper it here. + throw ex.getCause(); + } + } + }); + } + + private void verifyAppendIsBlocked(final StreamWriter writer) throws Exception { + Thread appendThread = + new Thread( + new Runnable() { + @Override + public void run() { + sendTestMessage(writer, new String[] {"A"}); + } + }); + // Start a separate thread to append and verify that it is still alive after 2 seoncds. + appendThread.start(); + TimeUnit.SECONDS.sleep(2); + assertTrue(appendThread.isAlive()); + appendThread.interrupt(); + } + + private void verifyAppendRequests(long appendCount) { + assertEquals(appendCount, testBigQueryWrite.getAppendRequests().size()); + for (int i = 0; i < appendCount; i++) { + AppendRowsRequest serverRequest = testBigQueryWrite.getAppendRequests().get(i); + assertTrue(serverRequest.getProtoRows().getRows().getSerializedRowsCount() > 0); + assertEquals(i, serverRequest.getOffset().getValue()); + if (i == 0) { + // First request received by server should have schema and stream name. + assertTrue(serverRequest.getProtoRows().hasWriterSchema()); + assertEquals(serverRequest.getWriteStream(), TEST_STREAM); + assertEquals(serverRequest.getTraceId(), TEST_TRACE_ID); + } else { + // Following request should not have schema and stream name. + assertFalse(serverRequest.getProtoRows().hasWriterSchema()); + assertEquals(serverRequest.getWriteStream(), ""); + assertEquals(serverRequest.getTraceId(), ""); + } + } + } + + @Test + public void testBuildBigQueryWriteClientInWriter() throws Exception { + StreamWriter writer = + StreamWriter.newBuilder(TEST_STREAM) + .setCredentialsProvider(NoCredentialsProvider.create()) + .setChannelProvider(serviceHelper.createChannelProvider()) + .setWriterSchema(createProtoSchema()) + .build(); + + testBigQueryWrite.addResponse(createAppendResponse(0)); + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + writer.close(); + } + + @Test + public void testAppendSuccess() throws Exception { + StreamWriter writer = getTestStreamWriter(); + + long appendCount = 100; + for (int i = 0; i < appendCount; i++) { + testBigQueryWrite.addResponse(createAppendResponse(i)); + } + + List> futures = new ArrayList<>(); + for (int i = 0; i < appendCount; i++) { + futures.add(writer.append(createProtoRows(new String[] {String.valueOf(i)}), i)); + } + + for (int i = 0; i < appendCount; i++) { + assertEquals(i, futures.get(i).get().getAppendResult().getOffset().getValue()); + } + + verifyAppendRequests(appendCount); + + writer.close(); + } + + @Test + public void testNoSchema() throws Exception { + StatusRuntimeException ex = + assertThrows( + StatusRuntimeException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriter.newBuilder(TEST_STREAM, client).build(); + } + }); + assertEquals(ex.getStatus().getCode(), Status.INVALID_ARGUMENT.getCode()); + assertTrue(ex.getStatus().getDescription().contains("Writer schema must be provided")); + } + + @Test + public void testInvalidTraceId() throws Exception { + assertThrows( + IllegalArgumentException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriter.newBuilder(TEST_STREAM).setTraceId("abc"); + } + }); + assertThrows( + IllegalArgumentException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriter.newBuilder(TEST_STREAM).setTraceId("abc:"); + } + }); + assertThrows( + IllegalArgumentException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + StreamWriter.newBuilder(TEST_STREAM).setTraceId(":abc"); + } + }); + } + + @Test + public void testAppendSuccessAndConnectionError() throws Exception { + StreamWriter writer = getTestStreamWriter(); + testBigQueryWrite.addResponse(createAppendResponse(0)); + testBigQueryWrite.addException(Status.INTERNAL.asException()); + + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + ApiFuture appendFuture2 = sendTestMessage(writer, new String[] {"B"}); + + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + ApiException actualError = assertFutureException(ApiException.class, appendFuture2); + assertEquals(Code.INTERNAL, actualError.getStatusCode().getCode()); + + writer.close(); + } + + @Test + public void testAppendSuccessAndInStreamError() throws Exception { + StreamWriter writer = getTestStreamWriter(); + testBigQueryWrite.addResponse(createAppendResponse(0)); + testBigQueryWrite.addResponse( + createAppendResponseWithError(Status.INVALID_ARGUMENT.getCode(), "test message")); + testBigQueryWrite.addResponse(createAppendResponse(1)); + + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + ApiFuture appendFuture2 = sendTestMessage(writer, new String[] {"B"}); + ApiFuture appendFuture3 = sendTestMessage(writer, new String[] {"C"}); + + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + StatusRuntimeException actualError = + assertFutureException(StatusRuntimeException.class, appendFuture2); + assertEquals(Status.Code.INVALID_ARGUMENT, actualError.getStatus().getCode()); + assertEquals("test message", actualError.getStatus().getDescription()); + assertEquals(1, appendFuture3.get().getAppendResult().getOffset().getValue()); + + writer.close(); + } + + @Test + public void longIdleBetweenAppends() throws Exception { + StreamWriter writer = getTestStreamWriter(); + testBigQueryWrite.addResponse(createAppendResponse(0)); + testBigQueryWrite.addResponse(createAppendResponse(1)); + + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + + // Sleep to create a long idle between appends. + TimeUnit.SECONDS.sleep(3); + + ApiFuture appendFuture2 = sendTestMessage(writer, new String[] {"B"}); + assertEquals(1, appendFuture2.get().getAppendResult().getOffset().getValue()); + + writer.close(); + } + + @Test + public void testAppendAfterUserClose() throws Exception { + StreamWriter writer = getTestStreamWriter(); + testBigQueryWrite.addResponse(createAppendResponse(0)); + + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + writer.close(); + ApiFuture appendFuture2 = sendTestMessage(writer, new String[] {"B"}); + + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + assertTrue(appendFuture2.isDone()); + StatusRuntimeException actualError = + assertFutureException(StatusRuntimeException.class, appendFuture2); + assertEquals(Status.Code.FAILED_PRECONDITION, actualError.getStatus().getCode()); + } + + @Test + public void testAppendAfterServerClose() throws Exception { + StreamWriter writer = getTestStreamWriter(); + testBigQueryWrite.addException(Status.INTERNAL.asException()); + + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + ApiException error1 = assertFutureException(ApiException.class, appendFuture1); + assertEquals(Code.INTERNAL, error1.getStatusCode().getCode()); + + ApiFuture appendFuture2 = sendTestMessage(writer, new String[] {"B"}); + assertTrue(appendFuture2.isDone()); + StatusRuntimeException error2 = + assertFutureException(StatusRuntimeException.class, appendFuture2); + assertEquals(Status.Code.FAILED_PRECONDITION, error2.getStatus().getCode()); + + writer.close(); + } + + @Test + public void userCloseWhileRequestInflight() throws Exception { + final StreamWriter writer = getTestStreamWriter(); + // Server will sleep 2 seconds before sending back the response. + testBigQueryWrite.setResponseSleep(Duration.ofSeconds(2)); + testBigQueryWrite.addResponse(createAppendResponse(0)); + + // Send a request and close the stream in separate thread while the request is inflight. + final ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + Thread closeThread = + new Thread( + new Runnable() { + @Override + public void run() { + writer.close(); + } + }); + closeThread.start(); + + // Due to the sleep on server, the append won't finish within 1 second even though stream + // is being closed. + assertThrows( + TimeoutException.class, + new ThrowingRunnable() { + @Override + public void run() throws Throwable { + appendFuture1.get(1, TimeUnit.SECONDS); + } + }); + + // Within 2 seconds, the request should be done and stream should be closed. + closeThread.join(2000); + assertTrue(appendFuture1.isDone()); + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + } + + @Test + public void serverCloseWhileRequestsInflight() throws Exception { + StreamWriter writer = getTestStreamWriter(); + // Server will sleep 2 seconds before closing the connection. + testBigQueryWrite.setResponseSleep(Duration.ofSeconds(2)); + testBigQueryWrite.addException(Status.INTERNAL.asException()); + + // Send 10 requests, so that there are 10 inflight requests. + int appendCount = 10; + List> futures = new ArrayList<>(); + for (int i = 0; i < appendCount; i++) { + futures.add(sendTestMessage(writer, new String[] {String.valueOf(i)})); + } + + // Server close should properly handle all inflight requests. + for (int i = 0; i < appendCount; i++) { + ApiException actualError = assertFutureException(ApiException.class, futures.get(i)); + assertEquals(Code.INTERNAL, actualError.getStatusCode().getCode()); + } + + writer.close(); + } + + @Test + public void testZeroMaxInflightRequests() throws Exception { + StreamWriter writer = + StreamWriter.newBuilder(TEST_STREAM, client) + .setWriterSchema(createProtoSchema()) + .setMaxInflightRequests(0) + .build(); + testBigQueryWrite.addResponse(createAppendResponse(0)); + verifyAppendIsBlocked(writer); + writer.close(); + } + + @Test + public void testZeroMaxInflightBytes() throws Exception { + StreamWriter writer = + StreamWriter.newBuilder(TEST_STREAM, client) + .setWriterSchema(createProtoSchema()) + .setMaxInflightBytes(0) + .build(); + testBigQueryWrite.addResponse(createAppendResponse(0)); + verifyAppendIsBlocked(writer); + writer.close(); + } + + @Test + public void testOneMaxInflightRequests() throws Exception { + StreamWriter writer = + StreamWriter.newBuilder(TEST_STREAM, client) + .setWriterSchema(createProtoSchema()) + .setMaxInflightRequests(1) + .build(); + // Server will sleep 1 second before every response. + testBigQueryWrite.setResponseSleep(Duration.ofSeconds(1)); + testBigQueryWrite.addResponse(createAppendResponse(0)); + + long appendStartTimeMs = System.currentTimeMillis(); + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {"A"}); + long appendElapsedMs = System.currentTimeMillis() - appendStartTimeMs; + assertTrue(appendElapsedMs >= 1000); + assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue()); + writer.close(); + } + + @Test + public void testAppendsWithTinyMaxInflightBytes() throws Exception { + StreamWriter writer = + StreamWriter.newBuilder(TEST_STREAM, client) + .setWriterSchema(createProtoSchema()) + .setMaxInflightBytes(1) + .build(); + // Server will sleep 100ms before every response. + testBigQueryWrite.setResponseSleep(Duration.ofMillis(100)); + long appendCount = 10; + for (int i = 0; i < appendCount; i++) { + testBigQueryWrite.addResponse(createAppendResponse(i)); + } + + List> futures = new ArrayList<>(); + long appendStartTimeMs = System.currentTimeMillis(); + for (int i = 0; i < appendCount; i++) { + futures.add(writer.append(createProtoRows(new String[] {String.valueOf(i)}), i)); + } + long appendElapsedMs = System.currentTimeMillis() - appendStartTimeMs; + assertTrue(appendElapsedMs >= 1000); + + for (int i = 0; i < appendCount; i++) { + assertEquals(i, futures.get(i).get().getAppendResult().getOffset().getValue()); + } + assertEquals(appendCount, testBigQueryWrite.getAppendRequests().size()); + for (int i = 0; i < appendCount; i++) { + assertEquals(i, testBigQueryWrite.getAppendRequests().get(i).getOffset().getValue()); + } + writer.close(); + } + + @Test + public void testMessageTooLarge() throws Exception { + StreamWriter writer = getTestStreamWriter(); + + String oversized = Strings.repeat("a", (int) (StreamWriter.getApiMaxRequestBytes() + 1)); + ApiFuture appendFuture1 = sendTestMessage(writer, new String[] {oversized}); + assertTrue(appendFuture1.isDone()); + StatusRuntimeException actualError = + assertFutureException(StatusRuntimeException.class, appendFuture1); + assertEquals(Status.Code.INVALID_ARGUMENT, actualError.getStatus().getCode()); + assertTrue(actualError.getStatus().getDescription().contains("MessageSize is too large")); + + writer.close(); + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryBigDecimalByteStringEncoderTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryBigDecimalByteStringEncoderTest.java new file mode 100644 index 0000000000..ffcb58fb85 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryBigDecimalByteStringEncoderTest.java @@ -0,0 +1,182 @@ +/* + * Copyright 2021 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.storage.v1.it; + +import static org.junit.Assert.assertEquals; + +import com.google.api.core.ApiFuture; +import com.google.cloud.ServiceOptions; +import com.google.cloud.bigquery.BigQuery; +import com.google.cloud.bigquery.DatasetInfo; +import com.google.cloud.bigquery.Field.Mode; +import com.google.cloud.bigquery.FieldValueList; +import com.google.cloud.bigquery.Schema; +import com.google.cloud.bigquery.StandardSQLTypeName; +import com.google.cloud.bigquery.StandardTableDefinition; +import com.google.cloud.bigquery.TableId; +import com.google.cloud.bigquery.TableInfo; +import com.google.cloud.bigquery.TableResult; +import com.google.cloud.bigquery.storage.v1beta2.AppendRowsResponse; +import com.google.cloud.bigquery.storage.v1beta2.AppendRowsResponse.AppendResult; +import com.google.cloud.bigquery.storage.v1beta2.BigDecimalByteStringEncoder; +import com.google.cloud.bigquery.storage.v1beta2.BigQueryWriteClient; +import com.google.cloud.bigquery.storage.v1beta2.JsonStreamWriter; +import com.google.cloud.bigquery.storage.v1beta2.TableFieldSchema; +import com.google.cloud.bigquery.storage.v1beta2.TableName; +import com.google.cloud.bigquery.storage.v1beta2.TableSchema; +import com.google.cloud.bigquery.testing.RemoteBigQueryHelper; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.math.BigDecimal; +import java.util.Iterator; +import java.util.concurrent.ExecutionException; +import java.util.logging.Logger; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +public class ITBigQueryBigDecimalByteStringEncoderTest { + private static final Logger LOG = + Logger.getLogger(ITBigQueryBigDecimalByteStringEncoderTest.class.getName()); + private static final String DATASET = RemoteBigQueryHelper.generateDatasetName(); + private static final String TABLE = "testtable"; + private static final String DESCRIPTION = "BigQuery Write Java manual client test dataset"; + + private static BigQueryWriteClient client; + private static TableInfo tableInfo; + private static BigQuery bigquery; + + @BeforeClass + public static void beforeClass() throws IOException { + client = BigQueryWriteClient.create(); + + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + bigquery = bigqueryHelper.getOptions().getService(); + DatasetInfo datasetInfo = + DatasetInfo.newBuilder(/* datasetId = */ DATASET).setDescription(DESCRIPTION).build(); + bigquery.create(datasetInfo); + tableInfo = + TableInfo.newBuilder( + TableId.of(DATASET, TABLE), + StandardTableDefinition.of( + Schema.of( + com.google.cloud.bigquery.Field.newBuilder( + "test_numeric_zero", StandardSQLTypeName.NUMERIC) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_numeric_one", StandardSQLTypeName.NUMERIC) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_numeric_repeated", StandardSQLTypeName.NUMERIC) + .setMode(Mode.REPEATED) + .build()))) + .build(); + bigquery.create(tableInfo); + } + + @AfterClass + public static void afterClass() { + if (client != null) { + client.close(); + } + if (bigquery != null) { + RemoteBigQueryHelper.forceDelete(bigquery, DATASET); + } + } + + @Test + public void TestBigDecimalEncoding() + throws IOException, InterruptedException, ExecutionException, + Descriptors.DescriptorValidationException { + TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, TABLE); + TableFieldSchema TEST_NUMERIC_ZERO = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_numeric_zero") + .build(); + TableFieldSchema TEST_NUMERIC_ONE = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_numeric_one") + .build(); + TableFieldSchema TEST_NUMERIC_REPEATED = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_numeric_repeated") + .build(); + TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, TEST_NUMERIC_ZERO) + .addFields(1, TEST_NUMERIC_ONE) + .addFields(2, TEST_NUMERIC_REPEATED) + .build(); + try (JsonStreamWriter jsonStreamWriter = + JsonStreamWriter.newBuilder(parent.toString(), tableSchema).build()) { + JSONObject row = new JSONObject(); + row.put( + "test_numeric_zero", + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("0"))); + row.put( + "test_numeric_one", + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("1.2"))); + row.put( + "test_numeric_repeated", + new JSONArray( + new byte[][] { + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("0")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("1.2")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("-1.2")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("99999999999999999999999999999.999999999")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString( + new BigDecimal("-99999999999999999999999999999.999999999")) + .toByteArray(), + })); + JSONArray jsonArr = new JSONArray(new JSONObject[] {row}); + ApiFuture response = jsonStreamWriter.append(jsonArr, -1); + AppendRowsResponse arr = response.get(); + AppendResult ar = arr.getAppendResult(); + boolean ho = ar.hasOffset(); + TableResult result = + bigquery.listTableData( + tableInfo.getTableId(), BigQuery.TableDataListOption.startIndex(0L)); + Iterator iter = result.getValues().iterator(); + FieldValueList currentRow; + currentRow = iter.next(); + assertEquals("0", currentRow.get(0).getStringValue()); + assertEquals("1.2", currentRow.get(1).getStringValue()); + assertEquals("0", currentRow.get(2).getRepeatedValue().get(0).getStringValue()); + assertEquals("1.2", currentRow.get(2).getRepeatedValue().get(1).getStringValue()); + assertEquals("-1.2", currentRow.get(2).getRepeatedValue().get(2).getStringValue()); + assertEquals( + "99999999999999999999999999999.999999999", + currentRow.get(2).getRepeatedValue().get(3).getStringValue()); + assertEquals( + "-99999999999999999999999999999.999999999", + currentRow.get(2).getRepeatedValue().get(4).getStringValue()); + } + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryTimeEncoderTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryTimeEncoderTest.java new file mode 100644 index 0000000000..d96c27c19d --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryTimeEncoderTest.java @@ -0,0 +1,190 @@ +/* + * Copyright 2021 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.storage.v1.it; + +import static org.junit.Assert.assertEquals; + +import com.google.api.core.ApiFuture; +import com.google.cloud.ServiceOptions; +import com.google.cloud.bigquery.BigQuery; +import com.google.cloud.bigquery.DatasetInfo; +import com.google.cloud.bigquery.Field.Mode; +import com.google.cloud.bigquery.FieldValueList; +import com.google.cloud.bigquery.Schema; +import com.google.cloud.bigquery.StandardSQLTypeName; +import com.google.cloud.bigquery.StandardTableDefinition; +import com.google.cloud.bigquery.TableId; +import com.google.cloud.bigquery.TableInfo; +import com.google.cloud.bigquery.TableResult; +import com.google.cloud.bigquery.storage.v1beta2.AppendRowsResponse; +import com.google.cloud.bigquery.storage.v1beta2.BigQueryWriteClient; +import com.google.cloud.bigquery.storage.v1beta2.CivilTimeEncoder; +import com.google.cloud.bigquery.storage.v1beta2.JsonStreamWriter; +import com.google.cloud.bigquery.storage.v1beta2.TableFieldSchema; +import com.google.cloud.bigquery.storage.v1beta2.TableName; +import com.google.cloud.bigquery.storage.v1beta2.TableSchema; +import com.google.cloud.bigquery.testing.RemoteBigQueryHelper; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.util.Iterator; +import java.util.concurrent.ExecutionException; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; +import org.threeten.bp.LocalDateTime; +import org.threeten.bp.LocalTime; + +public class ITBigQueryTimeEncoderTest { + private static final String DATASET = RemoteBigQueryHelper.generateDatasetName(); + private static final String TABLE = "testtable"; + private static final String DESCRIPTION = "BigQuery Write Java manual client test dataset"; + + private static BigQueryWriteClient client; + private static TableInfo tableInfo; + private static BigQuery bigquery; + + @BeforeClass + public static void beforeClass() throws IOException { + client = BigQueryWriteClient.create(); + + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + bigquery = bigqueryHelper.getOptions().getService(); + DatasetInfo datasetInfo = + DatasetInfo.newBuilder(/* datasetId = */ DATASET).setDescription(DESCRIPTION).build(); + bigquery.create(datasetInfo); + tableInfo = + TableInfo.newBuilder( + TableId.of(DATASET, TABLE), + StandardTableDefinition.of( + Schema.of( + com.google.cloud.bigquery.Field.newBuilder( + "test_str", StandardSQLTypeName.STRING) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_time_micros", StandardSQLTypeName.TIME) + .setMode(Mode.REPEATED) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_datetime_micros", StandardSQLTypeName.DATETIME) + .setMode(Mode.REPEATED) + .build()))) + .build(); + bigquery.create(tableInfo); + } + + @AfterClass + public static void afterClass() { + if (client != null) { + client.close(); + } + if (bigquery != null) { + RemoteBigQueryHelper.forceDelete(bigquery, DATASET); + } + } + + @Test + public void TestTimeEncoding() + throws IOException, InterruptedException, ExecutionException, + Descriptors.DescriptorValidationException { + TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, TABLE); + TableFieldSchema TEST_STRING = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_str") + .build(); + TableFieldSchema TEST_TIME = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.TIME) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_time_micros") + .build(); + TableFieldSchema TEST_DATETIME = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DATETIME) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_datetime_micros") + .build(); + TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, TEST_STRING) + .addFields(1, TEST_TIME) + .addFields(2, TEST_DATETIME) + .build(); + try (JsonStreamWriter jsonStreamWriter = + JsonStreamWriter.newBuilder(parent.toString(), tableSchema).build()) { + JSONObject row = new JSONObject(); + row.put("test_str", "Start of the day"); + row.put( + "test_time_micros", + new JSONArray( + new long[] { + CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(13, 14, 15, 16_000_000)), + CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(23, 59, 59, 999_999_000)), + CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(0, 0, 0, 0)), + CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(1, 2, 3, 4_000)), + CivilTimeEncoder.encodePacked64TimeMicros(LocalTime.of(5, 6, 7, 8_000)) + })); + row.put( + "test_datetime_micros", + new JSONArray( + new long[] { + CivilTimeEncoder.encodePacked64DatetimeMicros( + LocalDateTime.of(1, 1, 1, 12, 0, 0, 0)), + CivilTimeEncoder.encodePacked64DatetimeMicros( + LocalDateTime.of(1995, 5, 19, 10, 30, 45, 0)), + CivilTimeEncoder.encodePacked64DatetimeMicros( + LocalDateTime.of(2000, 1, 1, 0, 0, 0, 0)), + CivilTimeEncoder.encodePacked64DatetimeMicros( + LocalDateTime.of(2026, 3, 11, 5, 45, 12, 9_000_000)), + CivilTimeEncoder.encodePacked64DatetimeMicros( + LocalDateTime.of(2050, 1, 2, 3, 4, 5, 6_000)), + })); + JSONArray jsonArr = new JSONArray(new JSONObject[] {row}); + ApiFuture response = jsonStreamWriter.append(jsonArr, -1); + Assert.assertFalse(response.get().getAppendResult().hasOffset()); + TableResult result = + bigquery.listTableData( + tableInfo.getTableId(), BigQuery.TableDataListOption.startIndex(0L)); + Iterator iter = result.getValues().iterator(); + FieldValueList currentRow; + currentRow = iter.next(); + assertEquals("Start of the day", currentRow.get(0).getValue()); + assertEquals("13:14:15.016000", currentRow.get(1).getRepeatedValue().get(0).getStringValue()); + assertEquals("23:59:59.999999", currentRow.get(1).getRepeatedValue().get(1).getStringValue()); + assertEquals("00:00:00", currentRow.get(1).getRepeatedValue().get(2).getStringValue()); + assertEquals("01:02:03.000004", currentRow.get(1).getRepeatedValue().get(3).getStringValue()); + assertEquals("05:06:07.000008", currentRow.get(1).getRepeatedValue().get(4).getStringValue()); + + assertEquals( + "0001-01-01T12:00:00", currentRow.get(2).getRepeatedValue().get(0).getStringValue()); + assertEquals( + "1995-05-19T10:30:45", currentRow.get(2).getRepeatedValue().get(1).getStringValue()); + assertEquals( + "2000-01-01T00:00:00", currentRow.get(2).getRepeatedValue().get(2).getStringValue()); + assertEquals( + "2026-03-11T05:45:12.009000", + currentRow.get(2).getRepeatedValue().get(3).getStringValue()); + assertEquals( + "2050-01-02T03:04:05.000006", + currentRow.get(2).getRepeatedValue().get(4).getStringValue()); + } + } +} diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java new file mode 100644 index 0000000000..7d06ec1e53 --- /dev/null +++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java @@ -0,0 +1,531 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.storage.v1.it; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.api.core.ApiFuture; +import com.google.cloud.ServiceOptions; +import com.google.cloud.bigquery.*; +import com.google.cloud.bigquery.Schema; +import com.google.cloud.bigquery.storage.test.Test.*; +import com.google.cloud.bigquery.storage.v1.*; +import com.google.cloud.bigquery.testing.RemoteBigQueryHelper; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.math.BigDecimal; +import java.util.*; +import java.util.concurrent.*; +import java.util.logging.Logger; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; +import org.threeten.bp.LocalDateTime; + +/** Integration tests for BigQuery Write API. */ +public class ITBigQueryWriteManualClientTest { + private static final Logger LOG = + Logger.getLogger(ITBigQueryWriteManualClientTest.class.getName()); + private static final String DATASET = RemoteBigQueryHelper.generateDatasetName(); + private static final String DATASET_EU = RemoteBigQueryHelper.generateDatasetName(); + private static final String TABLE = "testtable"; + private static final String TABLE2 = "complicatedtable"; + private static final String DESCRIPTION = "BigQuery Write Java manual client test dataset"; + + private static BigQueryWriteClient client; + private static TableInfo tableInfo; + private static TableInfo tableInfo2; + private static TableInfo tableInfoEU; + private static String tableId; + private static String tableId2; + private static String tableIdEU; + private static BigQuery bigquery; + + @BeforeClass + public static void beforeClass() throws IOException { + client = BigQueryWriteClient.create(); + + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + bigquery = bigqueryHelper.getOptions().getService(); + DatasetInfo datasetInfo = + DatasetInfo.newBuilder(/* datasetId = */ DATASET).setDescription(DESCRIPTION).build(); + bigquery.create(datasetInfo); + LOG.info("Created test dataset: " + DATASET); + tableInfo = + TableInfo.newBuilder( + TableId.of(DATASET, TABLE), + StandardTableDefinition.of( + Schema.of( + com.google.cloud.bigquery.Field.newBuilder("foo", LegacySQLTypeName.STRING) + .setMode(Field.Mode.NULLABLE) + .build()))) + .build(); + com.google.cloud.bigquery.Field.Builder innerTypeFieldBuilder = + com.google.cloud.bigquery.Field.newBuilder( + "inner_type", + LegacySQLTypeName.RECORD, + com.google.cloud.bigquery.Field.newBuilder("value", LegacySQLTypeName.STRING) + .setMode(Field.Mode.REPEATED) + .build()); + + tableInfo2 = + TableInfo.newBuilder( + TableId.of(DATASET, TABLE2), + StandardTableDefinition.of( + Schema.of( + Field.newBuilder( + "nested_repeated_type", + LegacySQLTypeName.RECORD, + innerTypeFieldBuilder.setMode(Field.Mode.REPEATED).build()) + .setMode(Field.Mode.REPEATED) + .build(), + innerTypeFieldBuilder.setMode(Field.Mode.NULLABLE).build()))) + .build(); + bigquery.create(tableInfo); + bigquery.create(tableInfo2); + tableId = + String.format( + "projects/%s/datasets/%s/tables/%s", + ServiceOptions.getDefaultProjectId(), DATASET, TABLE); + tableId2 = + String.format( + "projects/%s/datasets/%s/tables/%s", + ServiceOptions.getDefaultProjectId(), DATASET, TABLE2); + DatasetInfo datasetInfoEU = + DatasetInfo.newBuilder(/* datasetId = */ DATASET_EU) + .setLocation("EU") + .setDescription(DESCRIPTION) + .build(); + bigquery.create(datasetInfoEU); + tableInfoEU = + TableInfo.newBuilder( + TableId.of(DATASET_EU, TABLE), + StandardTableDefinition.of( + Schema.of( + com.google.cloud.bigquery.Field.newBuilder("foo", LegacySQLTypeName.STRING) + .build()))) + .build(); + tableIdEU = + String.format( + "projects/%s/datasets/%s/tables/%s", + ServiceOptions.getDefaultProjectId(), DATASET_EU, TABLE); + bigquery.create(tableInfoEU); + } + + @AfterClass + public static void afterClass() { + if (client != null) { + client.close(); + } + + if (bigquery != null) { + RemoteBigQueryHelper.forceDelete(bigquery, DATASET); + LOG.info("Deleted test dataset: " + DATASET); + } + } + + ProtoRows CreateProtoRows(String[] messages) { + ProtoRows.Builder rows = ProtoRows.newBuilder(); + for (String message : messages) { + FooType foo = FooType.newBuilder().setFoo(message).build(); + rows.addSerializedRows(foo.toByteString()); + } + return rows.build(); + } + + ProtoRows CreateProtoRowsComplex(String[] messages) { + ProtoRows.Builder rows = ProtoRows.newBuilder(); + for (String message : messages) { + ComplicateType foo = + ComplicateType.newBuilder() + .setInnerType(InnerType.newBuilder().addValue(message).addValue(message).build()) + .build(); + rows.addSerializedRows(foo.toByteString()); + } + return rows.build(); + } + + @Test + public void testBatchWriteWithCommittedStreamEU() + throws IOException, InterruptedException, ExecutionException { + WriteStream writeStream = + client.createWriteStream( + CreateWriteStreamRequest.newBuilder() + .setParent(tableIdEU) + .setWriteStream( + WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build()) + .build()); + StreamWriter streamWriter = + StreamWriter.newBuilder(writeStream.getName()) + .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor())) + .build(); + LOG.info("Sending one message"); + + ApiFuture response = + streamWriter.append(CreateProtoRows(new String[] {"aaa"}), 0); + assertEquals(0, response.get().getAppendResult().getOffset().getValue()); + + LOG.info("Sending two more messages"); + ApiFuture response1 = + streamWriter.append(CreateProtoRows(new String[] {"bbb", "ccc"}), 1); + ApiFuture response2 = + streamWriter.append(CreateProtoRows(new String[] {"ddd"}), 3); + assertEquals(1, response1.get().getAppendResult().getOffset().getValue()); + assertEquals(3, response2.get().getAppendResult().getOffset().getValue()); + + TableResult result = + bigquery.listTableData( + tableInfoEU.getTableId(), BigQuery.TableDataListOption.startIndex(0L)); + Iterator iter = result.getValues().iterator(); + assertEquals("aaa", iter.next().get(0).getStringValue()); + assertEquals("bbb", iter.next().get(0).getStringValue()); + assertEquals("ccc", iter.next().get(0).getStringValue()); + assertEquals("ddd", iter.next().get(0).getStringValue()); + assertEquals(false, iter.hasNext()); + } + + @Test + public void testJsonStreamWriterCommittedStream() + throws IOException, InterruptedException, ExecutionException, + Descriptors.DescriptorValidationException { + String tableName = "JsonTable"; + TableInfo tableInfo = + TableInfo.newBuilder( + TableId.of(DATASET, tableName), + StandardTableDefinition.of( + Schema.of( + com.google.cloud.bigquery.Field.newBuilder( + "test_str", StandardSQLTypeName.STRING) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_numerics", StandardSQLTypeName.NUMERIC) + .setMode(Field.Mode.REPEATED) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_datetime", StandardSQLTypeName.DATETIME) + .build()))) + .build(); + bigquery.create(tableInfo); + TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, tableName); + WriteStream writeStream = + client.createWriteStream( + CreateWriteStreamRequest.newBuilder() + .setParent(parent.toString()) + .setWriteStream( + WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build()) + .build()); + try (JsonStreamWriter jsonStreamWriter = + JsonStreamWriter.newBuilder(writeStream.getName(), writeStream.getTableSchema()).build()) { + LOG.info("Sending one message"); + JSONObject row1 = new JSONObject(); + row1.put("test_str", "aaa"); + row1.put( + "test_numerics", + new JSONArray( + new byte[][] { + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("123.4")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("-9000000")) + .toByteArray() + })); + row1.put( + "test_datetime", + CivilTimeEncoder.encodePacked64DatetimeMicros(LocalDateTime.of(2020, 10, 1, 12, 0))); + JSONArray jsonArr1 = new JSONArray(new JSONObject[] {row1}); + + ApiFuture response1 = jsonStreamWriter.append(jsonArr1, -1); + + assertEquals(0, response1.get().getAppendResult().getOffset().getValue()); + + JSONObject row2 = new JSONObject(); + row1.put("test_str", "bbb"); + JSONObject row3 = new JSONObject(); + row2.put("test_str", "ccc"); + JSONArray jsonArr2 = new JSONArray(); + jsonArr2.put(row1); + jsonArr2.put(row2); + + JSONObject row4 = new JSONObject(); + row4.put("test_str", "ddd"); + JSONArray jsonArr3 = new JSONArray(); + jsonArr3.put(row4); + + LOG.info("Sending two more messages"); + ApiFuture response2 = jsonStreamWriter.append(jsonArr2, -1); + LOG.info("Sending one more message"); + ApiFuture response3 = jsonStreamWriter.append(jsonArr3, -1); + assertEquals(1, response2.get().getAppendResult().getOffset().getValue()); + assertEquals(3, response3.get().getAppendResult().getOffset().getValue()); + + TableResult result = + bigquery.listTableData( + tableInfo.getTableId(), BigQuery.TableDataListOption.startIndex(0L)); + Iterator iter = result.getValues().iterator(); + FieldValueList currentRow = iter.next(); + assertEquals("aaa", currentRow.get(0).getStringValue()); + assertEquals("-9000000", currentRow.get(1).getRepeatedValue().get(1).getStringValue()); + assertEquals("2020-10-01T12:00:00", currentRow.get(2).getStringValue()); + assertEquals("bbb", iter.next().get(0).getStringValue()); + assertEquals("ccc", iter.next().get(0).getStringValue()); + assertEquals("ddd", iter.next().get(0).getStringValue()); + assertEquals(false, iter.hasNext()); + } + } + + @Test + public void testJsonStreamWriterWithDefaultStream() + throws IOException, InterruptedException, ExecutionException, + Descriptors.DescriptorValidationException { + String tableName = "JsonTableDefaultStream"; + TableFieldSchema TEST_STRING = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.STRING) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_str") + .build(); + TableFieldSchema TEST_NUMERIC = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.NUMERIC) + .setMode(TableFieldSchema.Mode.REPEATED) + .setName("test_numerics") + .build(); + TableFieldSchema TEST_DATE = + TableFieldSchema.newBuilder() + .setType(TableFieldSchema.Type.DATETIME) + .setMode(TableFieldSchema.Mode.NULLABLE) + .setName("test_datetime") + .build(); + TableSchema tableSchema = + TableSchema.newBuilder() + .addFields(0, TEST_STRING) + .addFields(1, TEST_DATE) + .addFields(2, TEST_NUMERIC) + .build(); + TableInfo tableInfo = + TableInfo.newBuilder( + TableId.of(DATASET, tableName), + StandardTableDefinition.of( + Schema.of( + com.google.cloud.bigquery.Field.newBuilder( + "test_str", StandardSQLTypeName.STRING) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_numerics", StandardSQLTypeName.NUMERIC) + .setMode(Field.Mode.REPEATED) + .build(), + com.google.cloud.bigquery.Field.newBuilder( + "test_datetime", StandardSQLTypeName.DATETIME) + .build()))) + .build(); + bigquery.create(tableInfo); + TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, tableName); + try (JsonStreamWriter jsonStreamWriter = + JsonStreamWriter.newBuilder(parent.toString(), tableSchema).build()) { + LOG.info("Sending one message"); + JSONObject row1 = new JSONObject(); + row1.put("test_str", "aaa"); + row1.put( + "test_numerics", + new JSONArray( + new byte[][] { + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("123.4")) + .toByteArray(), + BigDecimalByteStringEncoder.encodeToNumericByteString(new BigDecimal("-9000000")) + .toByteArray() + })); + row1.put( + "test_datetime", + CivilTimeEncoder.encodePacked64DatetimeMicros(LocalDateTime.of(2020, 10, 1, 12, 0))); + JSONArray jsonArr1 = new JSONArray(new JSONObject[] {row1}); + + ApiFuture response1 = jsonStreamWriter.append(jsonArr1, -1); + + assertEquals(0, response1.get().getAppendResult().getOffset().getValue()); + + JSONObject row2 = new JSONObject(); + row1.put("test_str", "bbb"); + JSONObject row3 = new JSONObject(); + row2.put("test_str", "ccc"); + JSONArray jsonArr2 = new JSONArray(); + jsonArr2.put(row1); + jsonArr2.put(row2); + + JSONObject row4 = new JSONObject(); + row4.put("test_str", "ddd"); + JSONArray jsonArr3 = new JSONArray(); + jsonArr3.put(row4); + + LOG.info("Sending two more messages"); + ApiFuture response2 = jsonStreamWriter.append(jsonArr2, -1); + LOG.info("Sending one more message"); + ApiFuture response3 = jsonStreamWriter.append(jsonArr3, -1); + Assert.assertFalse(response2.get().getAppendResult().hasOffset()); + Assert.assertFalse(response3.get().getAppendResult().hasOffset()); + + TableResult result = + bigquery.listTableData( + tableInfo.getTableId(), BigQuery.TableDataListOption.startIndex(0L)); + Iterator iter = result.getValues().iterator(); + FieldValueList currentRow = iter.next(); + assertEquals("aaa", currentRow.get(0).getStringValue()); + assertEquals("-9000000", currentRow.get(1).getRepeatedValue().get(1).getStringValue()); + assertEquals("2020-10-01T12:00:00", currentRow.get(2).getStringValue()); + assertEquals("bbb", iter.next().get(0).getStringValue()); + assertEquals("ccc", iter.next().get(0).getStringValue()); + assertEquals("ddd", iter.next().get(0).getStringValue()); + assertEquals(false, iter.hasNext()); + } + } + + @Test + public void testComplicateSchemaWithPendingStream() + throws IOException, InterruptedException, ExecutionException { + LOG.info("Create a write stream"); + WriteStream writeStream = + client.createWriteStream( + CreateWriteStreamRequest.newBuilder() + .setParent(tableId2) + .setWriteStream(WriteStream.newBuilder().setType(WriteStream.Type.PENDING).build()) + .build()); + FinalizeWriteStreamResponse finalizeResponse = FinalizeWriteStreamResponse.getDefaultInstance(); + try (StreamWriter streamWriter = + StreamWriter.newBuilder(writeStream.getName()) + .setWriterSchema(ProtoSchemaConverter.convert(ComplicateType.getDescriptor())) + .build()) { + LOG.info("Sending two messages"); + ApiFuture response = + streamWriter.append(CreateProtoRowsComplex(new String[] {"aaa"}), 0L); + assertEquals(0, response.get().getAppendResult().getOffset().getValue()); + + ApiFuture response2 = + streamWriter.append(CreateProtoRowsComplex(new String[] {"bbb"}), 1L); + assertEquals(1, response2.get().getAppendResult().getOffset().getValue()); + + // Nothing showed up since rows are not committed. + TableResult result = + bigquery.listTableData( + tableInfo2.getTableId(), BigQuery.TableDataListOption.startIndex(0L)); + Iterator iter = result.getValues().iterator(); + assertEquals(false, iter.hasNext()); + + LOG.info("Finalize a write stream"); + finalizeResponse = + client.finalizeWriteStream( + FinalizeWriteStreamRequest.newBuilder().setName(writeStream.getName()).build()); + + ApiFuture response3 = + streamWriter.append(CreateProtoRows(new String[] {"ccc"}), 2L); + try { + response3.get(); + Assert.fail("Append to finalized stream should fail."); + } catch (Exception expected) { + LOG.info("Got exception: " + expected.toString()); + } + } + assertEquals(2, finalizeResponse.getRowCount()); + LOG.info("Commit a write stream"); + BatchCommitWriteStreamsResponse batchCommitWriteStreamsResponse = + client.batchCommitWriteStreams( + BatchCommitWriteStreamsRequest.newBuilder() + .setParent(tableId2) + .addWriteStreams(writeStream.getName()) + .build()); + assertEquals(true, batchCommitWriteStreamsResponse.hasCommitTime()); + TableResult queryResult = + bigquery.query( + QueryJobConfiguration.newBuilder("SELECT * from " + DATASET + '.' + TABLE2).build()); + Iterator queryIter = queryResult.getValues().iterator(); + assertTrue(queryIter.hasNext()); + assertEquals( + "[FieldValue{attribute=REPEATED, value=[FieldValue{attribute=PRIMITIVE, value=aaa}, FieldValue{attribute=PRIMITIVE, value=aaa}]}]", + queryIter.next().get(1).getRepeatedValue().toString()); + assertEquals( + "[FieldValue{attribute=REPEATED, value=[FieldValue{attribute=PRIMITIVE, value=bbb}, FieldValue{attribute=PRIMITIVE, value=bbb}]}]", + queryIter.next().get(1).getRepeatedValue().toString()); + assertFalse(queryIter.hasNext()); + } + + @Test + public void testStreamError() throws IOException, InterruptedException, ExecutionException { + WriteStream writeStream = + client.createWriteStream( + CreateWriteStreamRequest.newBuilder() + .setParent(tableId) + .setWriteStream( + WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build()) + .build()); + try (StreamWriter streamWriter = + StreamWriter.newBuilder(writeStream.getName()) + .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor())) + .build()) { + ApiFuture response = + streamWriter.append(CreateProtoRows(new String[] {"aaa"}), -1L); + assertEquals(0L, response.get().getAppendResult().getOffset().getValue()); + // Send in a bogus stream name should cause in connection error. + ApiFuture response2 = + streamWriter.append(CreateProtoRows(new String[] {"aaa"}), 100L); + try { + response2.get(); + Assert.fail("Should fail"); + } catch (ExecutionException e) { + assertThat(e.getCause().getMessage()) + .contains("OUT_OF_RANGE: The offset is beyond stream, expected offset 1, received 100"); + } + // We can keep sending requests on the same stream. + ApiFuture response3 = + streamWriter.append(CreateProtoRows(new String[] {"aaa"}), -1L); + assertEquals(1L, response3.get().getAppendResult().getOffset().getValue()); + } finally { + } + } + + @Test + public void testStreamReconnect() throws IOException, InterruptedException, ExecutionException { + WriteStream writeStream = + client.createWriteStream( + CreateWriteStreamRequest.newBuilder() + .setParent(tableId) + .setWriteStream( + WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build()) + .build()); + try (StreamWriter streamWriter = + StreamWriter.newBuilder(writeStream.getName()) + .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor())) + .build()) { + ApiFuture response = + streamWriter.append(CreateProtoRows(new String[] {"aaa"}), 0L); + assertEquals(0L, response.get().getAppendResult().getOffset().getValue()); + } + + try (StreamWriter streamWriter = + StreamWriter.newBuilder(writeStream.getName()) + .setWriterSchema(ProtoSchemaConverter.convert(FooType.getDescriptor())) + .build()) { + // Currently there is a bug that reconnection must wait 5 seconds to get the real row count. + Thread.sleep(5000L); + ApiFuture response = + streamWriter.append(CreateProtoRows(new String[] {"bbb"}), 1L); + assertEquals(1L, response.get().getAppendResult().getOffset().getValue()); + } + } +}