feat: Use location in BigQueryOption as the default for query (#3047)

* feat: Use location in BigQueryOption as the default for query * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * feat: Use location in BigQueryOption as the default for query * feat: Use location in BigQueryOption as the default for query * feat: Use location in BigQueryOption as the default for query --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · Dec 19, 2023 · 270f866 · 270f866
1 parent a710632
commit 270f866
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 12 deletions.
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
@@ -1339,6 +1339,9 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio
  if (requestInfo.isFastQuerySupported(null)) {
  String projectId = getOptions().getProjectId();
  QueryRequest content = requestInfo.toPb();
+ if (getOptions().getLocation() != null) {
+ content.setLocation(getOptions().getLocation());
+ }
  return queryRpc(projectId, content, options);
  }
  // Otherwise, fall back to the existing create query job logic
@@ -1437,12 +1440,14 @@ public TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOp
  String projectId =
  jobId.getProject() != null ? jobId.getProject() : getOptions().getProjectId();
  QueryRequest content = requestInfo.toPb();
- // Be careful when setting the location in JobId, if a location is specified in the JobId,
- // the job created by the query method will be in that location, even if the table to be
+ // Be careful when setting the location, if a location is specified in the BigQueryOption or
+ // JobId the job created by the query method will be in that location, even if the table to be
  // queried is in a different location. This may cause the query to fail with
  // "BigQueryException: Not found"
  if (jobId.getLocation() != null) {
  content.setLocation(jobId.getLocation());
+ } else if (getOptions().getLocation() != null) {
+ content.setLocation(getOptions().getLocation());
  }
 
  return queryRpc(projectId, content, options);

diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java
@@ -1994,6 +1994,47 @@ public void testFastQueryRequestCompleted() throws InterruptedException {
  QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(),
  requestPb.getDefaultDataset().getDatasetId());
  assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache());
+ assertNull(requestPb.getLocation());
+
+ verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture());
+ }
+
+ @Test
+ public void testFastQueryRequestCompletedWithLocation() throws InterruptedException {
+ com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
+ new com.google.api.services.bigquery.model.QueryResponse()
+ .setCacheHit(false)
+ .setJobComplete(true)
+ .setKind("bigquery#queryResponse")
+ .setPageToken(null)
+ .setRows(ImmutableList.of(TABLE_ROW))
+ .setSchema(TABLE_SCHEMA.toPb())
+ .setTotalBytesProcessed(42L)
+ .setTotalRows(BigInteger.valueOf(1L));
+
+ when(bigqueryRpcMock.queryRpc(eq(PROJECT), requestPbCapture.capture()))
+ .thenReturn(queryResponsePb);
+
+ BigQueryOptions options = createBigQueryOptionsForProjectWithLocation(PROJECT, rpcFactoryMock);
+ bigquery = options.getService();
+ TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY);
+ assertNull(result.getNextPage());
+ assertNull(result.getNextPageToken());
+ assertFalse(result.hasNextPage());
+ assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA);
+ assertThat(result.getTotalRows()).isEqualTo(1);
+ for (FieldValueList row : result.getValues()) {
+ assertThat(row.get(0).getBooleanValue()).isFalse();
+ assertThat(row.get(1).getLongValue()).isEqualTo(1);
+ }
+
+ QueryRequest requestPb = requestPbCapture.getValue();
+ assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.getQuery(), requestPb.getQuery());
+ assertEquals(
+ QUERY_JOB_CONFIGURATION_FOR_QUERY.getDefaultDataset().getDataset(),
+ requestPb.getDefaultDataset().getDatasetId());
+ assertEquals(QUERY_JOB_CONFIGURATION_FOR_QUERY.useQueryCache(), requestPb.getUseQueryCache());
+ assertEquals(LOCATION, requestPb.getLocation());
 
  verify(bigqueryRpcMock).queryRpc(eq(PROJECT), requestPbCapture.capture());
  }

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
@@ -6227,27 +6227,80 @@ public void testAlreadyExistJobExceptionHandling() throws InterruptedException {
 
  @Test
  public void testStatelessQueries() throws InterruptedException {
+ // Create local BigQuery to not contaminate global test parameters.
+ RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
+ BigQuery bigQuery = bigqueryHelper.getOptions().getService();
+
  // simulate setting the QUERY_PREVIEW_ENABLED environment variable
- bigquery.getOptions().setQueryPreviewEnabled("TRUE");
- assertNull(executeSimpleQuery().getJobId());
+ bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
+ assertNull(executeSimpleQuery(bigQuery).getJobId());
 
  // the flag should be case-insensitive
- bigquery.getOptions().setQueryPreviewEnabled("tRuE");
- assertNull(executeSimpleQuery().getJobId());
+ bigQuery.getOptions().setQueryPreviewEnabled("tRuE");
+ assertNull(executeSimpleQuery(bigQuery).getJobId());
 
  // any other values won't enable optional job creation mode
- bigquery.getOptions().setQueryPreviewEnabled("test_value");
- assertNotNull(executeSimpleQuery().getJobId());
+ bigQuery.getOptions().setQueryPreviewEnabled("test_value");
+ assertNotNull(executeSimpleQuery(bigQuery).getJobId());
 
  // reset the flag
- bigquery.getOptions().setQueryPreviewEnabled(null);
- assertNotNull(executeSimpleQuery().getJobId());
+ bigQuery.getOptions().setQueryPreviewEnabled(null);
+ assertNotNull(executeSimpleQuery(bigQuery).getJobId());
  }
 
- private TableResult executeSimpleQuery() throws InterruptedException {
+ private TableResult executeSimpleQuery(BigQuery bigQuery) throws InterruptedException {
  String query = "SELECT 1 as one";
  QueryJobConfiguration config = QueryJobConfiguration.newBuilder(query).build();
- TableResult result = bigquery.query(config);
+ TableResult result = bigQuery.query(config);
  return result;
  }
+
+ @Test
+ public void testStatelessQueriesWithLocation() throws Exception {
+ // This test validates BigQueryOption location is used for stateless query by verifying that the
+ // stateless query fails when the BigQueryOption location does not match the dataset location.
+ String location = "EU";
+ String wrongLocation = "US";
+
+ RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
+ BigQuery bigQuery =
+ bigqueryHelper.getOptions().toBuilder().setLocation(location).build().getService();
+
+ Dataset dataset =
+ bigQuery.create(
+ DatasetInfo.newBuilder("locationset_" + UUID.randomUUID().toString().replace("-", "_"))
+ .setLocation(location)
+ .build());
+ try {
+ TableId tableId = TableId.of(dataset.getDatasetId().getDataset(), "sometable");
+ Schema schema = Schema.of(Field.of("name", LegacySQLTypeName.STRING));
+ TableDefinition tableDef = StandardTableDefinition.of(schema);
+ Table table = bigQuery.create(TableInfo.newBuilder(tableId, tableDef).build());
+
+ String query =
+ String.format(
+ "SELECT * FROM `%s.%s.%s`",
+ table.getTableId().getProject(),
+ table.getTableId().getDataset(),
+ table.getTableId().getTable());
+
+ // Test stateless query when BigQueryOption location matches dataset location.
+ bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
+ TableResult tb = bigQuery.query(QueryJobConfiguration.of(query));
+ assertNull(tb.getJobId());
+
+ // Test stateless query when BigQueryOption location does not match dataset location.
+ try {
+ BigQuery bigQueryWrongLocation =
+ bigqueryHelper.getOptions().toBuilder().setLocation(wrongLocation).build().getService();
+ bigQueryWrongLocation.getOptions().setQueryPreviewEnabled("TRUE");
+ bigQueryWrongLocation.query(QueryJobConfiguration.of(query));
+ fail("querying a table with wrong location shouldn't work");
+ } catch (BigQueryException e) {
+ // Nothing to do
+ }
+ } finally {
+ bigQuery.delete(dataset.getDatasetId(), DatasetDeleteOption.deleteContents());
+ }
+ }
 }