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>

Author: PhongChuong

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);

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());
   }

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());
+    }
+  }
 }