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 01048689f..b5ca2e577 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 @@ -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 index ccde465af..00a898363 100644 --- 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 index a62e3c0a0..456475597 100644 --- 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 index 4f72e7e66..ed073a93a 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 @@ -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 {