Skip to content
Permalink
Browse files
fix: schema and totalRows for duplicate request ids (#804)
* fix: schema and totalRows for duplicate request ids

* fix: modify code and test case

* fix: address comment
  • Loading branch information
Praful Makani committed Oct 19, 2020
1 parent 137eeaa commit f2799dd9a4323460843aa6df7ac7c2f092c7f1b3
@@ -1239,7 +1239,8 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio
}

private TableResult queryRpc(
final String projectId, final QueryRequest content, JobOption... options) {
final String projectId, final QueryRequest content, JobOption... options)
throws InterruptedException {
com.google.api.services.bigquery.model.QueryResponse results;
try {
results =
@@ -1265,14 +1266,22 @@ public com.google.api.services.bigquery.model.QueryResponse call() {
throw new BigQueryException(bigQueryErrors);
}

Schema schema = results.getSchema() == null ? null : Schema.fromPb(results.getSchema());
Long numRows;
if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
numRows = 0L;
} else if (results.getNumDmlAffectedRows() != null) {
numRows = results.getNumDmlAffectedRows();
long numRows;
Schema schema;
if (results.getSchema() == null && results.getJobComplete()) {
JobId jobId = JobId.fromPb(results.getJobReference());
Job job = getJob(jobId, options);
TableResult tableResult = job.getQueryResults();
return tableResult;
} else {
numRows = results.getTotalRows().longValue();
schema = results.getSchema() == null ? null : Schema.fromPb(results.getSchema());
if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
numRows = 0L;
} else if (results.getNumDmlAffectedRows() != null) {
numRows = results.getNumDmlAffectedRows();
} else {
numRows = results.getTotalRows().longValue();
}
}

if (results.getPageToken() != null) {
@@ -1608,6 +1608,35 @@ public void testFastQueryMultipleRuns() throws InterruptedException {
assertFalse(result2.hasNextPage());
}

@Test
public void testFastQuerySinglePageDuplicateRequestIds() throws InterruptedException {
String query =
"SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID_FASTQUERY.getTable();
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build();
TableResult result = bigquery.query(config);
assertEquals(QUERY_RESULT_SCHEMA, result.getSchema());
assertEquals(2, result.getTotalRows());
assertNull(result.getNextPage());
assertNull(result.getNextPageToken());
assertFalse(result.hasNextPage());

TableResult result1 = bigquery.query(config);
assertEquals(QUERY_RESULT_SCHEMA, result1.getSchema());
assertEquals(2, result1.getTotalRows());
assertNull(result1.getNextPage());
assertNull(result1.getNextPageToken());
assertFalse(result1.hasNextPage());

config.toBuilder().setQuery(query).build();
TableResult result2 = bigquery.query(config);
assertEquals(QUERY_RESULT_SCHEMA, result2.getSchema());
assertEquals(2, result2.getTotalRows());
assertNull(result2.getNextPage());
assertNull(result2.getNextPageToken());
assertFalse(result2.hasNextPage());
}

@Test
public void testFastSQLQuery() throws InterruptedException {
String query =
@@ -1650,6 +1679,21 @@ public void testFastSQLQueryMultiPage() throws InterruptedException {
assertNotNull(result.getNextPage());
assertNotNull(result.getNextPageToken());
assertTrue(result.hasNextPage());

TableResult result1 = bigquery.query(config);
assertEquals(LARGE_TABLE_SCHEMA, result.getSchema());
assertEquals(313348, result.getTotalRows());
assertNotNull(result1.getNextPage());
assertNotNull(result1.getNextPageToken());
assertTrue(result1.hasNextPage());

config.toBuilder().setQuery(query).build();
TableResult result2 = bigquery.query(config);
assertEquals(LARGE_TABLE_SCHEMA, result2.getSchema());
assertEquals(313348, result2.getTotalRows());
assertNotNull(result2.getNextPage());
assertNotNull(result2.getNextPageToken());
assertTrue(result2.hasNextPage());
}

@Test

0 comments on commit f2799dd

Please sign in to comment.