Skip to content

Commit

Permalink
Cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
Evan Greco committed Oct 30, 2023
1 parent 51f1f51 commit 8c82567
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 102 deletions.
85 changes: 14 additions & 71 deletions ...c/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteNonQuotaRetryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,73 +16,32 @@

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.assertNotNull;
import static org.junit.Assert.assertTrue;

import com.google.api.core.ApiFuture;
import com.google.api.gax.retrying.RetrySettings;
import com.google.cloud.ServiceOptions;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.DatasetInfo;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.LegacySQLTypeName;
import com.google.cloud.bigquery.QueryJobConfiguration;
import com.google.cloud.bigquery.Schema;
import com.google.cloud.bigquery.StandardSQLTypeName;
import com.google.cloud.bigquery.StandardTableDefinition;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.TableInfo;
import com.google.cloud.bigquery.TableResult;
import com.google.cloud.bigquery.storage.test.Test.ComplicateType;
import com.google.cloud.bigquery.storage.test.Test.FooTimestampType;
import com.google.cloud.bigquery.storage.test.Test.FooType;
import com.google.cloud.bigquery.storage.test.Test.InnerType;
import com.google.cloud.bigquery.storage.test.Test.UpdatedFooType;
import com.google.cloud.bigquery.storage.v1.AppendRowsResponse;
import com.google.cloud.bigquery.storage.v1.BatchCommitWriteStreamsRequest;
import com.google.cloud.bigquery.storage.v1.BatchCommitWriteStreamsResponse;
import com.google.cloud.bigquery.storage.v1.BigDecimalByteStringEncoder;
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
import com.google.cloud.bigquery.storage.v1.CivilTimeEncoder;
import com.google.cloud.bigquery.storage.v1.ConnectionWorkerPool;
import com.google.cloud.bigquery.storage.v1.CreateWriteStreamRequest;
import com.google.cloud.bigquery.storage.v1.Exceptions.AppendSerializationError;
import com.google.cloud.bigquery.storage.v1.Exceptions.OffsetAlreadyExists;
import com.google.cloud.bigquery.storage.v1.Exceptions.OffsetOutOfRange;
import com.google.cloud.bigquery.storage.v1.Exceptions.SchemaMismatchedException;
import com.google.cloud.bigquery.storage.v1.Exceptions.StreamFinalizedException;
import com.google.cloud.bigquery.storage.v1.FinalizeWriteStreamRequest;
import com.google.cloud.bigquery.storage.v1.FinalizeWriteStreamResponse;
import com.google.cloud.bigquery.storage.v1.JsonStreamWriter;
import com.google.cloud.bigquery.storage.v1.ProtoRows;
import com.google.cloud.bigquery.storage.v1.ProtoSchemaConverter;
import com.google.cloud.bigquery.storage.v1.StreamWriter;
import com.google.cloud.bigquery.storage.v1.TableFieldSchema;
import com.google.cloud.bigquery.storage.v1.TableName;
import com.google.cloud.bigquery.storage.v1.TableSchema;
import com.google.cloud.bigquery.storage.v1.WriteStream;
import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import io.grpc.Status;
import io.grpc.Status.Code;
import java.io.IOException;
import java.math.BigDecimal;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Logger;
import org.json.JSONArray;
import org.json.JSONObject;
Expand All @@ -91,21 +50,17 @@
import org.junit.BeforeClass;
import org.junit.Test;
import org.threeten.bp.Duration;
import org.threeten.bp.LocalDateTime;

/** Integration tests for BigQuery Write API. */
public class ITBigQueryWriteNonQuotaRetryTest {
private static final Logger LOG =
Logger.getLogger(ITBigQueryWriteNonQuotaRetryTest.class.getName());
private static final String DATASET = 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 final String NON_QUOTA_RETRY_PROJECT_ID = "bq-write-api-java-retry-test";

private static BigQueryWriteClient client;
private static TableInfo tableInfo;
private static String tableId;
private static BigQuery bigquery;

@BeforeClass
Expand All @@ -118,27 +73,15 @@ public static void beforeClass() throws IOException {
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(
Field.newBuilder("foo", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build())))
.build();
Field.Builder innerTypeFieldBuilder =
Field.newBuilder(
"inner_type",
LegacySQLTypeName.RECORD,
Field.newBuilder("value", LegacySQLTypeName.STRING)
.setMode(Field.Mode.REPEATED)
.build());
TableInfo tableInfo = TableInfo.newBuilder(
TableId.of(DATASET, TABLE),
StandardTableDefinition.of(
Schema.of(
Field.newBuilder("foo", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build())))
.build();
bigquery.create(tableInfo);
tableId =
String.format(
"projects/%s/datasets/%s/tables/%s",
ServiceOptions.getDefaultProjectId(), DATASET, TABLE);
}

@AfterClass
Expand All @@ -155,8 +98,8 @@ public static void afterClass() {

@Test
public void testJsonStreamWriterCommittedStreamWithNonQuotaRetry()
throws IOException, InterruptedException, ExecutionException,
DescriptorValidationException {
throws IOException, InterruptedException,
DescriptorValidationException {
RetrySettings retrySettings =
RetrySettings.newBuilder()
.setInitialRetryDelay(Duration.ofMillis(500))
Expand All @@ -182,7 +125,7 @@ public void testJsonStreamWriterCommittedStreamWithNonQuotaRetry()
int totalRequest = 901;
int rowBatch = 1;
ArrayList<ApiFuture<AppendRowsResponse>> allResponses =
new ArrayList<ApiFuture<AppendRowsResponse>>(totalRequest);
new ArrayList<>(totalRequest);
try (JsonStreamWriter jsonStreamWriter =
JsonStreamWriter.newBuilder(writeStream.getName(), writeStream.getTableSchema())
.setRetrySettings(retrySettings)
Expand Down Expand Up @@ -213,8 +156,8 @@ public void testJsonStreamWriterCommittedStreamWithNonQuotaRetry()

@Test
public void testJsonStreamWriterDefaultStreamWithNonQuotaRetry()
throws IOException, InterruptedException, ExecutionException,
DescriptorValidationException {
throws IOException, InterruptedException,
DescriptorValidationException {
RetrySettings retrySettings =
RetrySettings.newBuilder()
.setInitialRetryDelay(Duration.ofMillis(500))
Expand Down Expand Up @@ -246,7 +189,7 @@ public void testJsonStreamWriterDefaultStreamWithNonQuotaRetry()
int totalRequest = 901;
int rowBatch = 1;
ArrayList<ApiFuture<AppendRowsResponse>> allResponses =
new ArrayList<ApiFuture<AppendRowsResponse>>(totalRequest);
new ArrayList<>(totalRequest);
try (JsonStreamWriter jsonStreamWriter =
JsonStreamWriter.newBuilder(parent.toString(), tableSchema)
.setIgnoreUnknownFields(true)
Expand Down
45 changes: 14 additions & 31 deletions .../src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteQuotaRetryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@

import com.google.api.core.ApiFuture;
import com.google.api.gax.retrying.RetrySettings;
import com.google.cloud.ServiceOptions;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.DatasetInfo;
import com.google.cloud.bigquery.Field;
Expand Down Expand Up @@ -58,13 +57,9 @@ public class ITBigQueryWriteQuotaRetryTest {
Logger.getLogger(ITBigQueryWriteQuotaRetryTest.class.getName());
private static final String DATASET = 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 final String QUOTA_RETRY_PROJECT_ID = "bq-writeapi-java-quota-retry";

private static BigQueryWriteClient client;
private static TableInfo tableInfo;
private static String tableId;
private static BigQuery bigquery;

@BeforeClass
Expand All @@ -77,27 +72,15 @@ public static void beforeClass() throws IOException {
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(
Field.newBuilder("foo", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build())))
.build();
Field.Builder innerTypeFieldBuilder =
Field.newBuilder(
"inner_type",
LegacySQLTypeName.RECORD,
Field.newBuilder("value", LegacySQLTypeName.STRING)
.setMode(Field.Mode.REPEATED)
.build());
TableInfo tableInfo = TableInfo.newBuilder(
TableId.of(DATASET, TABLE),
StandardTableDefinition.of(
Schema.of(
Field.newBuilder("foo", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build())))
.build();
bigquery.create(tableInfo);
tableId =
String.format(
"projects/%s/datasets/%s/tables/%s",
ServiceOptions.getDefaultProjectId(), DATASET, TABLE);
}

@AfterClass
Expand All @@ -114,8 +97,8 @@ public static void afterClass() {

@Test
public void testJsonStreamWriterCommittedStreamWithNonQuotaRetry()
throws IOException, InterruptedException, ExecutionException,
DescriptorValidationException {
throws IOException, InterruptedException,
DescriptorValidationException {
RetrySettings retrySettings =
RetrySettings.newBuilder()
.setInitialRetryDelay(Duration.ofMillis(500))
Expand All @@ -141,7 +124,7 @@ public void testJsonStreamWriterCommittedStreamWithNonQuotaRetry()
int totalRequest = 901;
int rowBatch = 1;
ArrayList<ApiFuture<AppendRowsResponse>> allResponses =
new ArrayList<ApiFuture<AppendRowsResponse>>(totalRequest);
new ArrayList<>(totalRequest);
try (JsonStreamWriter jsonStreamWriter =
JsonStreamWriter.newBuilder(writeStream.getName(), writeStream.getTableSchema())
.setRetrySettings(retrySettings)
Expand Down Expand Up @@ -172,8 +155,8 @@ public void testJsonStreamWriterCommittedStreamWithNonQuotaRetry()

@Test
public void testJsonStreamWriterDefaultStreamWithNonQuotaRetry()
throws IOException, InterruptedException, ExecutionException,
DescriptorValidationException {
throws IOException, InterruptedException,
DescriptorValidationException {
RetrySettings retrySettings =
RetrySettings.newBuilder()
.setInitialRetryDelay(Duration.ofMillis(500))
Expand Down Expand Up @@ -205,7 +188,7 @@ public void testJsonStreamWriterDefaultStreamWithNonQuotaRetry()
int totalRequest = 901;
int rowBatch = 1;
ArrayList<ApiFuture<AppendRowsResponse>> allResponses =
new ArrayList<ApiFuture<AppendRowsResponse>>(totalRequest);
new ArrayList<>(totalRequest);
try (JsonStreamWriter jsonStreamWriter =
JsonStreamWriter.newBuilder(parent.toString(), tableSchema)
.setIgnoreUnknownFields(true)
Expand Down

0 comments on commit 8c82567

Please sign in to comment.