connectionProperties);
+
+ /**
+ * Sets the default dataset. This dataset is used for all unqualified table names used in the
+ * query.
+ */
+ public abstract Builder setDefaultDataset(DatasetId datasetId);
+
+ /**
+ * Limits the bytes billed for this job. Queries that will have bytes billed beyond this limit
+ * will fail (without incurring a charge). If unspecified, this will be set to your project
+ * default.
+ *
+ * @param maximumBytesBilled maximum bytes billed for this job
+ */
+ public abstract Builder setMaximumBytesBilled(Long maximumBytesBilled);
+
+ /**
+ * Sets the maximum number of rows of data to return per page of results. Setting this flag to a
+ * small value such as 1000 and then paging through results might improve reliability when the
+ * query result set is large. In addition to this limit, responses are also limited to 10 MB. By
+ * default, there is no maximum row count, and only the byte limit applies.
+ *
+ * @param maxResults maxResults or {@code null} for none
+ */
+ public abstract Builder setMaxResults(Long maxResults);
+
+ /**
+ * Sets the number of rows in the buffer (a blocking queue) that query results are consumed
+ * from.
+ *
+ * @param numBufferedRows numBufferedRows or {@code null} for none
+ */
+ public abstract Builder setNumBufferedRows(Integer numBufferedRows);
+
+ /**
+ * Sets a ratio of the total number of records and the records returned in the current page.
+ * This value is checked before calling the Read API
+ *
+ * @param totalToPageRowCountRatio totalToPageRowCountRatio
+ */
+ public abstract Builder setTotalToPageRowCountRatio(Integer totalToPageRowCountRatio);
+
+ /**
+ * Sets the minimum result size for which the Read API will be enabled
+ *
+ * @param minResultSize minResultSize
+ */
+ public abstract Builder setMinResultSize(Integer minResultSize);
+
+ /**
+ * Sets the maximum records per page to be used for pagination. This is used as an input for the
+ * tabledata.list and jobs.getQueryResults RPC calls
+ *
+ * @param maxResultPerPage
+ */
+ public abstract Builder setMaxResultPerPage(Integer maxResultPerPage);
+
+ /**
+ * Sets whether to look for the result in the query cache. The query cache is a best-effort
+ * cache that will be flushed whenever tables in the query are modified. Moreover, the query
+ * cache is only available when {@link ConnectionSettings.Builder#setDestinationTable(TableId)}
+ * is not set.
+ *
+ * @see Query Caching
+ */
+ public abstract Builder setUseQueryCache(Boolean useQueryCache);
+
+ /**
+ * Sets whether nested and repeated fields should be flattened. If set to {@code false} {@link
+ * ConnectionSettings.Builder#setAllowLargeResults(Boolean)} must be {@code true}. By default
+ * results are flattened.
+ *
+ * @see Flatten
+ */
+ public abstract Builder setFlattenResults(Boolean flattenResults);
+
+ /* */
+ /**/
+ /**
+ * Sets the values necessary to determine whether table result will be read using the BigQuery
+ * Storage client Read API. The BigQuery Storage client Read API will be used to read the query
+ * result when the totalToFirstPageSizeRatio (default 3) and minimumTableSize (default 100 rows)
+ * conditions set are met. A ReadSession will be created using the Apache Arrow data format for
+ * serialization.
+ *
+ * It also sets the maximum number of table rows allowed in buffer before streaming them to
+ * the BigQueryResult.
+ *
+ * @param readClientConnectionConfiguration or {@code null} for none
+ */
+ /*
+ public abstract Builder setReadClientConnectionConfiguration(
+ ReadClientConnectionConfiguration readClientConnectionConfiguration);*/
+
+ /** Sets the clustering specification for the destination table. */
+ public abstract Builder setClustering(Clustering clustering);
+
+ /**
+ * Sets whether the job is allowed to create tables.
+ *
+ * @see
+ * Create Disposition
+ */
+ public abstract Builder setCreateDisposition(CreateDisposition createDisposition);
+
+ /**
+ * Sets the custom encryption configuration (e.g., Cloud KMS keys).
+ *
+ * @param destinationEncryptionConfiguration destinationEncryptionConfiguration or {@code null}
+ * for none
+ */
+ public abstract Builder setDestinationEncryptionConfiguration(
+ EncryptionConfiguration destinationEncryptionConfiguration);
+
+ /**
+ * Sets the table where to put query results. If not provided a new table is created. This value
+ * is required if {@link ConnectionSettings.Builder#setAllowLargeResults(Boolean)} is set to
+ * {@code true}.
+ */
+ public abstract Builder setDestinationTable(TableId destinationTable);
+
+ /**
+ * [Optional] Job timeout in milliseconds. If this time limit is exceeded, BigQuery may attempt
+ * to terminate the job.
+ *
+ * @param jobTimeoutMs jobTimeoutMs or {@code null} for none
+ */
+ public abstract Builder setJobTimeoutMs(Long jobTimeoutMs);
+
+ /**
+ * Limits the billing tier for this job. Queries that have resource usage beyond this tier will
+ * fail (without incurring a charge). If unspecified, this will be set to your project default.
+ *
+ * @param maximumBillingTier maximum billing tier for this job
+ */
+ public abstract Builder setMaximumBillingTier(Integer maximumBillingTier);
+
+ /**
+ * Sets a priority for the query. If not specified the priority is assumed to be {@link
+ * Priority#INTERACTIVE}.
+ */
+ public abstract Builder setPriority(Priority priority);
+
+ /**
+ * Sets whether the job is enabled to create arbitrarily large results. If {@code true} the
+ * query is allowed to create large results at a slight cost in performance. If {@code true}
+ * {@link ConnectionSettings.Builder#setDestinationTable(TableId)} must be provided.
+ *
+ * @see Returning
+ * Large Query Results
+ */
+ public abstract Builder setAllowLargeResults(Boolean allowLargeResults);
+
+ /**
+ * Sets whether to create a new session. If {@code true} a random session id will be generated
+ * by BigQuery. If false, runs query with an existing session_id passed in ConnectionProperty,
+ * otherwise runs query in non-session mode."
+ */
+ public abstract Builder setCreateSession(Boolean createSession);
+
+ /**
+ * Range partitioning specification for this table. Only one of timePartitioning and
+ * rangePartitioning should be specified.
+ *
+ * @param rangePartitioning rangePartitioning or {@code null} for none
+ */
+ public abstract Builder setRangePartitioning(RangePartitioning rangePartitioning);
+
+ /**
+ * [Experimental] Sets options allowing the schema of the destination table to be updated as a
+ * side effect of the query job. Schema update options are supported in two cases: when
+ * writeDisposition is WRITE_APPEND; when writeDisposition is WRITE_TRUNCATE and the destination
+ * table is a partition of a table, specified by partition decorators. For normal tables,
+ * WRITE_TRUNCATE will always overwrite the schema.
+ */
+ public abstract Builder setSchemaUpdateOptions(List schemaUpdateOptions);
+
+ /**
+ * Sets the external tables definitions. If querying external data sources outside of BigQuery,
+ * this value describes the data format, location and other properties of the data sources. By
+ * defining these properties, the data sources can be queried as if they were standard BigQuery
+ * tables.
+ */
+ public abstract Builder setTableDefinitions(
+ Map tableDefinitions);
+
+ /** Sets the time partitioning specification for the destination table. */
+ public abstract Builder setTimePartitioning(TimePartitioning timePartitioning);
+
+ /**
+ * Sets user defined function resources that can be used by this query. Function resources can
+ * either be defined inline ({@link UserDefinedFunction#inline(String)}) or loaded from a Google
+ * Cloud Storage URI ({@link UserDefinedFunction#fromUri(String)}.
+ */
+ public abstract Builder setUserDefinedFunctions(List userDefinedFunctions);
+
+ /**
+ * Sets the action that should occur if the destination table already exists.
+ *
+ * @see
+ * Write Disposition
+ */
+ public abstract Builder setWriteDisposition(WriteDisposition writeDisposition);
+
+ /** Creates a {@code ConnectionSettings} object. */
+ public abstract ConnectionSettings build();
+ }
+}
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
index ab9fdabb3..0ef1d1f94 100644
--- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
@@ -21,6 +21,7 @@
import com.google.api.services.bigquery.model.JobStatistics2;
import com.google.api.services.bigquery.model.JobStatistics3;
import com.google.api.services.bigquery.model.JobStatistics4;
+import com.google.api.services.bigquery.model.QueryParameter;
import com.google.cloud.StringEnumType;
import com.google.cloud.StringEnumValue;
import com.google.common.base.Function;
@@ -339,6 +340,7 @@ public static class QueryStatistics extends JobStatistics {
private final List queryPlan;
private final List timeline;
private final Schema schema;
+ private final List queryParameters;
/**
* StatementType represents possible types of SQL statements reported as part of the
@@ -421,6 +423,7 @@ static final class Builder extends JobStatistics.Builder queryPlan;
private List timeline;
private Schema schema;
+ private List queryParameters;
private Builder() {}
@@ -569,6 +572,11 @@ Builder setSchema(Schema schema) {
return self();
}
+ Builder setQueryParameters(List queryParameters) {
+ this.queryParameters = queryParameters;
+ return self();
+ }
+
@Override
QueryStatistics build() {
return new QueryStatistics(this);
@@ -595,6 +603,7 @@ private QueryStatistics(Builder builder) {
this.queryPlan = builder.queryPlan;
this.timeline = builder.timeline;
this.schema = builder.schema;
+ this.queryParameters = builder.queryParameters;
}
/** Returns query statistics specific to the use of BI Engine. */
@@ -715,6 +724,14 @@ public Schema getSchema() {
return schema;
}
+ /**
+ * Standard SQL only: Returns a list of undeclared query parameters detected during a dry run
+ * validation.
+ */
+ public List getQueryParameters() {
+ return queryParameters;
+ }
+
@Override
ToStringHelper toStringHelper() {
return super.toStringHelper()
@@ -725,7 +742,8 @@ ToStringHelper toStringHelper() {
.add("totalBytesProcessed", totalBytesProcessed)
.add("queryPlan", queryPlan)
.add("timeline", timeline)
- .add("schema", schema);
+ .add("schema", schema)
+ .add("queryParameters", queryParameters);
}
@Override
@@ -746,7 +764,8 @@ public final int hashCode() {
totalBytesBilled,
totalBytesProcessed,
queryPlan,
- schema);
+ schema,
+ queryParameters);
}
@Override
@@ -788,6 +807,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() {
if (schema != null) {
queryStatisticsPb.setSchema(schema.toPb());
}
+ if (queryParameters != null) {
+ queryStatisticsPb.setUndeclaredQueryParameters(queryParameters);
+ }
return super.toPb().setQuery(queryStatisticsPb);
}
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Parameter.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Parameter.java
new file mode 100644
index 000000000..9959feab9
--- /dev/null
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Parameter.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright 2022 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;
+
+import com.google.auto.value.AutoValue;
+import javax.annotation.Nullable;
+
+/* Wrapper class for query parameters */
+@AutoValue
+public abstract class Parameter {
+ Parameter() {
+ // Package private so users can't subclass it but AutoValue can.
+ }
+
+ /**
+ * Returns the name of the query parameter. If unset, this is a positional parameter. Otherwise,
+ * should be unique within a query.
+ *
+ * @return value or {@code null} for none
+ */
+ @Nullable
+ public abstract String getName();
+
+ /** Returns the value for a query parameter along with its type. */
+ public abstract QueryParameterValue getValue();
+
+ /** Returns a builder pre-populated using the current values of this field. */
+ public abstract Builder toBuilder();
+
+ /** Returns a builder for a {@code Parameter} object. */
+ public static Builder newBuilder() {
+ return new AutoValue_Parameter.Builder();
+ }
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+
+ /**
+ * [Optional] Sets the name of the query parameter. If unset, this is a positional parameter.
+ * Otherwise, should be unique within a query.
+ *
+ * @param name name or {@code null} for none
+ */
+ public abstract Builder setName(String name);
+
+ /**
+ * Sets the the value for a query parameter along with its type.
+ *
+ * @param parameter parameter or {@code null} for none
+ */
+ public abstract Builder setValue(QueryParameterValue parameter);
+
+ /** Creates a {@code Parameter} object. */
+ public abstract Parameter build();
+ }
+}
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java
index 48ec22caf..cc726bdd1 100644
--- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java
@@ -166,6 +166,11 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
this();
JobConfigurationQuery queryConfigurationPb = configurationPb.getQuery();
this.query = queryConfigurationPb.getQuery();
+ // Allows to get undeclaredqueryparameters in jobstatistics2
+ if (queryConfigurationPb.getQueryParameters() == null
+ && queryConfigurationPb.getParameterMode() != null) {
+ parameterMode = queryConfigurationPb.getParameterMode();
+ }
if (queryConfigurationPb.getQueryParameters() != null
&& !queryConfigurationPb.getQueryParameters().isEmpty()) {
if (queryConfigurationPb.getQueryParameters().get(0).getName() == null) {
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ReadClientConnectionConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ReadClientConnectionConfiguration.java
new file mode 100644
index 000000000..e0805a11e
--- /dev/null
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ReadClientConnectionConfiguration.java
@@ -0,0 +1,70 @@
+/*
+ * 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;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import javax.annotation.Nullable;
+
+/** Represents BigQueryStorage Read client connection information. */
+@AutoValue
+public abstract class ReadClientConnectionConfiguration implements Serializable {
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+
+ /**
+ * Sets the total row count to page row count ratio used to determine whether to us the
+ * BigQueryStorage Read client to fetch result sets after the first page.
+ */
+ @Nullable
+ public abstract Builder setTotalToPageRowCountRatio(Long ratio);
+
+ /**
+ * Sets the minimum number of table rows in the query results used to determine whether to us
+ * the BigQueryStorage Read client to fetch result sets after the first page.
+ */
+ @Nullable
+ public abstract Builder setMinResultSize(Long numRows);
+
+ /**
+ * Sets the maximum number of table rows allowed in buffer before streaming them to the
+ * BigQueryResult.
+ */
+ @Nullable
+ public abstract Builder setBufferSize(Long bufferSize);
+
+ /** Creates a {@code ReadClientConnectionConfiguration} object. */
+ public abstract ReadClientConnectionConfiguration build();
+ }
+
+ /** Returns the totalToPageRowCountRatio in this configuration. */
+ public abstract Long getTotalToPageRowCountRatio();
+
+ /** Returns the minResultSize in this configuration. */
+ public abstract Long getMinResultSize();
+
+ /** Returns the bufferSize in this configuration. */
+ public abstract Long getBufferSize();
+
+ public abstract Builder toBuilder();
+
+ /** Returns a builder for a {@code ReadClientConnectionConfiguration} object. */
+ public static Builder newBuilder() {
+ return new AutoValue_ReadClientConnectionConfiguration.Builder();
+ }
+}
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/BigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/BigQueryRpc.java
index 06488c5b4..871590ca4 100644
--- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/BigQueryRpc.java
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/BigQueryRpc.java
@@ -122,6 +122,13 @@ Boolean getBoolean(Map options) {
*/
Job create(Job job, Map options);
+ /**
+ * Creates a new query job.
+ *
+ * @throws BigQueryException upon failure
+ */
+ Job createJobForQuery(Job job);
+
/**
* Delete the requested dataset.
*
@@ -246,6 +253,14 @@ TableDataInsertAllResponse insertAll(
TableDataList listTableData(
String projectId, String datasetId, String tableId, Map options);
+ /**
+ * Lists the table's rows with a limit on how many rows of data to pre-fetch.
+ *
+ * @throws BigQueryException upon failure
+ */
+ TableDataList listTableDataWithRowLimit(
+ String projectId, String datasetId, String tableId, Integer rowLimit, String pageToken);
+
/**
* Returns the requested job or {@code null} if not found.
*
@@ -253,6 +268,13 @@ TableDataList listTableData(
*/
Job getJob(String projectId, String jobId, String location, Map options);
+ /**
+ * Returns the requested query job or {@code null} if not found.
+ *
+ * @throws BigQueryException upon failure
+ */
+ Job getQueryJob(String projectId, String jobId, String location);
+
/**
* Lists the project's jobs.
*
@@ -286,6 +308,15 @@ TableDataList listTableData(
GetQueryResultsResponse getQueryResults(
String projectId, String jobId, String location, Map options);
+ /**
+ * Returns results of the query with a limit on how many rows of data to pre-fetch associated with
+ * the provided job.
+ *
+ * @throws BigQueryException upon failure
+ */
+ GetQueryResultsResponse getQueryResultsWithRowLimit(
+ String projectId, String jobId, String location, Integer preFetchedRowLimit);
+
/**
* Runs a BigQuery SQL query synchronously and returns query results if the query completes within
* a specified timeout.
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
index 24d7dd6b0..d6b57a3da 100644
--- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
@@ -221,6 +221,19 @@ public Job create(Job job, Map options) {
}
}
+ @Override
+ public Job createJobForQuery(Job job) {
+ try {
+ String projectId =
+ job.getJobReference() != null
+ ? job.getJobReference().getProjectId()
+ : this.options.getProjectId();
+ return bigquery.jobs().insert(projectId, job).setPrettyPrint(false).execute();
+ } catch (IOException ex) {
+ throw translate(ex);
+ }
+ }
+
@Override
public boolean deleteDataset(String projectId, String datasetId, Map options) {
try {
@@ -515,6 +528,26 @@ public TableDataList listTableData(
}
}
+ @Override
+ public TableDataList listTableDataWithRowLimit(
+ String projectId,
+ String datasetId,
+ String tableId,
+ Integer maxResultPerPage,
+ String pageToken) {
+ try {
+ return bigquery
+ .tabledata()
+ .list(projectId, datasetId, tableId)
+ .setPrettyPrint(false)
+ .setMaxResults(Long.valueOf(maxResultPerPage))
+ .setPageToken(pageToken)
+ .execute();
+ } catch (IOException ex) {
+ throw translate(ex);
+ }
+ }
+
@Override
public Job getJob(String projectId, String jobId, String location, Map options) {
try {
@@ -534,6 +567,24 @@ public Job getJob(String projectId, String jobId, String location, Map> listJobs(String projectId, Map options) {
try {
@@ -644,6 +695,22 @@ public GetQueryResultsResponse getQueryResults(
}
}
+ @Override
+ public GetQueryResultsResponse getQueryResultsWithRowLimit(
+ String projectId, String jobId, String location, Integer maxResultPerPage) {
+ try {
+ return bigquery
+ .jobs()
+ .getQueryResults(projectId, jobId)
+ .setPrettyPrint(false)
+ .setLocation(location)
+ .setMaxResults(Long.valueOf(maxResultPerPage))
+ .execute();
+ } catch (IOException ex) {
+ throw translate(ex);
+ }
+ }
+
@Override
public QueryResponse queryRpc(String projectId, QueryRequest content) {
try {
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java
new file mode 100644
index 000000000..e4fdc9731
--- /dev/null
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionImplTest.java
@@ -0,0 +1,542 @@
+/*
+ * 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;
+
+import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.verify;
+
+import com.google.api.services.bigquery.model.*;
+import com.google.api.services.bigquery.model.QueryResponse;
+import com.google.cloud.ServiceOptions;
+import com.google.cloud.Tuple;
+import com.google.cloud.bigquery.spi.BigQueryRpcFactory;
+import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
+import com.google.common.collect.ImmutableList;
+import java.math.BigInteger;
+import java.sql.SQLException;
+import java.util.AbstractList;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingDeque;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class ConnectionImplTest {
+ private BigQueryOptions options;
+ private BigQueryRpcFactory rpcFactoryMock;
+ private BigQueryRpc bigqueryRpcMock;
+ private Connection connectionMock;
+ private BigQuery bigquery;
+ private ConnectionImpl connection;
+ private static final String PROJECT = "project";
+ private static final String JOB = "job";
+ private static final String LOCATION = "US";
+ private static final String DEFAULT_TEST_DATASET = "bigquery_test_dataset";
+ private static final String PAGE_TOKEN = "ABCD123";
+ private static final TableId TABLE_NAME = TableId.of(DEFAULT_TEST_DATASET, PROJECT);
+ private static final TableCell STRING_CELL = new TableCell().setV("Value");
+ private static final TableRow TABLE_ROW = new TableRow().setF(ImmutableList.of(STRING_CELL));
+ private static final String SQL_QUERY =
+ "SELECT county, state_name FROM bigquery_test_dataset.large_data_testing_table limit 2";
+ private static final String DRY_RUN_SQL =
+ "SELECT county, state_name FROM bigquery_test_dataset.large_data_testing_table where country = ?";
+ private static final int DEFAULT_PAGE_SIZE = 10000;
+ private ConnectionSettings connectionSettings;
+ private static final Schema QUERY_SCHEMA =
+ Schema.of(
+ Field.newBuilder("country", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("state_name", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .build());
+ private static final TableSchema FAST_QUERY_TABLESCHEMA = QUERY_SCHEMA.toPb();
+ private static final BigQueryResult BQ_RS_MOCK_RES =
+ new BigQueryResultImpl(QUERY_SCHEMA, 2, null, null);
+
+ private static final BigQueryResult BQ_RS_MOCK_RES_MULTI_PAGE =
+ new BigQueryResultImpl(QUERY_SCHEMA, 4, null, null);
+
+ private static final JobId QUERY_JOB = JobId.of(PROJECT, JOB).setLocation(LOCATION);
+ private static final GetQueryResultsResponse GET_QUERY_RESULTS_RESPONSE =
+ new GetQueryResultsResponse()
+ .setJobReference(QUERY_JOB.toPb())
+ .setRows(ImmutableList.of(TABLE_ROW))
+ .setJobComplete(true)
+ .setCacheHit(false)
+ .setPageToken(PAGE_TOKEN)
+ .setTotalBytesProcessed(42L)
+ .setTotalRows(BigInteger.valueOf(1L))
+ .setSchema(FAST_QUERY_TABLESCHEMA);
+
+ private BigQueryOptions createBigQueryOptionsForProject(
+ String project, BigQueryRpcFactory rpcFactory) {
+ return BigQueryOptions.newBuilder()
+ .setProjectId(project)
+ .setServiceRpcFactory(rpcFactory)
+ .setRetrySettings(ServiceOptions.getNoRetrySettings())
+ .build();
+ }
+
+ @Before
+ public void setUp() {
+ rpcFactoryMock = mock(BigQueryRpcFactory.class);
+ bigqueryRpcMock = mock(BigQueryRpc.class);
+ connectionMock = mock(Connection.class);
+ when(rpcFactoryMock.create(any(BigQueryOptions.class))).thenReturn(bigqueryRpcMock);
+ options = createBigQueryOptionsForProject(PROJECT, rpcFactoryMock);
+ bigquery = options.getService();
+
+ connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DEFAULT_TEST_DATASET))
+ .setNumBufferedRows(DEFAULT_PAGE_SIZE)
+ .build();
+ bigquery =
+ options
+ .toBuilder()
+ .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
+ .build()
+ .getService();
+ connection = (ConnectionImpl) bigquery.createConnection(connectionSettings);
+ assertNotNull(connection);
+ }
+
+ @Test
+ public void testFastQuerySinglePage() throws BigQuerySQLException {
+ com.google.api.services.bigquery.model.QueryResponse mockQueryRes =
+ new QueryResponse().setSchema(FAST_QUERY_TABLESCHEMA).setJobComplete(true);
+ when(bigqueryRpcMock.queryRpc(any(String.class), any(QueryRequest.class)))
+ .thenReturn(mockQueryRes);
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .processQueryResponseResults(any(QueryResponse.class));
+
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(connectionSpy, times(1))
+ .processQueryResponseResults(
+ any(com.google.api.services.bigquery.model.QueryResponse.class));
+ }
+
+ @Test
+ // NOTE: This doesn't truly paginates. Returns a response while mocking
+ // processQueryResponseResults
+ public void testFastQueryMultiplePages() throws BigQuerySQLException {
+ com.google.api.services.bigquery.model.QueryResponse mockQueryRes =
+ new QueryResponse()
+ .setSchema(FAST_QUERY_TABLESCHEMA)
+ .setJobComplete(true)
+ .setPageToken(PAGE_TOKEN);
+ when(bigqueryRpcMock.queryRpc(any(String.class), any(QueryRequest.class)))
+ .thenReturn(mockQueryRes);
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+
+ doReturn(BQ_RS_MOCK_RES_MULTI_PAGE)
+ .when(connectionSpy)
+ .processQueryResponseResults(
+ any(com.google.api.services.bigquery.model.QueryResponse.class));
+
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY);
+ assertEquals(res.getTotalRows(), 4);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(connectionSpy, times(1))
+ .processQueryResponseResults(
+ any(com.google.api.services.bigquery.model.QueryResponse.class));
+ }
+
+ @Test
+ public void testClose() throws BigQuerySQLException {
+ boolean cancelled = connection.close();
+ assertTrue(cancelled);
+ }
+
+ @Test
+ public void testQueryDryRun() throws BigQuerySQLException {
+ List queryParametersMock =
+ ImmutableList.of(
+ new QueryParameter().setParameterType(new QueryParameterType().setType("STRING")));
+ com.google.api.services.bigquery.model.JobStatistics2 queryMock =
+ new com.google.api.services.bigquery.model.JobStatistics2()
+ .setSchema(FAST_QUERY_TABLESCHEMA)
+ .setUndeclaredQueryParameters(queryParametersMock);
+ com.google.api.services.bigquery.model.JobStatistics jobStatsMock =
+ new com.google.api.services.bigquery.model.JobStatistics()
+ .setCreationTime(1234L)
+ .setStartTime(5678L)
+ .setQuery(queryMock);
+ com.google.api.services.bigquery.model.JobConfigurationQuery jobConfigurationQuery =
+ new com.google.api.services.bigquery.model.JobConfigurationQuery();
+ com.google.api.services.bigquery.model.JobConfiguration jobConfig =
+ new com.google.api.services.bigquery.model.JobConfiguration()
+ .setQuery(jobConfigurationQuery);
+ com.google.api.services.bigquery.model.Job mockDryRunJob =
+ new com.google.api.services.bigquery.model.Job()
+ .setStatistics(jobStatsMock)
+ .setConfiguration(jobConfig);
+ when(bigqueryRpcMock.createJobForQuery(any(com.google.api.services.bigquery.model.Job.class)))
+ .thenReturn(mockDryRunJob);
+ BigQueryDryRunResult dryRunResult = connection.dryRun(DRY_RUN_SQL);
+ assertEquals(1, dryRunResult.getQueryParameters().size());
+ assertEquals(QUERY_SCHEMA, dryRunResult.getSchema());
+ verify(bigqueryRpcMock, times(1))
+ .createJobForQuery(any(com.google.api.services.bigquery.model.Job.class));
+ }
+
+ @Test
+ public void testParseDataTask() throws InterruptedException {
+ List tableRows =
+ ImmutableList.of(
+ new TableRow()
+ .setF(
+ ImmutableList.of(
+ new TableCell().setV("Value1"), new TableCell().setV("Value2"))),
+ new TableRow()
+ .setF(
+ ImmutableList.of(
+ new TableCell().setV("Value3"), new TableCell().setV("Value4"))));
+
+ BlockingQueue, Boolean>> pageCache =
+ new LinkedBlockingDeque<>(2);
+ BlockingQueue> rpcResponseQueue = new LinkedBlockingDeque<>(2);
+ rpcResponseQueue.offer(Tuple.of(null, false));
+ // This call should populate page cache
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ connectionSpy.parseRpcDataAsync(tableRows, QUERY_SCHEMA, pageCache, rpcResponseQueue);
+ Tuple, Boolean> fvlTupple =
+ pageCache.take(); // wait for the parser thread to parse the data
+ assertNotNull(fvlTupple);
+ Iterable iterableFvl = fvlTupple.x();
+ int rowCnt = 0;
+ for (FieldValueList fvl : iterableFvl) {
+ assertEquals(2, fvl.size()); // both the rows should have 2 fields each
+ rowCnt++;
+ }
+ assertEquals(2, rowCnt); // row rows read
+
+ verify(connectionSpy, times(1))
+ .parseRpcDataAsync(
+ any(List.class), any(Schema.class), any(BlockingQueue.class), any(BlockingQueue.class));
+ }
+
+ @Test
+ public void testPopulateBuffer() throws InterruptedException {
+ List tableRows =
+ ImmutableList.of(
+ new TableRow()
+ .setF(
+ ImmutableList.of(
+ new TableCell().setV("Value1"), new TableCell().setV("Value2"))),
+ new TableRow()
+ .setF(
+ ImmutableList.of(
+ new TableCell().setV("Value3"), new TableCell().setV("Value4"))));
+
+ BlockingQueue, Boolean>> pageCache =
+ new LinkedBlockingDeque<>(2);
+ BlockingQueue> rpcResponseQueue = new LinkedBlockingDeque<>(2);
+ BlockingQueue> buffer = new LinkedBlockingDeque<>(5);
+ rpcResponseQueue.offer(Tuple.of(null, false));
+ // This call should populate page cache
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+
+ connectionSpy.parseRpcDataAsync(tableRows, QUERY_SCHEMA, pageCache, rpcResponseQueue);
+
+ verify(connectionSpy, times(1))
+ .parseRpcDataAsync(
+ any(List.class), any(Schema.class), any(BlockingQueue.class), any(BlockingQueue.class));
+
+ // now pass the pageCache to populateBuffer method
+ connectionSpy.populateBufferAsync(rpcResponseQueue, pageCache, buffer);
+ // check if buffer was populated with two rows async by using the blocking take method
+ AbstractList fvl1 = buffer.take();
+ assertNotNull(fvl1);
+ assertEquals(2, fvl1.size());
+ assertEquals("Value1", fvl1.get(0).getValue().toString());
+ assertEquals("Value2", fvl1.get(1).getValue().toString());
+ AbstractList fvl2 = buffer.take();
+ assertNotNull(fvl2);
+ assertEquals(2, fvl2.size());
+ assertEquals("Value3", fvl2.get(0).getValue().toString());
+ assertEquals("Value4", fvl2.get(1).getValue().toString());
+ verify(connectionSpy, times(1))
+ .populateBufferAsync(
+ any(BlockingQueue.class), any(BlockingQueue.class), any(BlockingQueue.class));
+ }
+
+ @Test
+ public void testNextPageTask() throws InterruptedException {
+ BlockingQueue> rpcResponseQueue = new LinkedBlockingDeque<>(2);
+ TableDataList mockTabledataList =
+ new TableDataList()
+ .setPageToken(PAGE_TOKEN)
+ .setRows(ImmutableList.of(TABLE_ROW))
+ .setTotalRows(1L);
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ doReturn(mockTabledataList)
+ .when(connectionSpy)
+ .tableDataListRpc(any(TableId.class), any(String.class));
+ connectionSpy.runNextPageTaskAsync(PAGE_TOKEN, TABLE_NAME, rpcResponseQueue);
+ Tuple tableDataListTuple = rpcResponseQueue.take();
+ assertNotNull(tableDataListTuple);
+ TableDataList tableDataList = tableDataListTuple.x();
+ assertNotNull(tableDataList);
+ assertEquals("ABCD123", tableDataList.getPageToken());
+ assertEquals(Long.valueOf(1), tableDataList.getTotalRows());
+ verify(connectionSpy, times(1))
+ .runNextPageTaskAsync(any(String.class), any(TableId.class), any(BlockingQueue.class));
+ }
+
+ @Test
+ public void testGetQueryResultsFirstPage() {
+ when(bigqueryRpcMock.getQueryResultsWithRowLimit(
+ any(String.class), any(String.class), any(String.class), any(Integer.class)))
+ .thenReturn(GET_QUERY_RESULTS_RESPONSE);
+ GetQueryResultsResponse response = connection.getQueryResultsFirstPage(QUERY_JOB);
+ assertNotNull(response);
+ assertEquals(GET_QUERY_RESULTS_RESPONSE, response);
+ verify(bigqueryRpcMock, times(1))
+ .getQueryResultsWithRowLimit(
+ any(String.class), any(String.class), any(String.class), any(Integer.class));
+ }
+
+ // calls executeSelect with a nonFast query and exercises createQueryJob
+ @Test
+ public void testLegacyQuerySinglePage() throws BigQuerySQLException {
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ com.google.api.services.bigquery.model.Job jobResponseMock =
+ new com.google.api.services.bigquery.model.Job()
+ // .setConfiguration(QUERY_JOB.g)
+ .setJobReference(QUERY_JOB.toPb())
+ .setId(JOB)
+ .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"));
+ // emulating a legacy query
+ doReturn(false).when(connectionSpy).isFastQuerySupported();
+ doReturn(GET_QUERY_RESULTS_RESPONSE)
+ .when(connectionSpy)
+ .getQueryResultsFirstPage(any(JobId.class));
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .getSubsequentQueryResultsWithJob(
+ any(Long.class),
+ any(Long.class),
+ any(JobId.class),
+ any(GetQueryResultsResponse.class),
+ any(Boolean.class));
+ when(bigqueryRpcMock.createJobForQuery(any(com.google.api.services.bigquery.model.Job.class)))
+ .thenReturn(jobResponseMock); // RPC call in createQueryJob
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(bigqueryRpcMock, times(1))
+ .createJobForQuery(any(com.google.api.services.bigquery.model.Job.class));
+ }
+
+ // exercises getSubsequentQueryResultsWithJob for fast running queries
+ @Test
+ public void testFastQueryLongRunning() throws SQLException {
+ List tableRows =
+ ImmutableList.of(
+ new TableRow()
+ .setF(
+ ImmutableList.of(
+ new TableCell().setV("Value1"), new TableCell().setV("Value2"))),
+ new TableRow()
+ .setF(
+ ImmutableList.of(
+ new TableCell().setV("Value3"), new TableCell().setV("Value4"))));
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ // emulating a fast query
+ doReturn(true).when(connectionSpy).isFastQuerySupported();
+ doReturn(GET_QUERY_RESULTS_RESPONSE)
+ .when(connectionSpy)
+ .getQueryResultsFirstPage(any(JobId.class));
+
+ doReturn(TABLE_NAME).when(connectionSpy).getDestinationTable(any(JobId.class));
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .tableDataList(any(GetQueryResultsResponse.class), any(JobId.class));
+
+ com.google.api.services.bigquery.model.QueryResponse mockQueryRes =
+ new QueryResponse()
+ .setSchema(FAST_QUERY_TABLESCHEMA)
+ .setJobComplete(false)
+ .setTotalRows(new BigInteger(String.valueOf(4L)))
+ .setJobReference(QUERY_JOB.toPb())
+ .setRows(tableRows);
+ when(bigqueryRpcMock.queryRpc(any(String.class), any(QueryRequest.class)))
+ .thenReturn(mockQueryRes);
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(bigqueryRpcMock, times(1)).queryRpc(any(String.class), any(QueryRequest.class));
+ }
+
+ @Test
+ // Emulates first page response using getQueryResultsFirstPage(jobId) and then subsequent pages
+ // using getQueryResultsFirstPage(jobId) getSubsequentQueryResultsWithJob(
+ public void testLegacyQueryMultiplePages() throws SQLException {
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ com.google.api.services.bigquery.model.JobStatistics jobStatistics =
+ new com.google.api.services.bigquery.model.JobStatistics();
+ // emulating a Legacy query
+ doReturn(false).when(connectionSpy).isFastQuerySupported();
+ doReturn(GET_QUERY_RESULTS_RESPONSE)
+ .when(connectionSpy)
+ .getQueryResultsFirstPage(any(JobId.class));
+ doReturn(TABLE_NAME).when(connectionSpy).getDestinationTable(any(JobId.class));
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .tableDataList(any(GetQueryResultsResponse.class), any(JobId.class));
+ com.google.api.services.bigquery.model.Job jobResponseMock =
+ new com.google.api.services.bigquery.model.Job()
+ .setJobReference(QUERY_JOB.toPb())
+ .setId(JOB)
+ .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"))
+ .setStatistics(jobStatistics);
+ when(bigqueryRpcMock.createJobForQuery(any(com.google.api.services.bigquery.model.Job.class)))
+ .thenReturn(jobResponseMock); // RPC call in createQueryJob
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(bigqueryRpcMock, times(1))
+ .createJobForQuery(any(com.google.api.services.bigquery.model.Job.class));
+ verify(connectionSpy, times(1))
+ .tableDataList(any(GetQueryResultsResponse.class), any(JobId.class));
+ }
+
+ @Test
+ public void testExecuteSelectSlow() throws BigQuerySQLException {
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ doReturn(false).when(connectionSpy).isFastQuerySupported();
+ com.google.api.services.bigquery.model.JobStatistics jobStatistics =
+ new com.google.api.services.bigquery.model.JobStatistics();
+ com.google.api.services.bigquery.model.Job jobResponseMock =
+ new com.google.api.services.bigquery.model.Job()
+ .setJobReference(QUERY_JOB.toPb())
+ .setId(JOB)
+ .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"))
+ .setStatistics(jobStatistics);
+
+ doReturn(jobResponseMock)
+ .when(connectionSpy)
+ .createQueryJob(SQL_QUERY, connectionSettings, null, null);
+ doReturn(GET_QUERY_RESULTS_RESPONSE)
+ .when(connectionSpy)
+ .getQueryResultsFirstPage(any(JobId.class));
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .getResultSet(
+ any(GetQueryResultsResponse.class),
+ any(JobId.class),
+ any(String.class),
+ any(Boolean.class));
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(connectionSpy, times(1))
+ .getResultSet(
+ any(GetQueryResultsResponse.class),
+ any(JobId.class),
+ any(String.class),
+ any(Boolean.class));
+ }
+
+ @Test
+ public void testExecuteSelectSlowWithParams() throws BigQuerySQLException {
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ List parameters = new ArrayList<>();
+ Map labels = new HashMap<>();
+ doReturn(false).when(connectionSpy).isFastQuerySupported();
+ com.google.api.services.bigquery.model.JobStatistics jobStatistics =
+ new com.google.api.services.bigquery.model.JobStatistics();
+ com.google.api.services.bigquery.model.Job jobResponseMock =
+ new com.google.api.services.bigquery.model.Job()
+ .setJobReference(QUERY_JOB.toPb())
+ .setId(JOB)
+ .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"))
+ .setStatistics(jobStatistics);
+
+ doReturn(jobResponseMock)
+ .when(connectionSpy)
+ .createQueryJob(SQL_QUERY, connectionSettings, parameters, labels);
+ doReturn(GET_QUERY_RESULTS_RESPONSE)
+ .when(connectionSpy)
+ .getQueryResultsFirstPage(any(JobId.class));
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .getResultSet(
+ any(GetQueryResultsResponse.class),
+ any(JobId.class),
+ any(String.class),
+ any(Boolean.class));
+ BigQueryResult res = connectionSpy.executeSelect(SQL_QUERY, parameters, labels);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(connectionSpy, times(1))
+ .getResultSet(
+ any(GetQueryResultsResponse.class),
+ any(JobId.class),
+ any(String.class),
+ any(Boolean.class));
+ }
+
+ @Test
+ public void testGetSubsequentQueryResultsWithJob() {
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ JobId jobId = mock(JobId.class);
+ BigQueryResultStats bqRsStats = mock(BigQueryResultStats.class);
+ doReturn(true)
+ .when(connectionSpy)
+ .useReadAPI(any(Long.class), any(Long.class), any(Schema.class), any(Boolean.class));
+ doReturn(BQ_RS_MOCK_RES)
+ .when(connectionSpy)
+ .highThroughPutRead(
+ any(TableId.class), any(Long.class), any(Schema.class), any(BigQueryResultStats.class));
+
+ doReturn(TABLE_NAME).when(connectionSpy).getDestinationTable(any(JobId.class));
+ doReturn(bqRsStats).when(connectionSpy).getBigQueryResultSetStats(any(JobId.class));
+ BigQueryResult res =
+ connectionSpy.getSubsequentQueryResultsWithJob(
+ 10000L, 100L, jobId, GET_QUERY_RESULTS_RESPONSE, false);
+ assertEquals(res.getTotalRows(), 2);
+ assertEquals(QUERY_SCHEMA, res.getSchema());
+ verify(connectionSpy, times(1))
+ .getSubsequentQueryResultsWithJob(10000L, 100L, jobId, GET_QUERY_RESULTS_RESPONSE, false);
+ }
+
+ @Test
+ public void testGetPageCacheSize() {
+ ConnectionImpl connectionSpy = Mockito.spy(connection);
+ // number of cached pages should be within a range
+ assertTrue(connectionSpy.getPageCacheSize(10000, QUERY_SCHEMA) >= 3);
+ assertTrue(connectionSpy.getPageCacheSize(100000000, QUERY_SCHEMA) <= 20);
+ verify(connectionSpy, times(2)).getPageCacheSize(any(Integer.class), any(Schema.class));
+ }
+}
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionSettingsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionSettingsTest.java
new file mode 100644
index 000000000..8523825bc
--- /dev/null
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ConnectionSettingsTest.java
@@ -0,0 +1,166 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.cloud.bigquery.JobInfo.CreateDisposition;
+import com.google.cloud.bigquery.JobInfo.SchemaUpdateOption;
+import com.google.cloud.bigquery.JobInfo.WriteDisposition;
+import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.Test;
+
+public class ConnectionSettingsTest {
+ private static final String TEST_PROJECT_ID = "test-project-id";
+ private static final DatasetId DATASET_ID = DatasetId.of("dataset");
+ private static final TableId TABLE_ID = TableId.of("dataset", "table");
+ private static final Long REQUEST_TIMEOUT = 10l;
+ private static final Integer NUM_BUFFERED_ROWS = 100;
+ private static final Long MAX_RESULTS = 1000l;
+ private static final List SOURCE_URIS = ImmutableList.of("uri1", "uri2");
+ private static final String KEY = "time_zone";
+ private static final String VALUE = "US/Eastern";
+ private static final ConnectionProperty CONNECTION_PROPERTY =
+ ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build();
+ private static final List CONNECTION_PROPERTIES =
+ ImmutableList.of(CONNECTION_PROPERTY);
+ private static final Field FIELD_SCHEMA1 =
+ Field.newBuilder("StringField", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("FieldDescription1")
+ .build();
+ private static final Field FIELD_SCHEMA2 =
+ Field.newBuilder("IntegerField", StandardSQLTypeName.INT64)
+ .setMode(Field.Mode.REPEATED)
+ .setDescription("FieldDescription2")
+ .build();
+ private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA1, FIELD_SCHEMA2);
+ private static final Integer MAX_BAD_RECORDS = 42;
+ private static final Boolean IGNORE_UNKNOWN_VALUES = true;
+ private static final String COMPRESSION = "GZIP";
+ private static final CsvOptions CSV_OPTIONS = CsvOptions.newBuilder().build();
+ private static final ExternalTableDefinition TABLE_CONFIGURATION =
+ ExternalTableDefinition.newBuilder(SOURCE_URIS, TABLE_SCHEMA, CSV_OPTIONS)
+ .setCompression(COMPRESSION)
+ .setIgnoreUnknownValues(IGNORE_UNKNOWN_VALUES)
+ .setMaxBadRecords(MAX_BAD_RECORDS)
+ .build();
+ private static final Map TABLE_DEFINITIONS =
+ ImmutableMap.of("tableName", TABLE_CONFIGURATION);
+ private static final CreateDisposition CREATE_DISPOSITION = CreateDisposition.CREATE_IF_NEEDED;
+ private static final WriteDisposition WRITE_DISPOSITION = WriteDisposition.WRITE_APPEND;
+ private static final Priority PRIORITY = Priority.BATCH;
+ private static final boolean ALLOW_LARGE_RESULTS = true;
+ private static final boolean USE_QUERY_CACHE = false;
+ private static final boolean FLATTEN_RESULTS = true;
+ private static final Integer MAX_BILLING_TIER = 123;
+ private static final Long MAX_BYTES_BILL = 12345L;
+ private static final List SCHEMA_UPDATE_OPTIONS =
+ ImmutableList.of(SchemaUpdateOption.ALLOW_FIELD_RELAXATION);
+ private static final List USER_DEFINED_FUNCTIONS =
+ ImmutableList.of(UserDefinedFunction.inline("Function"), UserDefinedFunction.fromUri("URI"));
+ private static final EncryptionConfiguration JOB_ENCRYPTION_CONFIGURATION =
+ EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
+ private static final TimePartitioning TIME_PARTITIONING =
+ TimePartitioning.of(TimePartitioning.Type.DAY);
+ private static final Clustering CLUSTERING =
+ Clustering.newBuilder().setFields(ImmutableList.of("Foo", "Bar")).build();
+ private static final Long TIMEOUT = 10L;
+ private static final RangePartitioning.Range RANGE =
+ RangePartitioning.Range.newBuilder().setStart(1L).setInterval(2L).setEnd(10L).build();
+ private static final RangePartitioning RANGE_PARTITIONING =
+ RangePartitioning.newBuilder().setField("IntegerField").setRange(RANGE).build();
+
+ private static final ConnectionSettings CONNECTION_SETTINGS =
+ ConnectionSettings.newBuilder()
+ .setRequestTimeout(REQUEST_TIMEOUT)
+ .setNumBufferedRows(NUM_BUFFERED_ROWS)
+ .setMaxResults(MAX_RESULTS)
+ .setUseQueryCache(USE_QUERY_CACHE)
+ .setTableDefinitions(TABLE_DEFINITIONS)
+ .setAllowLargeResults(ALLOW_LARGE_RESULTS)
+ .setCreateDisposition(CREATE_DISPOSITION)
+ .setDefaultDataset(DATASET_ID)
+ .setDestinationTable(TABLE_ID)
+ .setWriteDisposition(WRITE_DISPOSITION)
+ .setPriority(PRIORITY)
+ .setFlattenResults(FLATTEN_RESULTS)
+ .setUserDefinedFunctions(USER_DEFINED_FUNCTIONS)
+ .setMaximumBillingTier(MAX_BILLING_TIER)
+ .setMaximumBytesBilled(MAX_BYTES_BILL)
+ .setSchemaUpdateOptions(SCHEMA_UPDATE_OPTIONS)
+ .setDestinationEncryptionConfiguration(JOB_ENCRYPTION_CONFIGURATION)
+ .setTimePartitioning(TIME_PARTITIONING)
+ .setClustering(CLUSTERING)
+ .setJobTimeoutMs(TIMEOUT)
+ .setRangePartitioning(RANGE_PARTITIONING)
+ .setConnectionProperties(CONNECTION_PROPERTIES)
+ .build();
+
+ @Test
+ public void testToBuilder() {
+ compareConnectionSettings(CONNECTION_SETTINGS, CONNECTION_SETTINGS.toBuilder().build());
+ }
+
+ @Test
+ public void testToBuilderIncomplete() {
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DATASET_ID).build();
+ compareConnectionSettings(connectionSettings, connectionSettings.toBuilder().build());
+ }
+
+ @Test
+ public void testBuilder() {
+ assertEquals(REQUEST_TIMEOUT, CONNECTION_SETTINGS.getRequestTimeout());
+ assertEquals(NUM_BUFFERED_ROWS, CONNECTION_SETTINGS.getNumBufferedRows());
+ assertEquals(MAX_RESULTS, CONNECTION_SETTINGS.getMaxResults());
+ }
+
+ private void compareConnectionSettings(ConnectionSettings expected, ConnectionSettings value) {
+ assertEquals(expected, value);
+ assertEquals(expected.hashCode(), value.hashCode());
+ assertEquals(expected.toString(), value.toString());
+ assertEquals(expected.getRequestTimeout(), value.getRequestTimeout());
+ assertEquals(expected.getNumBufferedRows(), value.getNumBufferedRows());
+ assertEquals(expected.getMaxResults(), value.getMaxResults());
+ assertEquals(expected.getAllowLargeResults(), value.getAllowLargeResults());
+ assertEquals(expected.getCreateDisposition(), value.getCreateDisposition());
+ assertEquals(expected.getDefaultDataset(), value.getDefaultDataset());
+ assertEquals(expected.getDestinationTable(), value.getDestinationTable());
+ assertEquals(expected.getFlattenResults(), value.getFlattenResults());
+ assertEquals(expected.getPriority(), value.getPriority());
+ assertEquals(expected.getTableDefinitions(), value.getTableDefinitions());
+ assertEquals(expected.getUseQueryCache(), value.getUseQueryCache());
+ assertEquals(expected.getUserDefinedFunctions(), value.getUserDefinedFunctions());
+ assertEquals(expected.getWriteDisposition(), value.getWriteDisposition());
+ assertEquals(expected.getMaximumBillingTier(), value.getMaximumBillingTier());
+ assertEquals(expected.getMaximumBytesBilled(), value.getMaximumBytesBilled());
+ assertEquals(expected.getSchemaUpdateOptions(), value.getSchemaUpdateOptions());
+ assertEquals(
+ expected.getDestinationEncryptionConfiguration(),
+ value.getDestinationEncryptionConfiguration());
+ assertEquals(expected.getTimePartitioning(), value.getTimePartitioning());
+ assertEquals(expected.getClustering(), value.getClustering());
+ assertEquals(expected.getJobTimeoutMs(), value.getJobTimeoutMs());
+ assertEquals(expected.getRangePartitioning(), value.getRangePartitioning());
+ assertEquals(expected.getConnectionProperties(), value.getConnectionProperties());
+ }
+}
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
index f3b60cebe..348749b46 100644
--- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -49,10 +49,15 @@
import com.google.cloud.bigquery.BigQuery.JobOption;
import com.google.cloud.bigquery.BigQuery.TableField;
import com.google.cloud.bigquery.BigQuery.TableOption;
+import com.google.cloud.bigquery.BigQueryDryRunResult;
import com.google.cloud.bigquery.BigQueryError;
import com.google.cloud.bigquery.BigQueryException;
+import com.google.cloud.bigquery.BigQueryResult;
+import com.google.cloud.bigquery.BigQuerySQLException;
import com.google.cloud.bigquery.Clustering;
+import com.google.cloud.bigquery.Connection;
import com.google.cloud.bigquery.ConnectionProperty;
+import com.google.cloud.bigquery.ConnectionSettings;
import com.google.cloud.bigquery.CopyJobConfiguration;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.DatasetId;
@@ -62,6 +67,7 @@
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.FieldList;
import com.google.cloud.bigquery.FieldValue;
+import com.google.cloud.bigquery.FieldValue.Attribute;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.FormatOptions;
import com.google.cloud.bigquery.HivePartitioningOptions;
@@ -72,6 +78,9 @@
import com.google.cloud.bigquery.JobInfo;
import com.google.cloud.bigquery.JobStatistics;
import com.google.cloud.bigquery.JobStatistics.LoadStatistics;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
+import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
+import com.google.cloud.bigquery.JobStatistics.SessionInfo;
import com.google.cloud.bigquery.JobStatistics.TransactionInfo;
import com.google.cloud.bigquery.LegacySQLTypeName;
import com.google.cloud.bigquery.LoadJobConfiguration;
@@ -79,6 +88,7 @@
import com.google.cloud.bigquery.Model;
import com.google.cloud.bigquery.ModelId;
import com.google.cloud.bigquery.ModelInfo;
+import com.google.cloud.bigquery.Parameter;
import com.google.cloud.bigquery.ParquetOptions;
import com.google.cloud.bigquery.PolicyTags;
import com.google.cloud.bigquery.QueryJobConfiguration;
@@ -129,7 +139,11 @@
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Time;
import java.time.Instant;
+import java.time.LocalTime;
import java.time.Period;
import java.util.ArrayList;
import java.util.Collection;
@@ -273,6 +287,103 @@ public class ITBigQueryTest {
BIGNUMERIC_FIELD_SCHEMA3,
BIGNUMERIC_FIELD_SCHEMA4);
+ private static final Schema BQ_RESULTSET_SCHEMA =
+ Schema.of(
+ Field.newBuilder("TimestampField", StandardSQLTypeName.TIMESTAMP)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("TimestampDescription")
+ .build(),
+ Field.newBuilder("StringField", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("StringDescription")
+ .build(),
+ Field.newBuilder("IntegerArrayField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.REPEATED)
+ .setDescription("IntegerArrayDescription")
+ .build(),
+ Field.newBuilder("BooleanField", StandardSQLTypeName.BOOL)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BooleanDescription")
+ .build(),
+ Field.newBuilder("BytesField", StandardSQLTypeName.BYTES)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BytesDescription")
+ .build(),
+ Field.newBuilder(
+ "RecordField",
+ StandardSQLTypeName.STRUCT,
+ Field.newBuilder("TimestampField", StandardSQLTypeName.TIMESTAMP)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("TimestampDescription")
+ .build(),
+ Field.newBuilder("StringField", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("StringDescription")
+ .build(),
+ Field.newBuilder("IntegerArrayField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.REPEATED)
+ .setDescription("IntegerArrayDescription")
+ .build(),
+ Field.newBuilder("BooleanField", StandardSQLTypeName.BOOL)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BooleanDescription")
+ .build(),
+ Field.newBuilder("BytesField", StandardSQLTypeName.BYTES)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BytesDescription")
+ .build())
+ .setMode(Field.Mode.REQUIRED)
+ .setDescription("RecordDescription")
+ .build(),
+ Field.newBuilder("IntegerField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("IntegerDescription")
+ .build(),
+ Field.newBuilder("FloatField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("FloatDescription")
+ .build(),
+ Field.newBuilder("GeographyField", StandardSQLTypeName.GEOGRAPHY)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("GeographyDescription")
+ .build(),
+ Field.newBuilder("NumericField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("NumericDescription")
+ .build(),
+ Field.newBuilder("BigNumericField", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BigNumericDescription")
+ .build(),
+ Field.newBuilder("BigNumericField1", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BigNumeric1Description")
+ .build(),
+ Field.newBuilder("BigNumericField2", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BigNumeric2Description")
+ .build(),
+ Field.newBuilder("BigNumericField3", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BigNumeric3Description")
+ .build(),
+ Field.newBuilder("BigNumericField4", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BigNumeric4Description")
+ .build(),
+ Field.newBuilder("TimeField", StandardSQLTypeName.TIME)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("TimeDescription")
+ .build(),
+ Field.newBuilder("DateField", StandardSQLTypeName.DATE)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("DateDescription")
+ .build(),
+ Field.newBuilder("DateTimeField", StandardSQLTypeName.DATETIME)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("DateTimeDescription")
+ .build());
+
private static final Field DDL_TIMESTAMP_FIELD_SCHEMA =
Field.newBuilder("TimestampField", LegacySQLTypeName.TIMESTAMP)
.setDescription("TimestampDescription")
@@ -316,6 +427,55 @@ public class ITBigQueryTest {
Field.newBuilder("BooleanField", LegacySQLTypeName.BOOLEAN)
.setMode(Field.Mode.NULLABLE)
.build());
+
+ private static final Schema BQ_RESULTSET_EXPECTED_SCHEMA =
+ Schema.of(
+ Field.newBuilder("StringField", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("BigNumericField", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("BooleanField", StandardSQLTypeName.BOOL)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("BytesField", StandardSQLTypeName.BYTES)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("IntegerField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("TimestampField", StandardSQLTypeName.TIMESTAMP)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("FloatField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("NumericField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("TimeField", StandardSQLTypeName.TIME)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("DateField", StandardSQLTypeName.DATE)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("DateTimeField", StandardSQLTypeName.DATETIME)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("GeographyField", StandardSQLTypeName.GEOGRAPHY)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("BytesField_1", StandardSQLTypeName.BYTES)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("BooleanField_1", StandardSQLTypeName.BOOL)
+ .setMode(Field.Mode.NULLABLE)
+ .build(),
+ Field.newBuilder("IntegerArrayField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.REPEATED)
+ .build());
+
private static final Schema QUERY_RESULT_SCHEMA_BIGNUMERIC =
Schema.of(
Field.newBuilder("TimestampField", LegacySQLTypeName.TIMESTAMP)
@@ -360,6 +520,7 @@ public class ITBigQueryTest {
private static final String LOAD_FILE = "load.csv";
private static final String LOAD_FILE_LARGE = "load_large.csv";
private static final String JSON_LOAD_FILE = "load.json";
+ private static final String JSON_LOAD_FILE_BQ_RESULTSET = "load_bq_resultset.json";
private static final String JSON_LOAD_FILE_SIMPLE = "load_simple.json";
private static final String EXTRACT_FILE = "extract.csv";
private static final String EXTRACT_MODEL_FILE = "extract_model.csv";
@@ -368,7 +529,10 @@ public class ITBigQueryTest {
private static final TableId TABLE_ID_DDL = TableId.of(DATASET, "ddl_testing_table");
private static final TableId TABLE_ID_FASTQUERY = TableId.of(DATASET, "fastquery_testing_table");
private static final TableId TABLE_ID_LARGE = TableId.of(DATASET, "large_data_testing_table");
+ private static final TableId TABLE_ID_FASTQUERY_BQ_RESULTSET =
+ TableId.of(DATASET, "fastquery_testing_bq_resultset");
private static final String CSV_CONTENT = "StringValue1\nStringValue2\n";
+
private static final String JSON_CONTENT =
"{"
+ " \"TimestampField\": \"2014-08-19 07:41:35.220 -05:00\","
@@ -424,6 +588,64 @@ public class ITBigQueryTest {
+ " \"BigNumericField3\": \"578960446186580977117854925043439539266.34992332820282019728792003956564819967\","
+ " \"BigNumericField4\": \"-578960446186580977117854925043439539266.34992332820282019728792003956564819968\""
+ "}";
+
+ private static final String JSON_CONTENT_BQ_RESULTSET =
+ "{"
+ + " \"TimestampField\": null,"
+ + " \"StringField\": null,"
+ + " \"IntegerArrayField\": null,"
+ + " \"BooleanField\": null,"
+ + " \"BytesField\": null,"
+ + " \"RecordField\": {"
+ + " \"TimestampField\": null,"
+ + " \"StringField\": null,"
+ + " \"IntegerArrayField\": null,"
+ + " \"BooleanField\": null,"
+ + " \"BytesField\": null"
+ + " },"
+ + " \"IntegerField\": null,"
+ + " \"FloatField\": null,"
+ + " \"GeographyField\": null,"
+ + " \"NumericField\": null,"
+ + " \"BigNumericField\": null,"
+ + " \"BigNumericField1\": null,"
+ + " \"BigNumericField2\": null,"
+ + " \"BigNumericField3\": null,"
+ + " \"BigNumericField4\": null,"
+ + " \"TimeField\": null,"
+ + " \"DateField\": null,"
+ + " \"DateTimeField\": null"
+ + "}\n"
+ + "{"
+ + " \"TimestampField\": \"2018-08-19 12:11:35.123456 UTC\","
+ + " \"StringField\": \"StringValue1\","
+ + " \"IntegerArrayField\": [1,2,3,4],"
+ + " \"BooleanField\": \"false\","
+ + " \"BytesField\": \""
+ + BYTES_BASE64
+ + "\","
+ + " \"RecordField\": {"
+ + " \"TimestampField\": \"1969-07-20 20:18:04 UTC\","
+ + " \"StringField\": null,"
+ + " \"IntegerArrayField\": [1,0],"
+ + " \"BooleanField\": \"true\","
+ + " \"BytesField\": \""
+ + BYTES_BASE64
+ + "\""
+ + " },"
+ + " \"IntegerField\": \"1\","
+ + " \"FloatField\": \"10.1\","
+ + " \"GeographyField\": \"POINT(-122.35022 47.649154)\","
+ + " \"NumericField\": \"100\","
+ + " \"BigNumericField\": \"0.33333333333333333333333333333333333333\","
+ + " \"BigNumericField1\": \"1e-38\","
+ + " \"BigNumericField2\": \"-1e38\","
+ + " \"BigNumericField3\": \"578960446186580977117854925043439539266.34992332820282019728792003956564819967\","
+ + " \"BigNumericField4\": \"-578960446186580977117854925043439539266.34992332820282019728792003956564819968\","
+ + " \"TimeField\": \"12:11:35.123456\","
+ + " \"DateField\": \"2018-08-19\","
+ + " \"DateTimeField\": \"2018-08-19 12:11:35.123456\""
+ + "}";
private static final String JSON_CONTENT_SIMPLE =
"{"
+ " \"TimestampField\": \"2014-08-19 07:41:35.220 -05:00\","
@@ -476,6 +698,11 @@ public static void beforeClass() throws InterruptedException, IOException {
ITBigQueryTest.class.getClassLoader().getResourceAsStream("QueryTestData.csv");
storage.createFrom(
BlobInfo.newBuilder(BUCKET, LOAD_FILE_LARGE).setContentType("text/plain").build(), stream);
+ storage.create(
+ BlobInfo.newBuilder(BUCKET, JSON_LOAD_FILE_BQ_RESULTSET)
+ .setContentType("application/json")
+ .build(),
+ JSON_CONTENT_BQ_RESULTSET.getBytes(StandardCharsets.UTF_8));
DatasetInfo info =
DatasetInfo.newBuilder(DATASET).setDescription(DESCRIPTION).setLabels(LABELS).build();
bigquery.create(info);
@@ -509,6 +736,19 @@ public static void beforeClass() throws InterruptedException, IOException {
jobFastQuery = jobFastQuery.waitFor();
assertNull(jobFastQuery.getStatus().getError());
+ LoadJobConfiguration configFastQueryBQResultset =
+ LoadJobConfiguration.newBuilder(
+ TABLE_ID_FASTQUERY_BQ_RESULTSET,
+ "gs://" + BUCKET + "/" + JSON_LOAD_FILE_BQ_RESULTSET,
+ FormatOptions.json())
+ .setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
+ .setSchema(BQ_RESULTSET_SCHEMA)
+ .setLabels(labels)
+ .build();
+ Job jobFastQueryBQResultSet = bigquery.create(JobInfo.of(configFastQueryBQResultset));
+ jobFastQueryBQResultSet = jobFastQueryBQResultSet.waitFor();
+ assertNull(jobFastQueryBQResultSet.getStatus().getError());
+
LoadJobConfiguration configurationDDL =
LoadJobConfiguration.newBuilder(
TABLE_ID_DDL, "gs://" + BUCKET + "/" + JSON_LOAD_FILE_SIMPLE, FormatOptions.json())
@@ -712,6 +952,7 @@ public void testCreateTableWithRangePartitioning() {
}
}
+ /* TODO(prasmish): replicate this test case for executeSelect on the relevant part */
@Test
public void testJsonType() throws InterruptedException {
String tableName = "test_create_table_jsontype";
@@ -819,6 +1060,7 @@ public void testJsonType() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate this test case for executeSelect on the relevant part */
@Test
public void testIntervalType() throws InterruptedException {
String tableName = "test_create_table_intervaltype";
@@ -1701,6 +1943,7 @@ public void testInsertAllWithErrors() {
assertTrue(bigquery.delete(TableId.of(DATASET, tableName)));
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testListAllTableData() {
Page rows = bigquery.listTableData(TABLE_ID);
@@ -2036,6 +2279,7 @@ public void testAuthorizeDataset() {
assertEquals(sharedDatasetAcl, updatedDataset.getAcl());
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testSingleStatementsQueryException() throws InterruptedException {
String invalidQuery =
@@ -2052,6 +2296,7 @@ public void testSingleStatementsQueryException() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testMultipleStatementsQueryException() throws InterruptedException {
String invalidQuery =
@@ -2070,6 +2315,7 @@ public void testMultipleStatementsQueryException() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQuery() throws InterruptedException {
String query = "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID.getTable();
@@ -2102,6 +2348,16 @@ public void testQuery() throws InterruptedException {
assertNotNull(statistics.getQueryPlan());
}
+ @Test
+ public void testExecuteSelectDefaultConnectionSettings() throws SQLException {
+ // Use the default connection settings
+ Connection connection = bigquery.createConnection();
+ String query = "SELECT corpus FROM `bigquery-public-data.samples.shakespeare` GROUP BY corpus;";
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ assertEquals(42, bigQueryResult.getTotalRows());
+ }
+
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryTimeStamp() throws InterruptedException {
String query = "SELECT TIMESTAMP '2022-01-24T23:54:25.095574Z'";
@@ -2136,6 +2392,7 @@ public void testQueryTimeStamp() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryCaseInsensitiveSchemaFieldByGetName() throws InterruptedException {
String query = "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID.getTable();
@@ -2164,6 +2421,7 @@ public void testQueryCaseInsensitiveSchemaFieldByGetName() throws InterruptedExc
assertEquals(2, rowCount);
}
+ /* TODO(prasmish): replicate bigquery.query part of the test case for executeSelect - modify this test case */
@Test
public void testQueryExternalHivePartitioningOptionAutoLayout() throws InterruptedException {
String tableName = "test_queryexternalhivepartition_autolayout_table";
@@ -2198,6 +2456,7 @@ public void testQueryExternalHivePartitioningOptionAutoLayout() throws Interrupt
assertTrue(bigquery.delete(tableId));
}
+ /* TODO(prasmish): replicate bigquery.query part of the test case for executeSelect - modify this test case */
@Test
public void testQueryExternalHivePartitioningOptionCustomLayout() throws InterruptedException {
String tableName = "test_queryexternalhivepartition_customlayout_table";
@@ -2233,6 +2492,463 @@ public void testQueryExternalHivePartitioningOptionCustomLayout() throws Interru
assertTrue(bigquery.delete(tableId));
}
+ @Test
+ public void testConnectionImplDryRun() throws SQLException {
+ String query =
+ String.format(
+ "select StringField, BigNumericField, BooleanField, BytesField, IntegerField, TimestampField, FloatField, NumericField, TimeField, DateField, DateTimeField , GeographyField, RecordField.BytesField, RecordField.BooleanField, IntegerArrayField from %s where StringField = ? order by TimestampField",
+ TABLE_ID_FASTQUERY_BQ_RESULTSET.getTable());
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setCreateSession(true)
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryDryRunResult bigQueryDryRunResultSet = connection.dryRun(query);
+ assertNotNull(bigQueryDryRunResultSet.getSchema());
+ assertEquals(
+ BQ_RESULTSET_EXPECTED_SCHEMA, bigQueryDryRunResultSet.getSchema()); // match the schema
+ List queryParameters = bigQueryDryRunResultSet.getQueryParameters();
+ assertEquals(StandardSQLTypeName.STRING, queryParameters.get(0).getValue().getType());
+ QueryStatistics queryStatistics = bigQueryDryRunResultSet.getStatistics().getQueryStatistics();
+ assertNotNull(queryStatistics);
+ SessionInfo sessionInfo = bigQueryDryRunResultSet.getStatistics().getSessionInfo();
+ assertNotNull(sessionInfo.getSessionId());
+ assertEquals(StatementType.SELECT, queryStatistics.getStatementType());
+ }
+
+ @Test
+ // This test case test the order of the records, making sure that the result is not jumbled up due
+ // to the multithreaded BigQueryResult implementation
+ public void testBQResultSetMultiThreadedOrder() throws SQLException {
+ String query =
+ "SELECT date FROM "
+ + TABLE_ID_LARGE.getTable()
+ + " where date is not null order by date asc limit 300000";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setNumBufferedRows(10000) // page size
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ assertTrue(rs.next());
+ ++cnt;
+ java.sql.Date lastDate = rs.getDate(0);
+ while (rs.next()) {
+ assertNotNull(rs.getDate(0));
+ assertTrue(rs.getDate(0).getTime() >= lastDate.getTime()); // sorted order is maintained
+ lastDate = rs.getDate(0);
+ ++cnt;
+ }
+ assertEquals(300000, cnt); // total 300000 rows should be read
+ }
+
+ @Test
+ public void testBQResultSetPaginationSlowQuery() throws SQLException {
+ String query =
+ "SELECT date, county, state_name, confirmed_cases, deaths FROM "
+ + TABLE_ID_LARGE.getTable()
+ + " where date is not null and county is not null and state_name is not null order by date limit 300000";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setNumBufferedRows(10000) // page size
+ .setJobTimeoutMs(
+ 15000L) // So that ConnectionImpl.isFastQuerySupported returns false, and the slow
+ // query route gets executed
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ while (rs.next()) { // pagination starts after approx 120,000 records
+ assertNotNull(rs.getDate(0));
+ assertNotNull(rs.getString(1));
+ assertNotNull(rs.getString(2));
+ assertTrue(rs.getInt(3) >= 0);
+ assertTrue(rs.getInt(4) >= 0);
+ ++cnt;
+ }
+ assertEquals(300000, cnt); // total 300000 rows should be read
+ }
+
+ @Test
+ public void testExecuteSelectSinglePageTableRow() throws SQLException {
+ String query =
+ "select StringField, BigNumericField, BooleanField, BytesField, IntegerField, TimestampField, FloatField, "
+ + "NumericField, TimeField, DateField, DateTimeField , GeographyField, RecordField.BytesField, RecordField.BooleanField, IntegerArrayField from "
+ + TABLE_ID_FASTQUERY_BQ_RESULTSET.getTable()
+ + " order by TimestampField";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ Schema sc = bigQueryResult.getSchema();
+
+ assertEquals(BQ_RESULTSET_EXPECTED_SCHEMA, sc); // match the schema
+ assertEquals(2, bigQueryResult.getTotalRows()); // Expecting 2 rows
+
+ assertTrue(rs.next()); // first row
+ // checking for the null or 0 column values
+ assertNull(rs.getString("StringField"));
+ assertTrue(rs.getDouble("BigNumericField") == 0.0d);
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertNull(rs.getBytes("BytesField"));
+ assertEquals(rs.getInt("IntegerField"), 0);
+ assertNull(rs.getTimestamp("TimestampField"));
+ assertNull(rs.getDate("DateField"));
+ assertTrue(rs.getDouble("FloatField") == 0.0d);
+ assertTrue(rs.getDouble("NumericField") == 0.0d);
+ assertNull(rs.getTime("TimeField"));
+ assertNull(rs.getString("DateTimeField"));
+ assertNull(rs.getString("GeographyField"));
+ assertNull(rs.getBytes("BytesField_1"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+
+ assertTrue(rs.next()); // second row
+ // second row is non null, comparing the values
+ assertEquals("StringValue1", rs.getString("StringField"));
+ assertTrue(rs.getDouble("BigNumericField") == 0.3333333333333333d);
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertNotNull(rs.getBytes("BytesField"));
+ assertEquals(1, rs.getInt("IntegerField"));
+ assertEquals(1534680695123L, rs.getTimestamp("TimestampField").getTime());
+ assertEquals(java.sql.Date.valueOf("2018-08-19"), rs.getDate("DateField"));
+ assertTrue(rs.getDouble("FloatField") == 10.1d);
+ assertTrue(rs.getDouble("NumericField") == 100.0d);
+ assertEquals(Time.valueOf(LocalTime.of(12, 11, 35, 123456)), rs.getTime("TimeField"));
+ assertEquals("2018-08-19T12:11:35.123456", rs.getString("DateTimeField"));
+ assertEquals("POINT(-122.35022 47.649154)", rs.getString("GeographyField"));
+ assertNotNull(rs.getBytes("BytesField_1"));
+ assertTrue(rs.getBoolean("BooleanField_1"));
+ assertTrue(
+ rs.getObject("IntegerArrayField") instanceof com.google.cloud.bigquery.FieldValueList);
+ FieldValueList integerArrayFieldValue =
+ (com.google.cloud.bigquery.FieldValueList) rs.getObject("IntegerArrayField");
+ assertEquals(4, integerArrayFieldValue.size()); // Array has 4 elements
+ assertEquals(3, (integerArrayFieldValue.get(2).getNumericValue()).intValue());
+
+ assertFalse(rs.next()); // no 3rd row in the table
+ }
+
+ @Test
+ public void testConnectionClose() throws SQLException {
+ String query =
+ "SELECT date, county, state_name, confirmed_cases, deaths FROM "
+ + TABLE_ID_LARGE.getTable()
+ + " where date is not null and county is not null and state_name is not null order by date limit 300000";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setNumBufferedRows(10000) // page size
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ while (rs.next()) {
+ ++cnt;
+ if (cnt > 57000) { // breaking at 57K, query reads 300K
+ assertTrue(connection.close()); // we should be able to cancel the connection
+ }
+ }
+ assertTrue(
+ cnt < 60000); // Few extra records are still read (generally ~10) even after canceling, as
+ // the backgrounds threads are still active while the interrupt occurs and the
+ // buffer and pageCache are cleared
+ }
+
+ @Test
+ public void testBQResultSetPagination() throws SQLException {
+ String query =
+ "SELECT date, county, state_name, confirmed_cases, deaths FROM "
+ + TABLE_ID_LARGE.getTable()
+ + " where date is not null and county is not null and state_name is not null order by date limit 300000";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setNumBufferedRows(10000) // page size
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ while (rs.next()) { // pagination starts after approx 120,000 records
+ assertNotNull(rs.getDate(0));
+ assertNotNull(rs.getString(1));
+ assertNotNull(rs.getString(2));
+ assertTrue(rs.getInt(3) >= 0);
+ assertTrue(rs.getInt(4) >= 0);
+ ++cnt;
+ }
+ assertEquals(300000, cnt); // total 300000 rows should be read
+ }
+
+ @Test
+ public void testReadAPIIterationAndOrder()
+ throws SQLException { // use read API to read 300K records and check the order
+ String query =
+ "SELECT date, county, state_name, confirmed_cases, deaths FROM "
+ + TABLE_ID_LARGE.getTable()
+ + " where date is not null and county is not null and state_name is not null order by confirmed_cases asc limit 300000";
+
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setPriority(
+ QueryJobConfiguration.Priority
+ .INTERACTIVE) // required for this integration test so that isFastQuerySupported
+ // returns false
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ int lasConfirmedCases = Integer.MIN_VALUE;
+ while (rs.next()) { // pagination starts after approx 120,000 records
+ assertNotNull(rs.getDate(0));
+ assertNotNull(rs.getString(1));
+ assertNotNull(rs.getString(2));
+ assertTrue(rs.getInt(3) >= 0);
+ assertTrue(rs.getInt(4) >= 0);
+
+ // check if the records are sorted
+ assertTrue(rs.getInt(3) >= lasConfirmedCases);
+ lasConfirmedCases = rs.getInt(3);
+ ++cnt;
+ }
+ assertEquals(300000, cnt); // total 300000 rows should be read
+ connection.close();
+ }
+
+ @Test
+ public void testReadAPIConnectionMultiClose()
+ throws
+ SQLException { // use read API to read 300K records, then closes the connection. This test
+ // repeats it multiple times and assets if the connection was closed
+ String query =
+ "SELECT date, county, state_name, confirmed_cases, deaths FROM "
+ + TABLE_ID_LARGE.getTable()
+ + " where date is not null and county is not null and state_name is not null order by confirmed_cases asc limit 300000";
+
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setPriority(
+ QueryJobConfiguration.Priority
+ .INTERACTIVE) // required for this integration test so that isFastQuerySupported
+ // returns false
+ .build();
+ int closeCnt = 0, runCnt = 3;
+ for (int run = 0; run < runCnt; run++) {
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ while (rs.next()) { // pagination starts after approx 120,000 records
+ assertNotNull(rs.getDate(0));
+ ++cnt;
+ }
+ assertEquals(300000, cnt); // total 300000 rows should be read
+ assertTrue(connection.close()); // check if connection closed
+ closeCnt++;
+ }
+ assertEquals(
+ closeCnt, runCnt); // check if the connection closed for the required number of times
+ }
+
+ @Test
+ public void testExecuteSelectSinglePageTableRowColInd() throws SQLException {
+ String query =
+ "select StringField, BigNumericField, BooleanField, BytesField, IntegerField, TimestampField, FloatField, "
+ + "NumericField, TimeField, DateField, DateTimeField , GeographyField, RecordField.BytesField, RecordField.BooleanField, IntegerArrayField from "
+ + TABLE_ID_FASTQUERY_BQ_RESULTSET.getTable()
+ + " order by TimestampField";
+ /*
+ Column Index mapping for ref:
+ StringField, 0 BigNumericField, 1 BooleanField, 2 BytesField, 3 IntegerField, 4 TimestampField, 5 FloatField, " 6
+ NumericField, 7 TimeField, 8 DateField, 9 DateTimeField , 10 GeographyField, 11 RecordField.BytesField, 12 RecordField.BooleanField, 13 IntegerArrayField 14
+ */
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ ResultSet rs = bigQueryResult.getResultSet();
+ Schema sc = bigQueryResult.getSchema();
+
+ assertEquals(BQ_RESULTSET_EXPECTED_SCHEMA, sc); // match the schema
+ assertEquals(2, bigQueryResult.getTotalRows()); // Expecting 2 rows
+ while (rs.next()) {
+ assertEquals(rs.getString(0), rs.getString("StringField"));
+ assertTrue(rs.getDouble(1) == rs.getDouble("BigNumericField"));
+ assertEquals(rs.getBoolean(2), rs.getBoolean("BooleanField"));
+ if (rs.getBytes(3) == null) { // both overloads should be null
+ assertEquals(rs.getBytes(3), rs.getBytes("BytesField"));
+ } else { // value in String representation should be the same
+ assertEquals(
+ new String(rs.getBytes(3), StandardCharsets.UTF_8),
+ new String(rs.getBytes("BytesField"), StandardCharsets.UTF_8));
+ }
+ assertEquals(rs.getInt(4), rs.getInt("IntegerField"));
+ assertEquals(rs.getTimestamp(5), rs.getTimestamp("TimestampField"));
+ assertEquals(rs.getDate(9), rs.getDate("DateField"));
+ assertTrue(rs.getDouble("FloatField") == rs.getDouble(6));
+ assertTrue(rs.getDouble("NumericField") == rs.getDouble(7));
+ assertEquals(rs.getTime(8), rs.getTime("TimeField"));
+ assertEquals(rs.getString(10), rs.getString("DateTimeField"));
+ assertEquals(rs.getString(11), rs.getString("GeographyField"));
+ if (rs.getBytes(12) == null) { // both overloads should be null
+ assertEquals(rs.getBytes(12), rs.getBytes("BytesField_1"));
+ } else { // value in String representation should be the same
+ assertEquals(
+ new String(rs.getBytes(12), StandardCharsets.UTF_8),
+ new String(rs.getBytes("BytesField_1"), StandardCharsets.UTF_8));
+ }
+ assertEquals(rs.getBoolean(13), rs.getBoolean("BooleanField_1"));
+ assertTrue(
+ rs.getObject("IntegerArrayField") instanceof com.google.cloud.bigquery.FieldValueList);
+ FieldValueList integerArrayFieldValue =
+ (com.google.cloud.bigquery.FieldValueList) rs.getObject("IntegerArrayField");
+ assertTrue(rs.getObject(14) instanceof com.google.cloud.bigquery.FieldValueList);
+ FieldValueList integerArrayFieldValueColInd =
+ (com.google.cloud.bigquery.FieldValueList) rs.getObject(14);
+ assertEquals(
+ integerArrayFieldValue.size(),
+ integerArrayFieldValueColInd.size()); // Array has 4 elements
+ if (integerArrayFieldValue.size() == 4) { // as we are picking the third index
+ assertEquals(
+ (integerArrayFieldValue.get(2).getNumericValue()).intValue(),
+ (integerArrayFieldValueColInd.get(2).getNumericValue()).intValue());
+ }
+ }
+ }
+
+ @Test
+ public void testExecuteSelectStruct() throws SQLException {
+ String query = "select (STRUCT(\"Vancouver\" as city, 5 as years)) as address";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ assertEquals(1, bigQueryResult.getTotalRows());
+
+ Schema schema = bigQueryResult.getSchema();
+ assertEquals("address", schema.getFields().get(0).getName());
+ assertEquals(Field.Mode.NULLABLE, schema.getFields().get(0).getMode());
+ // Backend is currently returning LegacySQLTypeName. Tracking bug: b/202977620
+ assertEquals(LegacySQLTypeName.RECORD, schema.getFields().get(0).getType());
+ assertEquals("city", schema.getFields().get(0).getSubFields().get(0).getName());
+ assertEquals(
+ LegacySQLTypeName.STRING, schema.getFields().get(0).getSubFields().get(0).getType());
+ assertEquals(Field.Mode.NULLABLE, schema.getFields().get(0).getSubFields().get(0).getMode());
+ assertEquals("years", schema.getFields().get(0).getSubFields().get(1).getName());
+ assertEquals(
+ LegacySQLTypeName.INTEGER, schema.getFields().get(0).getSubFields().get(1).getType());
+ assertEquals(Field.Mode.NULLABLE, schema.getFields().get(0).getSubFields().get(1).getMode());
+
+ ResultSet rs = bigQueryResult.getResultSet();
+ assertTrue(rs.next());
+ FieldValueList addressFieldValue =
+ (com.google.cloud.bigquery.FieldValueList) rs.getObject("address");
+ assertEquals(rs.getObject("address"), rs.getObject(0));
+ assertEquals("Vancouver", addressFieldValue.get(0).getStringValue());
+ assertEquals(5, addressFieldValue.get(1).getLongValue());
+ assertFalse(rs.next()); // only 1 row of data
+ }
+
+ @Test
+ public void testExecuteSelectStructSubField() throws SQLException {
+ String query =
+ "select address.city from (select (STRUCT(\"Vancouver\" as city, 5 as years)) as address)";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ assertEquals(1, bigQueryResult.getTotalRows());
+
+ Schema schema = bigQueryResult.getSchema();
+ assertEquals("city", schema.getFields().get(0).getName());
+ assertEquals(Field.Mode.NULLABLE, schema.getFields().get(0).getMode());
+ // Backend is currently returning LegacySQLTypeName. Tracking bug: b/202977620
+ assertEquals(LegacySQLTypeName.STRING, schema.getFields().get(0).getType());
+ assertNull(
+ schema.getFields().get(0).getSubFields()); // this is a String field without any subfields
+
+ ResultSet rs = bigQueryResult.getResultSet();
+ assertTrue(rs.next());
+ String cityFieldValue = rs.getString("city");
+ assertEquals(rs.getString("city"), rs.getObject(0));
+ assertEquals("Vancouver", cityFieldValue);
+ assertFalse(rs.next()); // only 1 row of data
+ }
+
+ @Test
+ public void testExecuteSelectArray() throws SQLException {
+ String query = "SELECT [1,2,3]";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ assertEquals(1, bigQueryResult.getTotalRows());
+
+ Schema schema = bigQueryResult.getSchema();
+ assertEquals("f0_", schema.getFields().get(0).getName());
+ assertEquals(Field.Mode.REPEATED, schema.getFields().get(0).getMode());
+ assertEquals(LegacySQLTypeName.INTEGER, schema.getFields().get(0).getType());
+ assertNull(schema.getFields().get(0).getSubFields()); // no subfields for Integers
+
+ ResultSet rs = bigQueryResult.getResultSet();
+ assertTrue(rs.next());
+ FieldValueList arrayFieldValue = (com.google.cloud.bigquery.FieldValueList) rs.getObject(0);
+ assertEquals(1, arrayFieldValue.get(0).getLongValue());
+ assertEquals(2, arrayFieldValue.get(1).getLongValue());
+ assertEquals(3, arrayFieldValue.get(2).getLongValue());
+ }
+
+ @Test
+ public void testExecuteSelectArrayOfStruct() throws SQLException {
+ String query =
+ "SELECT [STRUCT(\"Vancouver\" as city, 5 as years), STRUCT(\"Boston\" as city, 10 as years)]";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ assertEquals(1, bigQueryResult.getTotalRows());
+
+ Schema schema = bigQueryResult.getSchema();
+ assertEquals("f0_", schema.getFields().get(0).getName());
+ assertEquals(Field.Mode.REPEATED, schema.getFields().get(0).getMode());
+ // Backend is currently returning LegacySQLTypeName. Tracking bug: b/202977620
+ // Verify the field metadata of the two subfields of the struct
+ assertEquals(LegacySQLTypeName.RECORD, schema.getFields().get(0).getType());
+ assertEquals("city", schema.getFields().get(0).getSubFields().get(0).getName());
+ assertEquals(
+ LegacySQLTypeName.STRING, schema.getFields().get(0).getSubFields().get(0).getType());
+ assertEquals(Field.Mode.NULLABLE, schema.getFields().get(0).getSubFields().get(0).getMode());
+ assertEquals("years", schema.getFields().get(0).getSubFields().get(1).getName());
+ assertEquals(
+ LegacySQLTypeName.INTEGER, schema.getFields().get(0).getSubFields().get(1).getType());
+ assertEquals(Field.Mode.NULLABLE, schema.getFields().get(0).getSubFields().get(1).getMode());
+
+ ResultSet rs = bigQueryResult.getResultSet();
+ assertTrue(rs.next());
+ FieldValueList arrayOfStructFieldValue =
+ (com.google.cloud.bigquery.FieldValueList) rs.getObject(0);
+ // Verify the values of the two structs in the array
+ assertEquals(Attribute.RECORD, arrayOfStructFieldValue.get(0).getAttribute());
+ assertEquals(
+ "Vancouver", arrayOfStructFieldValue.get(0).getRecordValue().get(0).getStringValue());
+ assertEquals(5, arrayOfStructFieldValue.get(0).getRecordValue().get(1).getLongValue());
+ assertEquals(Attribute.RECORD, arrayOfStructFieldValue.get(1).getAttribute());
+ assertEquals("Boston", arrayOfStructFieldValue.get(1).getRecordValue().get(0).getStringValue());
+ assertEquals(10, arrayOfStructFieldValue.get(1).getRecordValue().get(1).getLongValue());
+ }
+
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testFastQueryMultipleRuns() throws InterruptedException {
String query =
@@ -2265,6 +2981,7 @@ public void testFastQueryMultipleRuns() throws InterruptedException {
assertFalse(result2.hasNextPage());
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testFastQuerySinglePageDuplicateRequestIds() throws InterruptedException {
String query =
@@ -2294,6 +3011,7 @@ public void testFastQuerySinglePageDuplicateRequestIds() throws InterruptedExcep
assertFalse(result2.hasNextPage());
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testFastSQLQuery() throws InterruptedException {
String query =
@@ -2323,6 +3041,7 @@ public void testFastSQLQuery() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testFastSQLQueryMultiPage() throws InterruptedException {
String query =
@@ -2449,6 +3168,7 @@ public void testFastQuerySlowDDL() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testFastQueryHTTPException() throws InterruptedException {
String queryInvalid =
@@ -2511,10 +3231,36 @@ public void testQuerySessionSupport() throws InterruptedException {
remoteJobWithSession = remoteJobWithSession.waitFor();
assertNull(remoteJobWithSession.getStatus().getError());
Job queryJobWithSession = bigquery.getJob(remoteJobWithSession.getJobId());
- JobStatistics.QueryStatistics statisticsWithSession = queryJobWithSession.getStatistics();
+ QueryStatistics statisticsWithSession = queryJobWithSession.getStatistics();
assertEquals(sessionId, statisticsWithSession.getSessionInfo().getSessionId());
}
+ // TODO: uncomment this testcase when executeUpdate is implemented
+ // @Test
+ // public void testExecuteSelectWithSession() throws BigQuerySQLException {
+ // String query = "CREATE TEMPORARY TABLE temptable AS SELECT 17 as foo";
+ // ConnectionSettings connectionSettings =
+ // ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).setCreateSession(true).build();
+ // Connection connection = bigquery.createConnection(connectionSettings);
+ // BigQueryResult bigQueryResult = connection.execute(query);
+ // BigQueryResultStats stats = bigQueryResult.getBigQueryResultStats();
+ // assertNotNull(stats.getSessionInfo().getSessionId());
+ // }
+
+ @Test
+ public void testExecuteSelectSessionSupport() throws BigQuerySQLException {
+ String query = "SELECT 17 as foo";
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .setCreateSession(true)
+ .build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ BigQueryResult bigQueryResult = connection.executeSelect(query);
+ String sessionId = bigQueryResult.getBigQueryResultStats().getSessionInfo().getSessionId();
+ assertNotNull(sessionId);
+ }
+
@Test
public void testDmlStatistics() throws InterruptedException {
String tableName = TABLE_ID_FASTQUERY.getTable();
@@ -2535,6 +3281,7 @@ public void testDmlStatistics() throws InterruptedException {
assertEquals(2L, statistics.getDmlStats().getUpdatedRowCount().longValue());
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testTransactionInfo() throws InterruptedException {
String tableName = TABLE_ID_FASTQUERY.getTable();
@@ -2556,6 +3303,7 @@ public void testTransactionInfo() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testScriptStatistics() throws InterruptedException {
String script =
@@ -2706,6 +3454,27 @@ public void testPositionalQueryParameters() throws InterruptedException {
}
}
+ /* TODO(prasmish): expand below test case with all the fields shown in the above test case */
+ @Test
+ public void testExecuteSelectWithPositionalQueryParameters() throws BigQuerySQLException {
+ String query =
+ "SELECT TimestampField, StringField FROM "
+ + TABLE_ID.getTable()
+ + " WHERE StringField = ?"
+ + " AND TimestampField > ?";
+ QueryParameterValue stringParameter = QueryParameterValue.string("stringValue");
+ QueryParameterValue timestampParameter =
+ QueryParameterValue.timestamp("2014-01-01 07:00:00.000000+00:00");
+ Parameter stringParam = Parameter.newBuilder().setValue(stringParameter).build();
+ Parameter timeStampParam = Parameter.newBuilder().setValue(timestampParameter).build();
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ List parameters = ImmutableList.of(stringParam, timeStampParam);
+ BigQueryResult rs = connection.executeSelect(query, parameters);
+ assertEquals(2, rs.getTotalRows());
+ }
+
@Test
public void testNamedQueryParameters() throws InterruptedException {
String query =
@@ -2728,6 +3497,30 @@ public void testNamedQueryParameters() throws InterruptedException {
assertEquals(2, Iterables.size(result.getValues()));
}
+ @Test
+ public void testExecuteSelectWithNamedQueryParameters() throws BigQuerySQLException {
+ String query =
+ "SELECT TimestampField, StringField, BooleanField FROM "
+ + TABLE_ID.getTable()
+ + " WHERE StringField = @stringParam"
+ + " AND IntegerField IN UNNEST(@integerList)";
+ QueryParameterValue stringParameter = QueryParameterValue.string("stringValue");
+ QueryParameterValue intArrayParameter =
+ QueryParameterValue.array(new Integer[] {3, 4}, Integer.class);
+ Parameter stringParam =
+ Parameter.newBuilder().setName("stringParam").setValue(stringParameter).build();
+ Parameter intArrayParam =
+ Parameter.newBuilder().setName("integerList").setValue(intArrayParameter).build();
+
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder().setDefaultDataset(DatasetId.of(DATASET)).build();
+ Connection connection = bigquery.createConnection(connectionSettings);
+ List parameters = ImmutableList.of(stringParam, intArrayParam);
+ BigQueryResult rs = connection.executeSelect(query, parameters);
+ assertEquals(2, rs.getTotalRows());
+ }
+
+ /* TODO(prasmish): replicate relevant parts of the test case for executeSelect */
@Test
public void testStructNamedQueryParameters() throws InterruptedException {
QueryParameterValue booleanValue = QueryParameterValue.bool(true);
@@ -2782,6 +3575,7 @@ private static void assertsFieldValue(FieldValue record) {
assertEquals("test-stringField", record.getRecordValue().get("stringField").getStringValue());
}
+ /* TODO(prasmish): replicate relevant parts of the test case for executeSelect */
@Test
public void testNestedStructNamedQueryParameters() throws InterruptedException {
QueryParameterValue booleanValue = QueryParameterValue.bool(true);
@@ -2825,6 +3619,7 @@ public void testNestedStructNamedQueryParameters() throws InterruptedException {
}
}
+ /* TODO(prasmish): replicate relevant parts of the test case for executeSelect */
@Test
public void testBytesParameter() throws Exception {
String query = "SELECT BYTE_LENGTH(@p) AS length";
@@ -3108,6 +3903,7 @@ public void testCopyJobWithLabels() throws InterruptedException {
assertTrue(remoteTable.delete());
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryJob() throws InterruptedException, TimeoutException {
String tableName = "test_query_job_table";
@@ -3152,6 +3948,7 @@ public void testQueryJob() throws InterruptedException, TimeoutException {
assertNotNull(statistics.getQueryPlan());
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryJobWithConnectionProperties() throws InterruptedException {
String tableName = "test_query_job_table_connection_properties";
@@ -3171,6 +3968,7 @@ public void testQueryJobWithConnectionProperties() throws InterruptedException {
assertTrue(bigquery.delete(destinationTable));
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryJobWithLabels() throws InterruptedException, TimeoutException {
String tableName = "test_query_job_table";
@@ -3194,6 +3992,7 @@ public void testQueryJobWithLabels() throws InterruptedException, TimeoutExcepti
}
}
+ /* TODO(prasmish): replicate the entire test case for executeSelect */
@Test
public void testQueryJobWithRangePartitioning() throws InterruptedException {
String tableName = "test_query_job_table_rangepartitioning";
@@ -3304,6 +4103,7 @@ public void testQueryJobWithDryRun() throws InterruptedException, TimeoutExcepti
.build();
Job remoteJob = bigquery.create(JobInfo.of(configuration));
assertNull(remoteJob.getJobId().getJob());
+ remoteJob.getStatistics();
assertEquals(DONE, remoteJob.getStatus().getState());
assertNotNull(remoteJob.getConfiguration());
}
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java
new file mode 100644
index 000000000..d672967b1
--- /dev/null
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITNightlyBigQueryTest.java
@@ -0,0 +1,610 @@
+/*
+ * Copyright 2022 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.it;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQueryError;
+import com.google.cloud.bigquery.BigQueryException;
+import com.google.cloud.bigquery.BigQueryResult;
+import com.google.cloud.bigquery.BigQuerySQLException;
+import com.google.cloud.bigquery.Connection;
+import com.google.cloud.bigquery.ConnectionSettings;
+import com.google.cloud.bigquery.Dataset;
+import com.google.cloud.bigquery.DatasetId;
+import com.google.cloud.bigquery.DatasetInfo;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.InsertAllRequest;
+import com.google.cloud.bigquery.InsertAllResponse;
+import com.google.cloud.bigquery.Parameter;
+import com.google.cloud.bigquery.QueryParameterValue;
+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.TableDefinition;
+import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.TableInfo;
+import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.BaseEncoding;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.nio.charset.StandardCharsets;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Time;
+import java.time.LocalTime;
+import java.time.ZoneId;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.TimeZone;
+import java.util.concurrent.ExecutionException;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import org.apache.arrow.vector.util.JsonStringArrayList;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+public class ITNightlyBigQueryTest {
+ private static final Logger logger = Logger.getLogger(ITNightlyBigQueryTest.class.getName());
+ private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
+ private static final String TABLE = "TEMP_RS_TEST_TABLE";
+ private static final byte[] BYTES = "TestByteValue".getBytes(StandardCharsets.UTF_8);
+ private static final String BYTES_BASE64 = BaseEncoding.base64().encode(BYTES);
+ // Script will populate NUM_BATCHES*REC_PER_BATCHES number of records (eg: 100*10000 = 1M)
+ private static final int NUM_BATCHES = 55;
+ private static final int REC_PER_BATCHES = 10000;
+ private static final int LIMIT_RECS = 500000; // We can plan to read ~ 500K / 1M records
+ private static final int MULTI_LIMIT_RECS =
+ 300000; // Used for multiquery testcase, a lower limit like 300K should be fine
+ private static int rowCnt = 0;
+ private static BigQuery bigquery;
+ private static final String BASE_QUERY =
+ "select StringField, GeographyField, BooleanField, BigNumericField, IntegerField, NumericField, BytesField, "
+ + "TimestampField, TimeField, DateField, IntegerArrayField, RecordField.BooleanField, RecordField.StringField ,"
+ + " JSONField, JSONField.hello, JSONField.id from %s.%s order by IntegerField asc LIMIT %s";
+ private static final String POSITIONAL_QUERY =
+ String.format(
+ "select RecordField.BooleanField, RecordField.StringField, StringField, BooleanField, BytesField, IntegerField, GeographyField, NumericField, BigNumericField, TimeField, DateField, TimestampField, JSONField from %s.%s where DateField = ? and BooleanField = ? and IntegerField > ? and NumericField > ? LIMIT %s",
+ DATASET, TABLE, MULTI_LIMIT_RECS);
+ private static final String QUERY = String.format(BASE_QUERY, DATASET, TABLE, LIMIT_RECS);
+ private static final String MULTI_QUERY =
+ String.format(BASE_QUERY, DATASET, TABLE, MULTI_LIMIT_RECS);
+ private static final String INVALID_QUERY =
+ String.format(
+ "select into %s.%s order by IntegerField asc LIMIT %s", DATASET, TABLE, LIMIT_RECS);
+
+ private static final Schema BQ_SCHEMA =
+ Schema.of(
+ Field.newBuilder("TimestampField", StandardSQLTypeName.TIMESTAMP)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("TimestampDescription")
+ .build(),
+ Field.newBuilder("StringField", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("StringDescription")
+ .build(),
+ Field.newBuilder("IntegerArrayField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.REPEATED)
+ .setDescription("IntegerArrayDescription")
+ .build(),
+ Field.newBuilder("BooleanField", StandardSQLTypeName.BOOL)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BooleanDescription")
+ .build(),
+ Field.newBuilder("BytesField", StandardSQLTypeName.BYTES)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BytesDescription")
+ .build(),
+ Field.newBuilder(
+ "RecordField",
+ StandardSQLTypeName.STRUCT,
+ Field.newBuilder("StringField", StandardSQLTypeName.STRING)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("StringDescription")
+ .build(),
+ Field.newBuilder("BooleanField", StandardSQLTypeName.BOOL)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BooleanDescription")
+ .build())
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("RecordDescription")
+ .build(),
+ Field.newBuilder("IntegerField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("IntegerDescription")
+ .build(),
+ Field.newBuilder("GeographyField", StandardSQLTypeName.GEOGRAPHY)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("GeographyDescription")
+ .build(),
+ Field.newBuilder("NumericField", StandardSQLTypeName.NUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("NumericDescription")
+ .build(),
+ Field.newBuilder("BigNumericField", StandardSQLTypeName.BIGNUMERIC)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("BigNumericDescription")
+ .build(),
+ Field.newBuilder("TimeField", StandardSQLTypeName.TIME)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("TimeDescription")
+ .build(),
+ Field.newBuilder("DateField", StandardSQLTypeName.DATE)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("DateDescription")
+ .build(),
+ Field.newBuilder("JSONField", StandardSQLTypeName.JSON)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("JSONFieldDescription")
+ .build(),
+ Field.newBuilder("IntervalField", StandardSQLTypeName.INTERVAL)
+ .setMode(Field.Mode.NULLABLE)
+ .setDescription("IntervalFieldDescription")
+ .build());
+
+ @Rule public Timeout globalTimeout = Timeout.seconds(1800); // setting 30 mins as the timeout
+
+ @BeforeClass
+ public static void beforeClass() throws InterruptedException, IOException {
+ RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
+ bigquery = bigqueryHelper.getOptions().getService();
+ createDataset(DATASET);
+ createTable(DATASET, TABLE, BQ_SCHEMA);
+ populateTestRecords(DATASET, TABLE);
+ }
+
+ @AfterClass
+ public static void afterClass() throws ExecutionException, InterruptedException {
+ try {
+ if (bigquery != null) {
+ deleteTable(DATASET, TABLE);
+ RemoteBigQueryHelper.forceDelete(bigquery, DATASET);
+ } else {
+ fail("Error clearing the test dataset");
+ }
+ } catch (BigQueryException e) {
+ fail("Error clearing the test dataset " + e);
+ }
+ }
+
+ @Test
+ public void testInvalidQuery() throws BigQuerySQLException {
+ Connection connection = getConnection();
+ try {
+ BigQueryResult bigQueryResult = connection.executeSelect(INVALID_QUERY);
+ fail("BigQuerySQLException was expected");
+ } catch (BigQuerySQLException ex) {
+ assertNotNull(ex.getMessage());
+ assertTrue(ex.getMessage().toLowerCase().contains("unexpected keyword into"));
+ } finally {
+ connection.close();
+ }
+ }
+
+ /*
+ This tests for the order of the records as well as the value of the records using testForAllDataTypeValues
+ */
+ @Test
+ public void testIterateAndOrder() throws SQLException {
+ Connection connection = getConnection();
+ BigQueryResult bigQueryResult = connection.executeSelect(QUERY);
+ logger.log(Level.INFO, "Query used: {0}", QUERY);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+
+ int prevIntegerFieldVal = 0;
+ while (rs.next()) {
+ if (cnt == 0) { // first row is supposed to be null
+ assertNull(rs.getString("StringField"));
+ assertNull(rs.getString("GeographyField"));
+ Object intAryField = rs.getObject("IntegerArrayField");
+ if (intAryField instanceof JsonStringArrayList) {
+ assertEquals(
+ new JsonStringArrayList(),
+ ((JsonStringArrayList) intAryField)); // null array is returned as an empty array
+ }
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d == rs.getDouble("BigNumericField"));
+ assertTrue(0 == rs.getInt("IntegerField"));
+ assertTrue(0L == rs.getLong("NumericField"));
+ assertNull(rs.getBytes("BytesField"));
+ assertNull(rs.getTimestamp("TimestampField"));
+ assertNull(rs.getTime("TimeField"));
+ assertNull(rs.getDate("DateField"));
+ assertNull(rs.getString("JSONField"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertNull(rs.getString("StringField_1"));
+ assertNull(rs.getString("hello")); // equivalent of testJsonType
+ assertEquals(0, rs.getInt("id"));
+
+ } else { // remaining rows are supposed to be non null
+ assertNotNull(rs.getString("StringField"));
+ assertNotNull(rs.getString("GeographyField"));
+ assertNotNull(rs.getObject("IntegerArrayField"));
+ assertTrue(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d < rs.getDouble("BigNumericField"));
+ assertTrue(0 < rs.getInt("IntegerField"));
+ assertTrue(0L < rs.getLong("NumericField"));
+ assertNotNull(rs.getBytes("BytesField"));
+ assertNotNull(rs.getTimestamp("TimestampField"));
+ assertNotNull(rs.getTime("TimeField"));
+ assertNotNull(rs.getDate("DateField"));
+ assertNotNull(rs.getString("JSONField"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertNotNull(rs.getString("StringField_1"));
+
+ // check the order of the records
+ assertTrue(prevIntegerFieldVal < rs.getInt("IntegerField"));
+ prevIntegerFieldVal = rs.getInt("IntegerField");
+
+ testForAllDataTypeValues(rs, cnt); // asserts the value of each row
+ }
+ ++cnt;
+ }
+ assertEquals(LIMIT_RECS, cnt); // all the records were retrieved
+ connection.close();
+ }
+
+ /*
+ This tests for the order of the records using default connection settings as well as the value of the records using testForAllDataTypeValues
+ */
+ @Test
+ public void testIterateAndOrderDefaultConnSettings() throws SQLException {
+ Connection connection = bigquery.createConnection();
+ BigQueryResult bigQueryResult = connection.executeSelect(QUERY);
+ logger.log(Level.INFO, "Query used: {0}", QUERY);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+
+ int prevIntegerFieldVal = 0;
+ while (rs.next()) {
+ if (cnt == 0) { // first row is supposed to be null
+ assertNull(rs.getString("StringField"));
+ assertNull(rs.getString("GeographyField"));
+ Object intAryField = rs.getObject("IntegerArrayField");
+ if (intAryField instanceof JsonStringArrayList) {
+ assertEquals(
+ new JsonStringArrayList(),
+ ((JsonStringArrayList) intAryField)); // null array is returned as an empty array
+ }
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d == rs.getDouble("BigNumericField"));
+ assertTrue(0 == rs.getInt("IntegerField"));
+ assertTrue(0L == rs.getLong("NumericField"));
+ assertNull(rs.getBytes("BytesField"));
+ assertNull(rs.getTimestamp("TimestampField"));
+ assertNull(rs.getTime("TimeField"));
+ assertNull(rs.getDate("DateField"));
+ assertNull(rs.getString("JSONField"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertNull(rs.getString("StringField_1"));
+ assertNull(rs.getString("hello")); // equivalent of testJsonType
+ assertEquals(0, rs.getInt("id"));
+
+ } else { // remaining rows are supposed to be non null
+ assertNotNull(rs.getString("StringField"));
+ assertNotNull(rs.getString("GeographyField"));
+ assertNotNull(rs.getObject("IntegerArrayField"));
+ assertTrue(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d < rs.getDouble("BigNumericField"));
+ assertTrue(0 < rs.getInt("IntegerField"));
+ assertTrue(0L < rs.getLong("NumericField"));
+ assertNotNull(rs.getBytes("BytesField"));
+ assertNotNull(rs.getTimestamp("TimestampField"));
+ assertNotNull(rs.getTime("TimeField"));
+ assertNotNull(rs.getDate("DateField"));
+ assertNotNull(rs.getString("JSONField"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertNotNull(rs.getString("StringField_1"));
+
+ // check the order of the records
+ assertTrue(prevIntegerFieldVal < rs.getInt("IntegerField"));
+ prevIntegerFieldVal = rs.getInt("IntegerField");
+
+ testForAllDataTypeValues(rs, cnt); // asserts the value of each row
+ }
+ ++cnt;
+ }
+ assertEquals(LIMIT_RECS, cnt); // all the records were retrieved
+ connection.close();
+ }
+
+ @Test
+ public void testMultipleRuns() throws SQLException {
+
+ Connection connection = getConnection();
+ BigQueryResult bigQueryResult = connection.executeSelect(MULTI_QUERY);
+ logger.log(Level.INFO, "Query used: {0}", MULTI_QUERY);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ int totalCnt = 0;
+
+ int prevIntegerFieldVal = 0;
+ while (rs.next()) {
+ if (cnt == 0) { // first row is supposed to be null
+ assertNull(rs.getString("StringField"));
+ assertNull(rs.getString("GeographyField"));
+ Object intAryField = rs.getObject("IntegerArrayField");
+ if (intAryField instanceof JsonStringArrayList) {
+ assertEquals(
+ new JsonStringArrayList(),
+ ((JsonStringArrayList) intAryField)); // null array is returned as an empty array
+ }
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d == rs.getDouble("BigNumericField"));
+ assertTrue(0 == rs.getInt("IntegerField"));
+ assertTrue(0L == rs.getLong("NumericField"));
+ assertNull(rs.getBytes("BytesField"));
+ assertNull(rs.getTimestamp("TimestampField"));
+ assertNull(rs.getTime("TimeField"));
+ assertNull(rs.getDate("DateField"));
+ assertNull(rs.getString("JSONField"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertNull(rs.getString("StringField_1"));
+ assertNull(rs.getString("hello")); // equivalent of testJsonType
+ assertEquals(0, rs.getInt("id"));
+
+ } else { // remaining rows are supposed to be non null
+ // check the order of the records
+ assertTrue(prevIntegerFieldVal < rs.getInt("IntegerField"));
+ prevIntegerFieldVal = rs.getInt("IntegerField");
+
+ testForAllDataTypeValues(rs, cnt); // asserts the value of each row
+ }
+ ++cnt;
+ }
+ connection.close();
+ totalCnt += cnt;
+ // Repeat the same run
+ connection = getConnection();
+ bigQueryResult = connection.executeSelect(MULTI_QUERY);
+ rs = bigQueryResult.getResultSet();
+ cnt = 0;
+ prevIntegerFieldVal = 0;
+ while (rs.next()) {
+ if (cnt == 0) { // first row is supposed to be null
+ assertNull(rs.getString("StringField"));
+ assertNull(rs.getString("GeographyField"));
+ Object intAryField = rs.getObject("IntegerArrayField");
+ if (intAryField instanceof JsonStringArrayList) {
+ assertEquals(
+ new JsonStringArrayList(),
+ ((JsonStringArrayList) intAryField)); // null array is returned as an empty array
+ }
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d == rs.getDouble("BigNumericField"));
+ assertTrue(0 == rs.getInt("IntegerField"));
+ assertTrue(0L == rs.getLong("NumericField"));
+ assertNull(rs.getBytes("BytesField"));
+ assertNull(rs.getTimestamp("TimestampField"));
+ assertNull(rs.getTime("TimeField"));
+ assertNull(rs.getDate("DateField"));
+ assertNull(rs.getString("JSONField"));
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertNull(rs.getString("StringField_1"));
+ assertNull(rs.getString("hello")); // equivalent of testJsonType
+ assertEquals(0, rs.getInt("id"));
+
+ } else { // remaining rows are supposed to be non null
+ // check the order of the records
+ assertTrue(prevIntegerFieldVal < rs.getInt("IntegerField"));
+ prevIntegerFieldVal = rs.getInt("IntegerField");
+
+ testForAllDataTypeValues(rs, cnt); // asserts the value of each row
+ }
+ ++cnt;
+ }
+ connection.close();
+ totalCnt += cnt;
+ assertEquals(MULTI_LIMIT_RECS * 2, totalCnt);
+ }
+
+ @Test
+ public void testPositionalParams()
+ throws SQLException { // Bypasses Read API as it doesnt support Positional Params
+ Connection connection = getConnection();
+ Parameter dateParam =
+ Parameter.newBuilder().setValue(QueryParameterValue.date("2022-01-01")).build();
+ Parameter boolParam = Parameter.newBuilder().setValue(QueryParameterValue.bool(true)).build();
+ Parameter intParam = Parameter.newBuilder().setValue(QueryParameterValue.int64(1)).build();
+ Parameter numericParam =
+ Parameter.newBuilder().setValue(QueryParameterValue.numeric(new BigDecimal(100))).build();
+ List parameters = ImmutableList.of(dateParam, boolParam, intParam, numericParam);
+
+ BigQueryResult bigQueryResult = connection.executeSelect(POSITIONAL_QUERY, parameters);
+ logger.log(Level.INFO, "Query used: {0}", POSITIONAL_QUERY);
+ ResultSet rs = bigQueryResult.getResultSet();
+ int cnt = 0;
+ while (rs.next()) {
+ assertFalse(rs.getBoolean("BooleanField"));
+ assertTrue(0.0d <= rs.getDouble("BigNumericField"));
+ assertTrue(0 <= rs.getInt("IntegerField"));
+ assertTrue(0L <= rs.getLong("NumericField"));
+ assertNotNull(rs.getBytes("BytesField"));
+ assertNotNull(rs.getTimestamp("TimestampField"));
+ assertNotNull(rs.getTime("TimeField"));
+ assertNotNull(rs.getDate("DateField"));
+ assertNotNull(rs.getString("JSONField"));
+ assertTrue(rs.getBoolean("BooleanField_1"));
+ assertNotNull(rs.getString("StringField_1"));
+ ++cnt;
+ }
+ connection.close();
+ assertEquals(MULTI_LIMIT_RECS, cnt);
+ }
+
+ // asserts the value of each row
+ private static void testForAllDataTypeValues(ResultSet rs, int cnt) throws SQLException {
+ // Testing JSON type
+ assertEquals("\"world\"", rs.getString("hello")); // BQ stores the value as "world"
+ assertEquals(100, rs.getInt("id"));
+ assertEquals("{\"hello\":\"world\",\"id\":100}", rs.getString("JSONField"));
+
+ // String and Geography types
+ assertEquals(String.format("String Val %s", cnt), rs.getString("StringField"));
+ assertEquals("POINT(1 2)", rs.getString("GeographyField"));
+
+ // Array type tests
+ if (rs.getObject("IntegerArrayField") instanceof JsonStringArrayList) {
+ JsonStringArrayList ary = (JsonStringArrayList) rs.getObject("IntegerArrayField");
+ assertEquals(3, ary.size());
+ assertEquals(1, ary.get(0).intValue());
+ assertEquals(2, ary.get(1).intValue());
+ assertEquals(3, ary.get(2).intValue());
+ }
+
+ // BigNumeric, int and Numeric
+ assertTrue(10000000L + cnt == rs.getDouble("BigNumericField"));
+ assertEquals(1 + cnt, rs.getInt("IntegerField"));
+ assertEquals(100 + cnt, rs.getLong("NumericField"));
+ // Test Byte field
+ assertEquals("TestByteValue", new String(rs.getBytes("BytesField"), StandardCharsets.UTF_8));
+
+ // Struct Fields
+ assertFalse(rs.getBoolean("BooleanField_1"));
+ assertEquals(String.format("Str Val %s", cnt), rs.getString("StringField_1"));
+
+ // Timestamp, Time, DateTime and Date fields
+ assertEquals(1649064795000L, rs.getTimestamp("TimestampField").getTime());
+ assertEquals(
+ java.sql.Date.valueOf("2022-01-01").toString(), rs.getDate("DateField").toString());
+ // Time is represented independent of a specific date and timezone. For example a 12:11:35 (GMT)
+ // is returned as
+ // 17:11:35 (GMT+5:30) . So we need to adjust the offset
+ int offset =
+ TimeZone.getTimeZone(ZoneId.systemDefault())
+ .getOffset(new java.util.Date().getTime()); // offset in seconds
+ assertEquals(
+ Time.valueOf(LocalTime.of(12, 11, 35)).getTime() + offset,
+ rs.getTime("TimeField").getTime());
+ }
+
+ private static void populateTestRecords(String datasetName, String tableName) {
+ TableId tableId = TableId.of(datasetName, tableName);
+ for (int batchCnt = 1; batchCnt <= NUM_BATCHES; batchCnt++) {
+ addBatchRecords(tableId);
+ }
+ }
+
+ private static void addBatchRecords(TableId tableId) {
+ Map nullRow = new HashMap<>();
+ try {
+ InsertAllRequest.Builder reqBuilder = InsertAllRequest.newBuilder(tableId);
+ if (rowCnt == 0) {
+ reqBuilder.addRow(nullRow);
+ }
+ for (int i = 0; i < REC_PER_BATCHES; i++) {
+ reqBuilder.addRow(getNextRow());
+ }
+ InsertAllResponse response = bigquery.insertAll(reqBuilder.build());
+
+ if (response.hasErrors()) {
+ // If any of the insertions failed, this lets you inspect the errors
+ for (Map.Entry> entry : response.getInsertErrors().entrySet()) {
+ logger.log(Level.WARNING, "Exception while adding records {0}", entry.getValue());
+ }
+ fail("Response has errors");
+ }
+ } catch (BigQueryException e) {
+ logger.log(Level.WARNING, "Exception while adding records {0}", e);
+ fail("Error in addBatchRecords");
+ }
+ }
+
+ private static void createTable(String datasetName, String tableName, Schema schema) {
+ try {
+ TableId tableId = TableId.of(datasetName, tableName);
+ TableDefinition tableDefinition = StandardTableDefinition.of(schema);
+ TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();
+ Table table = bigquery.create(tableInfo);
+ assertTrue(table.exists());
+ } catch (BigQueryException e) {
+ fail("Table was not created. \n" + e);
+ }
+ }
+
+ public static void deleteTable(String datasetName, String tableName) {
+ try {
+ assertTrue(bigquery.delete(TableId.of(datasetName, tableName)));
+ } catch (BigQueryException e) {
+ fail("Table was not deleted. \n" + e);
+ }
+ }
+
+ public static void createDataset(String datasetName) {
+ try {
+ DatasetInfo datasetInfo = DatasetInfo.newBuilder(datasetName).build();
+ Dataset newDataset = bigquery.create(datasetInfo);
+ assertNotNull(newDataset.getDatasetId().getDataset());
+ } catch (BigQueryException e) {
+ fail("Dataset was not created. \n" + e);
+ }
+ }
+
+ public static void deleteDataset(String datasetName) {
+ try {
+ DatasetInfo datasetInfo = DatasetInfo.newBuilder(datasetName).build();
+ assertTrue(bigquery.delete(datasetInfo.getDatasetId()));
+ } catch (BigQueryException e) {
+ fail("Dataset was not deleted. \n" + e);
+ }
+ }
+
+ private Connection getConnection() {
+
+ ConnectionSettings connectionSettings =
+ ConnectionSettings.newBuilder()
+ .setDefaultDataset(DatasetId.of(DATASET))
+ .build(); // Read API is enabled by default
+ return bigquery.createConnection(connectionSettings);
+ }
+
+ private static Map getNextRow() {
+ rowCnt++;
+ Map row = new HashMap<>();
+ Map structVal = new HashMap<>();
+ structVal.put("StringField", "Str Val " + rowCnt);
+ structVal.put("BooleanField", false);
+ row.put("RecordField", structVal); // struct
+ row.put("TimestampField", "2022-04-04 15:03:15.000 +05:30");
+ row.put("StringField", "String Val " + rowCnt);
+ row.put("IntegerArrayField", new int[] {1, 2, 3});
+ row.put("BooleanField", true);
+ row.put("BytesField", BYTES_BASE64);
+ row.put("IntegerField", 1 + rowCnt);
+ row.put("GeographyField", "POINT(1 2)");
+ row.put("NumericField", 100 + rowCnt);
+ row.put("BigNumericField", 10000000L + rowCnt);
+ row.put("TimeField", "12:11:35");
+ row.put("DateField", "2022-01-01");
+ row.put("JSONField", "{\"hello\":\"world\",\"id\":100}");
+ row.put("IntervalField", "10000-0 3660000 87840000:0:0");
+ return row;
+ }
+}
diff --git a/pom.xml b/pom.xml
index ad8867220..5283f06a9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -66,6 +66,21 @@
pom
import