Feat: Add queryId to TableResult (#3106)

* 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 <gcf-owl-bot[bot]@users.noreply.github.com>

google-cloud-bigquery/clirr-ignored-differences.xml

@@ -14,6 +14,19 @@
     <method>com.google.api.services.bigquery.model.GetQueryResultsResponse getQueryResultsWithRowLimit(java.lang.String, java.lang.String, java.lang.String, java.lang.Integer)</method>
     <justification>getQueryResultsWithRowLimit is just used by ConnectionImpl at the moment so it should be fine to update the signature instead of writing an overloaded method</justification>
   </difference>
+  <difference>
+    <differenceType>7004</differenceType>
+    <className>com/google/cloud/bigquery/TableResult*</className>
+    <method>*TableResult(*)</method>
+    <justification>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</justification>
+  </difference>
+  <difference>
+    <differenceType>7005</differenceType>
+    <className>com/google/cloud/bigquery/TableResult*</className>
+    <method>*TableResult(*)</method>
+    <to>*TableResult(*)</to>
+    <justification>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</justification>
+  </difference>
   <difference>
     <differenceType>7013</differenceType>
     <className>com/google/cloud/bigquery/ExternalTableDefinition*</className>

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<? extends Page<FieldValueList>, 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<? extends Page<FieldValueList>, 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

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<FieldValueList>(null, "", null));
+    super(schema, 0, new PageImpl<FieldValueList>(null, "", null), null);
   }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java

@@ -37,6 +37,8 @@ public class TableResult implements Page<FieldValueList>, Serializable {
   private final Page<FieldValueList> 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<FieldValueList> pageNoSchema) {
+  public TableResult(
+      Schema schema, long totalRows, Page<FieldValueList> 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<FieldValueList> pageNoSchema, JobId jobId) {
+      Schema schema,
+      long totalRows,
+      Page<FieldValueList> 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;
   }

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

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 =

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

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableTest.java

@@ -248,9 +248,10 @@ public void testInsertComplete() {
   @Test
   public void testList() {
     Page<FieldValueList> 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<FieldValueList> dataPage = table.list();
     assertThat(dataPage.getValues()).containsExactlyElementsIn(ROWS).inOrder();
     dataPage = table.list(SCHEMA);
@@ -263,9 +264,9 @@ public void testList() {
   public void testListWithOptions() {
     Page<FieldValueList> 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<FieldValueList> dataPage = table.list(BigQuery.TableDataListOption.pageSize(10L));
     assertThat(dataPage.getValues()).containsExactlyElementsIn(ROWS).inOrder();
 

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