From 2156f023b4ab95bc7ec669545b5709317555fdac Mon Sep 17 00:00:00 2001 From: Phong Chuong <147636638+PhongChuong@users.noreply.github.com> Date: Tue, 6 Feb 2024 11:12:25 -0500 Subject: [PATCH] Feat: Add queryId to TableResult (#3106) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Feat: Add queryId to TableResult * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Feat: Add queryId to TableResult * Feat: Add IT for TableResult with JobId and QueryId. * Feat: Add IT for TableResult with JobId and QueryId. --------- Co-authored-by: Owl Bot --- .../clirr-ignored-differences.xml | 13 ++++ .../google/cloud/bigquery/BigQueryImpl.java | 8 ++- .../cloud/bigquery/EmptyTableResult.java | 2 +- .../google/cloud/bigquery/TableResult.java | 19 +++++- .../com/google/cloud/bigquery/JobTest.java | 2 +- .../cloud/bigquery/SerializationTest.java | 2 +- .../cloud/bigquery/TableResultTest.java | 4 +- .../com/google/cloud/bigquery/TableTest.java | 9 +-- .../cloud/bigquery/it/ITBigQueryTest.java | 59 ++++++++++++++++--- 9 files changed, 95 insertions(+), 23 deletions(-) diff --git a/google-cloud-bigquery/clirr-ignored-differences.xml b/google-cloud-bigquery/clirr-ignored-differences.xml index 9c69fd6a6..75441b537 100644 --- a/google-cloud-bigquery/clirr-ignored-differences.xml +++ b/google-cloud-bigquery/clirr-ignored-differences.xml @@ -14,6 +14,19 @@ com.google.api.services.bigquery.model.GetQueryResultsResponse getQueryResultsWithRowLimit(java.lang.String, java.lang.String, java.lang.String, java.lang.Integer) getQueryResultsWithRowLimit is just used by ConnectionImpl at the moment so it should be fine to update the signature instead of writing an overloaded method + + 7004 + com/google/cloud/bigquery/TableResult* + *TableResult(*) + It should be fine to update TableResult constructors since it is used to return results to the user and users should not directly construct TableResult objects + + + 7005 + com/google/cloud/bigquery/TableResult* + *TableResult(*) + *TableResult(*) + It should be fine to update TableResult constructors since it is used to return results to the user and users should not directly construct TableResult objects + 7013 com/google/cloud/bigquery/ExternalTableDefinition* 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 dcbcd19b0..30fd24f93 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 @@ -1156,7 +1156,7 @@ public TableResult listTableData( public TableResult listTableData(TableId tableId, Schema schema, TableDataListOption... options) { Tuple, Long> data = listTableData(tableId, schema, getOptions(), optionMap(options)); - return new TableResult(schema, data.y(), data.x()); + return new TableResult(schema, data.y(), data.x(), null); } private static Tuple, Long> listTableData( @@ -1410,7 +1410,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() { // cache first page of result transformTableData(results.getRows(), schema)), // Return the JobID of the successful job - jobId); + jobId, + results.getQueryId()); } // only 1 page of result return new TableResult( @@ -1421,7 +1422,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() { null, transformTableData(results.getRows(), schema)), // Return the JobID of the successful job - results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null); + results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null, + results.getQueryId()); } @Override diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/EmptyTableResult.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/EmptyTableResult.java index 7cb5e1932..a5c6c4785 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/EmptyTableResult.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/EmptyTableResult.java @@ -27,6 +27,6 @@ public class EmptyTableResult extends TableResult { /** An empty {@code TableResult} to avoid making API requests to unlistable tables. */ @InternalApi("Exposed for testing") public EmptyTableResult(@Nullable Schema schema) { - super(schema, 0, new PageImpl(null, "", null)); + super(schema, 0, new PageImpl(null, "", null), null); } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java index 325e1fd0b..1631d3bd5 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java @@ -37,6 +37,8 @@ public class TableResult implements Page, Serializable { private final Page pageNoSchema; @Nullable private JobId jobId = null; + @Nullable private final String queryId; + // package-private so job id is not set outside the package. void setJobId(@Nullable JobId jobId) { this.jobId = jobId; @@ -46,24 +48,35 @@ public JobId getJobId() { return jobId; } + public String getQueryId() { + return queryId; + } + /** * If {@code schema} is non-null, {@code TableResult} adds the schema to {@code FieldValueList}s * when iterating through them. {@code pageNoSchema} must not be null. */ @InternalApi("Exposed for testing") - public TableResult(Schema schema, long totalRows, Page pageNoSchema) { + public TableResult( + Schema schema, long totalRows, Page pageNoSchema, String queryId) { this.schema = schema; this.totalRows = totalRows; this.pageNoSchema = checkNotNull(pageNoSchema); + this.queryId = queryId; } @InternalApi("Exposed for testing") public TableResult( - Schema schema, long totalRows, Page pageNoSchema, JobId jobId) { + Schema schema, + long totalRows, + Page pageNoSchema, + JobId jobId, + String queryId) { this.schema = schema; this.totalRows = totalRows; this.pageNoSchema = checkNotNull(pageNoSchema); this.jobId = jobId; + this.queryId = queryId; } /** Returns the schema of the results. Null if the schema is not supplied. */ @@ -92,7 +105,7 @@ public String getNextPageToken() { @Override public TableResult getNextPage() { if (pageNoSchema.hasNextPage()) { - return new TableResult(schema, totalRows, pageNoSchema.getNextPage()); + return new TableResult(schema, totalRows, pageNoSchema.getNextPage(), queryId); } return null; } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobTest.java index b4b17b7d9..08c4137d6 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobTest.java @@ -309,7 +309,7 @@ public void testWaitForAndGetQueryResults() throws InterruptedException { Job completedJob = expectedJob.toBuilder().setStatus(new JobStatus(JobStatus.State.RUNNING)).build(); Page singlePage = Pages.empty(); - TableResult result = new TableResult(Schema.of(), 1, singlePage); + TableResult result = new TableResult(Schema.of(), 1, singlePage, null); QueryResponse completedQuery = QueryResponse.newBuilder() .setCompleted(true) diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java index 0b0cd0207..d31b282e7 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java @@ -206,7 +206,7 @@ public class SerializationTest extends BaseSerializationTest { private static final FieldValue FIELD_VALUE = FieldValue.of(FieldValue.Attribute.PRIMITIVE, "value"); private static final TableResult TABLE_RESULT = - new TableResult(Schema.of(), 0L, new PageImpl(null, "", ImmutableList.of())); + new TableResult(Schema.of(), 0L, new PageImpl(null, "", ImmutableList.of()), null); private static final BigQuery BIGQUERY = BigQueryOptions.newBuilder().setProjectId("p1").build().getService(); private static final Dataset DATASET = diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java index 35a167af1..89b7d6837 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java @@ -53,7 +53,7 @@ private static FieldValueList newFieldValueList(String s) { @Test public void testNullSchema() { - TableResult result = new TableResult(null, 3, INNER_PAGE_0); + TableResult result = new TableResult(null, 3, INNER_PAGE_0, null); assertThat(result.getSchema()).isNull(); assertThat(result.hasNextPage()).isTrue(); assertThat(result.getNextPageToken()).isNotNull(); @@ -75,7 +75,7 @@ public void testNullSchema() { @Test public void testSchema() { - TableResult result = new TableResult(SCHEMA, 3, INNER_PAGE_0); + TableResult result = new TableResult(SCHEMA, 3, INNER_PAGE_0, null); assertThat(result.getSchema()).isEqualTo(SCHEMA); assertThat(result.hasNextPage()).isTrue(); assertThat(result.getNextPageToken()).isNotNull(); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableTest.java index b93ed770b..9aa195331 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableTest.java @@ -248,9 +248,10 @@ public void testInsertComplete() { @Test public void testList() { Page page = new PageImpl<>(null, "c", ROWS); - when(bigquery.listTableData(TABLE_ID1)).thenReturn(new TableResult(null, ROWS.size(), page)); + when(bigquery.listTableData(TABLE_ID1)) + .thenReturn(new TableResult(null, ROWS.size(), page, null)); when(bigquery.listTableData(TABLE_ID1, SCHEMA)) - .thenReturn(new TableResult(SCHEMA, ROWS.size(), page)); + .thenReturn(new TableResult(SCHEMA, ROWS.size(), page, null)); Page dataPage = table.list(); assertThat(dataPage.getValues()).containsExactlyElementsIn(ROWS).inOrder(); dataPage = table.list(SCHEMA); @@ -263,9 +264,9 @@ public void testList() { public void testListWithOptions() { Page page = new PageImpl<>(null, "c", ROWS); when(bigquery.listTableData(TABLE_ID1, BigQuery.TableDataListOption.pageSize(10L))) - .thenReturn(new TableResult(null, ROWS.size(), page)); + .thenReturn(new TableResult(null, ROWS.size(), page, null)); when(bigquery.listTableData(TABLE_ID1, SCHEMA, BigQuery.TableDataListOption.pageSize(10L))) - .thenReturn(new TableResult(SCHEMA, ROWS.size(), page)); + .thenReturn(new TableResult(SCHEMA, ROWS.size(), page, null)); Page dataPage = table.list(BigQuery.TableDataListOption.pageSize(10L)); assertThat(dataPage.getValues()).containsExactlyElementsIn(ROWS).inOrder(); 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 ba30ea0d4..8b5a45913 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 @@ -6312,21 +6312,29 @@ public void testStatelessQueries() throws InterruptedException { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); BigQuery bigQuery = bigqueryHelper.getOptions().getService(); - // simulate setting the QUERY_PREVIEW_ENABLED environment variable + // Simulate setting the QUERY_PREVIEW_ENABLED environment variable. bigQuery.getOptions().setQueryPreviewEnabled("TRUE"); - assertNull(executeSimpleQuery(bigQuery).getJobId()); + TableResult tableResult = executeSimpleQuery(bigQuery); + assertNotNull(tableResult.getQueryId()); + assertNull(tableResult.getJobId()); - // the flag should be case-insensitive + // The flag should be case-insensitive. bigQuery.getOptions().setQueryPreviewEnabled("tRuE"); - assertNull(executeSimpleQuery(bigQuery).getJobId()); + tableResult = executeSimpleQuery(bigQuery); + assertNotNull(tableResult.getQueryId()); + assertNull(tableResult.getJobId()); - // any other values won't enable optional job creation mode + // Any other values won't enable optional job creation mode. bigQuery.getOptions().setQueryPreviewEnabled("test_value"); - assertNotNull(executeSimpleQuery(bigQuery).getJobId()); + tableResult = executeSimpleQuery(bigQuery); + assertNotNull(tableResult.getQueryId()); + assertNotNull(tableResult.getJobId()); - // reset the flag + // Reset the flag. bigQuery.getOptions().setQueryPreviewEnabled(null); - assertNotNull(executeSimpleQuery(bigQuery).getJobId()); + tableResult = executeSimpleQuery(bigQuery); + assertNotNull(tableResult.getQueryId()); + assertNotNull(tableResult.getJobId()); } private TableResult executeSimpleQuery(BigQuery bigQuery) throws InterruptedException { @@ -6336,6 +6344,41 @@ private TableResult executeSimpleQuery(BigQuery bigQuery) throws InterruptedExce return result; } + @Test + public void testTableResultJobIdAndQueryId() throws InterruptedException { + // For stateless queries, jobId and queryId are populated based on the following criteria: + // 1. For stateless queries, then queryId is populated. + // 2. For queries that fails the requirements to be stateless, then jobId is populated and + // queryId is not. + // 3. For explicitly created jobs, then jobId is populated and queryId is not populated. + + // Test scenario 1. + // Create local BigQuery for test scenario 1 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"); + String query = "SELECT 1 as one"; + QueryJobConfiguration configStateless = QueryJobConfiguration.newBuilder(query).build(); + TableResult result = bigQuery.query(configStateless); + assertNull(result.getJobId()); + assertNotNull(result.getQueryId()); + + // Test scenario 2 by failing stateless check by setting job timeout. + QueryJobConfiguration configQueryWithJob = + QueryJobConfiguration.newBuilder(query).setJobTimeoutMs(1L).build(); + result = bigQuery.query(configQueryWithJob); + assertNotNull(result.getJobId()); + assertNull(result.getQueryId()); + + // Test scenario 3. + QueryJobConfiguration configWithJob = QueryJobConfiguration.newBuilder(query).build(); + Job job = bigQuery.create(JobInfo.of(JobId.of(), configWithJob)); + result = job.getQueryResults(); + assertNotNull(result.getJobId()); + assertNull(result.getQueryId()); + } + @Test public void testStatelessQueriesWithLocation() throws Exception { // This test validates BigQueryOption location is used for stateless query by verifying that the