feat: add fast query path support when empty jobId object is passed (#2349)

Neenu1995 · gcf-owl-bot[bot] · web-flow · commit 42c083ac680c · 2022-11-11T14:16:08.000-05:00
* feat: add fast query path support with job location * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: fix NPE * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: respect user specified location * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: fix test * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore add IT tests * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: allow users to specify project ID with JobID * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: fix conditions * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
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
@@ -1300,7 +1300,7 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio
     // If all parameters passed in configuration are supported by the query() method on the backend,
     // put on fast path
     QueryRequestInfo requestInfo = new QueryRequestInfo(configuration);
-    if (requestInfo.isFastQuerySupported()) {
+    if (requestInfo.isFastQuerySupported(null)) {
       String projectId = getOptions().getProjectId();
       QueryRequest content = requestInfo.toPb();
       return queryRpc(projectId, content, options);
@@ -1385,6 +1385,27 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
   public TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOption... options)
       throws InterruptedException, JobException {
     Job.checkNotDryRun(configuration, "query");
+    // If all parameters passed in configuration are supported by the query() method on the backend,
+    // put on fast path
+    QueryRequestInfo requestInfo = new QueryRequestInfo(configuration);
+    if (requestInfo.isFastQuerySupported(jobId)) {
+      // Be careful when setting the projectID in JobId, if a projectID is specified in the JobId,
+      // the job created by the query method will use that project. This may cause the query to
+      // fail with "Access denied" if the project do not have enough permissions to run the job.
+
+      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
+      // 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());
+      }
+
+      return queryRpc(projectId, content, options);
+    }
     return create(JobInfo.of(jobId, configuration), options).getQueryResults();
   }
 
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java
@@ -57,7 +57,15 @@ final class QueryRequestInfo {
     this.useQueryCache = config.useQueryCache();
   }
 
-  boolean isFastQuerySupported() {
+  boolean isFastQuerySupported(JobId jobId) {
+    // Fast query path is not possible if job is specified in the JobID object
+    // Respect Job field value in JobId specified by user.
+    // Specifying it will force the query to take the slower path.
+    if (jobId != null) {
+      if (jobId.getJob() != null) {
+        return false;
+      }
+    }
     return config.getClustering() == null
         && config.getCreateDisposition() == null
         && config.getDestinationEncryptionConfiguration() == null
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java
@@ -150,8 +150,12 @@ public class QueryRequestInfoTest {
 
   @Test
   public void testIsFastQuerySupported() {
-    assertEquals(false, REQUEST_INFO.isFastQuerySupported());
-    assertEquals(true, REQUEST_INFO_SUPPORTED.isFastQuerySupported());
+    JobId jobIdSupported = JobId.newBuilder().build();
+    JobId jobIdNotSupported = JobId.newBuilder().setJob("random-job-id").build();
+    assertEquals(false, REQUEST_INFO.isFastQuerySupported(jobIdSupported));
+    assertEquals(true, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdSupported));
+    assertEquals(false, REQUEST_INFO.isFastQuerySupported(jobIdNotSupported));
+    assertEquals(false, REQUEST_INFO_SUPPORTED.isFastQuerySupported(jobIdNotSupported));
   }
 
   @Test
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
@@ -180,6 +180,7 @@ public class ITBigQueryTest {
   private static final Long EXPIRATION_MS = 86400000L;
   private static final Logger LOG = Logger.getLogger(ITBigQueryTest.class.getName());
   private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
+  private static final String UK_DATASET = RemoteBigQueryHelper.generateDatasetName();
   private static final String DESCRIPTION = "Test dataset";
   private static final String OTHER_DATASET = RemoteBigQueryHelper.generateDatasetName();
   private static final String MODEL_DATASET = RemoteBigQueryHelper.generateDatasetName();
@@ -533,6 +534,8 @@ public class ITBigQueryTest {
   private static final TableId TABLE_ID = TableId.of(DATASET, "testing_table");
   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_FASTQUERY_UK =
+      TableId.of(UK_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");
@@ -717,6 +720,7 @@ public static void beforeClass() throws InterruptedException, IOException {
     DatasetInfo info3 =
         DatasetInfo.newBuilder(ROUTINE_DATASET).setDescription("java routine lifecycle").build();
     bigquery.create(info3);
+
     LoadJobConfiguration configuration =
         LoadJobConfiguration.newBuilder(
                 TABLE_ID, "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json())
@@ -781,6 +785,7 @@ public static void beforeClass() throws InterruptedException, IOException {
   public static void afterClass() throws ExecutionException, InterruptedException {
     if (bigquery != null) {
       RemoteBigQueryHelper.forceDelete(bigquery, DATASET);
+      RemoteBigQueryHelper.forceDelete(bigquery, UK_DATASET);
       RemoteBigQueryHelper.forceDelete(bigquery, MODEL_DATASET);
       RemoteBigQueryHelper.forceDelete(bigquery, ROUTINE_DATASET);
     }
@@ -3284,6 +3289,86 @@ public void testFastSQLQuery() throws InterruptedException {
     }
   }
 
+  @Test
+  public void testProjectIDFastSQLQueryWithJobId() throws InterruptedException {
+    String random_project_id = "RANDOM_PROJECT_" + UUID.randomUUID().toString().replace('-', '_');
+    System.out.println(random_project_id);
+    String query =
+        "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID_FASTQUERY.getTable();
+    // With incorrect projectID in jobid
+    // The job will be created with the specified(incorrect) projectID
+    // hence failing the operation
+    JobId jobIdWithProjectId = JobId.newBuilder().setProject(random_project_id).build();
+    QueryJobConfiguration configSelect =
+        QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
+    try {
+      bigquery.query(configSelect, jobIdWithProjectId);
+    } catch (Exception exception) {
+      // error message for non-existent project
+      assertTrue(exception.getMessage().contains("Cannot parse  as CloudRegion"));
+      assertEquals(BigQueryException.class, exception.getClass());
+    }
+  }
+
+  @Test
+  public void testLocationFastSQLQueryWithJobId() throws InterruptedException {
+    DatasetInfo infoUK =
+        DatasetInfo.newBuilder(UK_DATASET)
+            .setDescription(DESCRIPTION)
+            .setLocation("europe-west1")
+            .setLabels(LABELS)
+            .build();
+    bigquery.create(infoUK);
+
+    TableDefinition tableDefinition = StandardTableDefinition.of(SIMPLE_SCHEMA);
+    TableInfo tableInfo = TableInfo.newBuilder(TABLE_ID_FASTQUERY_UK, tableDefinition).build();
+    bigquery.create(tableInfo);
+
+    String insert =
+        "INSERT " + UK_DATASET + "." + TABLE_ID_FASTQUERY_UK.getTable() + " VALUES('Anna');";
+
+    QueryJobConfiguration config =
+        QueryJobConfiguration.newBuilder(insert)
+            .setDefaultDataset(DatasetId.of(UK_DATASET))
+            .build();
+    TableResult result = bigquery.query(config);
+    assertEquals(SIMPLE_SCHEMA, result.getSchema());
+    assertEquals(1, result.getTotalRows());
+    assertNull(result.getNextPage());
+    assertNull(result.getNextPageToken());
+    assertFalse(result.hasNextPage());
+    // Verify correctness of table content
+    for (FieldValueList row : result.getValues()) {
+      FieldValue stringCell = row.get(0);
+      assertEquals(stringCell, row.get("StringField"));
+      assertEquals("Anna", stringCell.getStringValue());
+    }
+    // With incorrect location in jobid
+    // The job will be created with the specified(incorrect) location
+    // hence failing the operation
+    String query = "SELECT StringField FROM " + TABLE_ID_FASTQUERY_UK.getTable();
+    JobId jobIdWithLocation = JobId.newBuilder().setLocation("us-west1").build();
+    QueryJobConfiguration configSelect =
+        QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(UK_DATASET)).build();
+    try {
+      bigquery.query(configSelect, jobIdWithLocation);
+    } catch (BigQueryException exception) {
+      assertTrue(exception.getMessage().contains("Not found"));
+      assertEquals(BigQueryException.class, exception.getClass());
+    }
+
+    // Without location in jobID, the query job defaults to the location of the dataset
+    JobId jobIdNoLocation = JobId.newBuilder().build();
+    QueryJobConfiguration configNoLocation =
+        QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(UK_DATASET)).build();
+    TableResult resultNoLocation = bigquery.query(configNoLocation, jobIdNoLocation);
+    for (FieldValueList row : resultNoLocation.getValues()) {
+      FieldValue stringCell = row.get(0);
+      assertEquals(stringCell, row.get("StringField"));
+      assertEquals("Anna", stringCell.getStringValue());
+    }
+  }
+
   /* TODO(prasmish): replicate the entire test case for executeSelect */
   @Test
   public void testFastSQLQueryMultiPage() throws InterruptedException {
