Skip to content

Commit 2156f02

Browse files
PhongChuonggcf-owl-bot[bot]
PhongChuong
and
gcf-owl-bot[bot]
authored
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>
1 parent 3abdc70 commit 2156f02

File tree

9 files changed

+95
-23
lines changed

9 files changed

+95
-23
lines changed

β€Žgoogle-cloud-bigquery/clirr-ignored-differences.xml

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,19 @@
1414
<method>com.google.api.services.bigquery.model.GetQueryResultsResponse getQueryResultsWithRowLimit(java.lang.String, java.lang.String, java.lang.String, java.lang.Integer)</method>
1515
<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>
1616
</difference>
17+
<difference>
18+
<differenceType>7004</differenceType>
19+
<className>com/google/cloud/bigquery/TableResult*</className>
20+
<method>*TableResult(*)</method>
21+
<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>
22+
</difference>
23+
<difference>
24+
<differenceType>7005</differenceType>
25+
<className>com/google/cloud/bigquery/TableResult*</className>
26+
<method>*TableResult(*)</method>
27+
<to>*TableResult(*)</to>
28+
<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>
29+
</difference>
1730
<difference>
1831
<differenceType>7013</differenceType>
1932
<className>com/google/cloud/bigquery/ExternalTableDefinition*</className>

β€Žgoogle-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1156,7 +1156,7 @@ public TableResult listTableData(
11561156
public TableResult listTableData(TableId tableId, Schema schema, TableDataListOption... options) {
11571157
Tuple<? extends Page<FieldValueList>, Long> data =
11581158
listTableData(tableId, schema, getOptions(), optionMap(options));
1159-
return new TableResult(schema, data.y(), data.x());
1159+
return new TableResult(schema, data.y(), data.x(), null);
11601160
}
11611161

11621162
private static Tuple<? extends Page<FieldValueList>, Long> listTableData(
@@ -1410,7 +1410,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
14101410
// cache first page of result
14111411
transformTableData(results.getRows(), schema)),
14121412
// Return the JobID of the successful job
1413-
jobId);
1413+
jobId,
1414+
results.getQueryId());
14141415
}
14151416
// only 1 page of result
14161417
return new TableResult(
@@ -1421,7 +1422,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
14211422
null,
14221423
transformTableData(results.getRows(), schema)),
14231424
// Return the JobID of the successful job
1424-
results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null);
1425+
results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null,
1426+
results.getQueryId());
14251427
}
14261428

14271429
@Override

β€Žgoogle-cloud-bigquery/src/main/java/com/google/cloud/bigquery/EmptyTableResult.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,6 @@ public class EmptyTableResult extends TableResult {
2727
/** An empty {@code TableResult} to avoid making API requests to unlistable tables. */
2828
@InternalApi("Exposed for testing")
2929
public EmptyTableResult(@Nullable Schema schema) {
30-
super(schema, 0, new PageImpl<FieldValueList>(null, "", null));
30+
super(schema, 0, new PageImpl<FieldValueList>(null, "", null), null);
3131
}
3232
}

β€Žgoogle-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ public class TableResult implements Page<FieldValueList>, Serializable {
3737
private final Page<FieldValueList> pageNoSchema;
3838
@Nullable private JobId jobId = null;
3939

40+
@Nullable private final String queryId;
41+
4042
// package-private so job id is not set outside the package.
4143
void setJobId(@Nullable JobId jobId) {
4244
this.jobId = jobId;
@@ -46,24 +48,35 @@ public JobId getJobId() {
4648
return jobId;
4749
}
4850

51+
public String getQueryId() {
52+
return queryId;
53+
}
54+
4955
/**
5056
* If {@code schema} is non-null, {@code TableResult} adds the schema to {@code FieldValueList}s
5157
* when iterating through them. {@code pageNoSchema} must not be null.
5258
*/
5359
@InternalApi("Exposed for testing")
54-
public TableResult(Schema schema, long totalRows, Page<FieldValueList> pageNoSchema) {
60+
public TableResult(
61+
Schema schema, long totalRows, Page<FieldValueList> pageNoSchema, String queryId) {
5562
this.schema = schema;
5663
this.totalRows = totalRows;
5764
this.pageNoSchema = checkNotNull(pageNoSchema);
65+
this.queryId = queryId;
5866
}
5967

6068
@InternalApi("Exposed for testing")
6169
public TableResult(
62-
Schema schema, long totalRows, Page<FieldValueList> pageNoSchema, JobId jobId) {
70+
Schema schema,
71+
long totalRows,
72+
Page<FieldValueList> pageNoSchema,
73+
JobId jobId,
74+
String queryId) {
6375
this.schema = schema;
6476
this.totalRows = totalRows;
6577
this.pageNoSchema = checkNotNull(pageNoSchema);
6678
this.jobId = jobId;
79+
this.queryId = queryId;
6780
}
6881

6982
/** Returns the schema of the results. Null if the schema is not supplied. */
@@ -92,7 +105,7 @@ public String getNextPageToken() {
92105
@Override
93106
public TableResult getNextPage() {
94107
if (pageNoSchema.hasNextPage()) {
95-
return new TableResult(schema, totalRows, pageNoSchema.getNextPage());
108+
return new TableResult(schema, totalRows, pageNoSchema.getNextPage(), queryId);
96109
}
97110
return null;
98111
}

β€Žgoogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,7 @@ public void testWaitForAndGetQueryResults() throws InterruptedException {
309309
Job completedJob =
310310
expectedJob.toBuilder().setStatus(new JobStatus(JobStatus.State.RUNNING)).build();
311311
Page<FieldValueList> singlePage = Pages.empty();
312-
TableResult result = new TableResult(Schema.of(), 1, singlePage);
312+
TableResult result = new TableResult(Schema.of(), 1, singlePage, null);
313313
QueryResponse completedQuery =
314314
QueryResponse.newBuilder()
315315
.setCompleted(true)

β€Žgoogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SerializationTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ public class SerializationTest extends BaseSerializationTest {
206206
private static final FieldValue FIELD_VALUE =
207207
FieldValue.of(FieldValue.Attribute.PRIMITIVE, "value");
208208
private static final TableResult TABLE_RESULT =
209-
new TableResult(Schema.of(), 0L, new PageImpl(null, "", ImmutableList.of()));
209+
new TableResult(Schema.of(), 0L, new PageImpl(null, "", ImmutableList.of()), null);
210210
private static final BigQuery BIGQUERY =
211211
BigQueryOptions.newBuilder().setProjectId("p1").build().getService();
212212
private static final Dataset DATASET =

β€Žgoogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ private static FieldValueList newFieldValueList(String s) {
5353

5454
@Test
5555
public void testNullSchema() {
56-
TableResult result = new TableResult(null, 3, INNER_PAGE_0);
56+
TableResult result = new TableResult(null, 3, INNER_PAGE_0, null);
5757
assertThat(result.getSchema()).isNull();
5858
assertThat(result.hasNextPage()).isTrue();
5959
assertThat(result.getNextPageToken()).isNotNull();
@@ -75,7 +75,7 @@ public void testNullSchema() {
7575

7676
@Test
7777
public void testSchema() {
78-
TableResult result = new TableResult(SCHEMA, 3, INNER_PAGE_0);
78+
TableResult result = new TableResult(SCHEMA, 3, INNER_PAGE_0, null);
7979
assertThat(result.getSchema()).isEqualTo(SCHEMA);
8080
assertThat(result.hasNextPage()).isTrue();
8181
assertThat(result.getNextPageToken()).isNotNull();

β€Žgoogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableTest.java

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -248,9 +248,10 @@ public void testInsertComplete() {
248248
@Test
249249
public void testList() {
250250
Page<FieldValueList> page = new PageImpl<>(null, "c", ROWS);
251-
when(bigquery.listTableData(TABLE_ID1)).thenReturn(new TableResult(null, ROWS.size(), page));
251+
when(bigquery.listTableData(TABLE_ID1))
252+
.thenReturn(new TableResult(null, ROWS.size(), page, null));
252253
when(bigquery.listTableData(TABLE_ID1, SCHEMA))
253-
.thenReturn(new TableResult(SCHEMA, ROWS.size(), page));
254+
.thenReturn(new TableResult(SCHEMA, ROWS.size(), page, null));
254255
Page<FieldValueList> dataPage = table.list();
255256
assertThat(dataPage.getValues()).containsExactlyElementsIn(ROWS).inOrder();
256257
dataPage = table.list(SCHEMA);
@@ -263,9 +264,9 @@ public void testList() {
263264
public void testListWithOptions() {
264265
Page<FieldValueList> page = new PageImpl<>(null, "c", ROWS);
265266
when(bigquery.listTableData(TABLE_ID1, BigQuery.TableDataListOption.pageSize(10L)))
266-
.thenReturn(new TableResult(null, ROWS.size(), page));
267+
.thenReturn(new TableResult(null, ROWS.size(), page, null));
267268
when(bigquery.listTableData(TABLE_ID1, SCHEMA, BigQuery.TableDataListOption.pageSize(10L)))
268-
.thenReturn(new TableResult(SCHEMA, ROWS.size(), page));
269+
.thenReturn(new TableResult(SCHEMA, ROWS.size(), page, null));
269270
Page<FieldValueList> dataPage = table.list(BigQuery.TableDataListOption.pageSize(10L));
270271
assertThat(dataPage.getValues()).containsExactlyElementsIn(ROWS).inOrder();
271272

β€Žgoogle-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6312,21 +6312,29 @@ public void testStatelessQueries() throws InterruptedException {
63126312
RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
63136313
BigQuery bigQuery = bigqueryHelper.getOptions().getService();
63146314

6315-
// simulate setting the QUERY_PREVIEW_ENABLED environment variable
6315+
// Simulate setting the QUERY_PREVIEW_ENABLED environment variable.
63166316
bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
6317-
assertNull(executeSimpleQuery(bigQuery).getJobId());
6317+
TableResult tableResult = executeSimpleQuery(bigQuery);
6318+
assertNotNull(tableResult.getQueryId());
6319+
assertNull(tableResult.getJobId());
63186320

6319-
// the flag should be case-insensitive
6321+
// The flag should be case-insensitive.
63206322
bigQuery.getOptions().setQueryPreviewEnabled("tRuE");
6321-
assertNull(executeSimpleQuery(bigQuery).getJobId());
6323+
tableResult = executeSimpleQuery(bigQuery);
6324+
assertNotNull(tableResult.getQueryId());
6325+
assertNull(tableResult.getJobId());
63226326

6323-
// any other values won't enable optional job creation mode
6327+
// Any other values won't enable optional job creation mode.
63246328
bigQuery.getOptions().setQueryPreviewEnabled("test_value");
6325-
assertNotNull(executeSimpleQuery(bigQuery).getJobId());
6329+
tableResult = executeSimpleQuery(bigQuery);
6330+
assertNotNull(tableResult.getQueryId());
6331+
assertNotNull(tableResult.getJobId());
63266332

6327-
// reset the flag
6333+
// Reset the flag.
63286334
bigQuery.getOptions().setQueryPreviewEnabled(null);
6329-
assertNotNull(executeSimpleQuery(bigQuery).getJobId());
6335+
tableResult = executeSimpleQuery(bigQuery);
6336+
assertNotNull(tableResult.getQueryId());
6337+
assertNotNull(tableResult.getJobId());
63306338
}
63316339

63326340
private TableResult executeSimpleQuery(BigQuery bigQuery) throws InterruptedException {
@@ -6336,6 +6344,41 @@ private TableResult executeSimpleQuery(BigQuery bigQuery) throws InterruptedExce
63366344
return result;
63376345
}
63386346

6347+
@Test
6348+
public void testTableResultJobIdAndQueryId() throws InterruptedException {
6349+
// For stateless queries, jobId and queryId are populated based on the following criteria:
6350+
// 1. For stateless queries, then queryId is populated.
6351+
// 2. For queries that fails the requirements to be stateless, then jobId is populated and
6352+
// queryId is not.
6353+
// 3. For explicitly created jobs, then jobId is populated and queryId is not populated.
6354+
6355+
// Test scenario 1.
6356+
// Create local BigQuery for test scenario 1 to not contaminate global test parameters.
6357+
RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
6358+
BigQuery bigQuery = bigqueryHelper.getOptions().getService();
6359+
// Simulate setting the QUERY_PREVIEW_ENABLED environment variable.
6360+
bigQuery.getOptions().setQueryPreviewEnabled("TRUE");
6361+
String query = "SELECT 1 as one";
6362+
QueryJobConfiguration configStateless = QueryJobConfiguration.newBuilder(query).build();
6363+
TableResult result = bigQuery.query(configStateless);
6364+
assertNull(result.getJobId());
6365+
assertNotNull(result.getQueryId());
6366+
6367+
// Test scenario 2 by failing stateless check by setting job timeout.
6368+
QueryJobConfiguration configQueryWithJob =
6369+
QueryJobConfiguration.newBuilder(query).setJobTimeoutMs(1L).build();
6370+
result = bigQuery.query(configQueryWithJob);
6371+
assertNotNull(result.getJobId());
6372+
assertNull(result.getQueryId());
6373+
6374+
// Test scenario 3.
6375+
QueryJobConfiguration configWithJob = QueryJobConfiguration.newBuilder(query).build();
6376+
Job job = bigQuery.create(JobInfo.of(JobId.of(), configWithJob));
6377+
result = job.getQueryResults();
6378+
assertNotNull(result.getJobId());
6379+
assertNull(result.getQueryId());
6380+
}
6381+
63396382
@Test
63406383
public void testStatelessQueriesWithLocation() throws Exception {
63416384
// This test validates BigQueryOption location is used for stateless query by verifying that the

0 commit comments

Comments
 (0)