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 index 7d329e88f..dcbcd19b0 100644 --- 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 index 04bfa6d9f..62160bef2 100644 --- 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 index 7bdb7a7fb..9a0e3e107 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 @@ -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()); + } + } }