Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat: Add queryId to TableResult #3106

Merged
merged 6 commits into from Feb 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Expand Up @@ -53,7 +53,7 @@ If you are using Maven without the BOM, add this to your dependencies:
If you are using Gradle 5.x or later, add this to your dependencies:

```Groovy
implementation platform('com.google.cloud:libraries-bom:26.30.0')
implementation platform('com.google.cloud:libraries-bom:26.31.0')

implementation 'com.google.cloud:google-cloud-bigquery'
```
Expand Down
13 changes: 13 additions & 0 deletions google-cloud-bigquery/clirr-ignored-differences.xml
Expand Up @@ -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>
Expand Down
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down
Expand Up @@ -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);
}
}
Expand Up @@ -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;
Expand All @@ -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(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than solely extending the existing constructor, should we just provide an second signature or a builder? It would mean you can avoid having to touch all the callsites where queryId isn't relevant for construction.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As discussed offline. I think creating a builder would be the preferred option. I'll follow up with another PR to switch to using a builder.

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. */
Expand Down Expand Up @@ -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;
}
Expand Down
Expand Up @@ -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)
Expand Down
Expand Up @@ -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 =
Expand Down
Expand Up @@ -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();
Expand All @@ -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();
Expand Down
Expand Up @@ -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);
Expand All @@ -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();

Expand Down
Expand Up @@ -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());
Neenu1995 marked this conversation as resolved.
Show resolved Hide resolved
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());
Neenu1995 marked this conversation as resolved.
Show resolved Hide resolved
}

private TableResult executeSimpleQuery(BigQuery bigQuery) throws InterruptedException {
Expand All @@ -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
Expand Down