From 27311417db08a315a5fb9a576f001db3bf7a0431 Mon Sep 17 00:00:00 2001 From: Praful Makani Date: Thu, 15 Oct 2020 17:03:45 +0530 Subject: [PATCH 1/3] fix: schema and totalRows for duplicate request ids --- .../google/cloud/bigquery/BigQueryImpl.java | 22 +++++++++----- .../cloud/bigquery/it/ITBigQueryTest.java | 29 +++++++++++++++++++ 2 files changed, 44 insertions(+), 7 deletions(-) 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 76067be0d..590849bd2 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 @@ -1265,14 +1265,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()); + QueryResponse result = getQueryResults(jobId, getOptions(), optionMap(options)); + numRows = result.getTotalRows(); + schema = result.getSchema(); } else { - numRows = results.getTotalRows().longValue(); + schema = 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) { 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 05332ba6a..9bde02e62 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 @@ -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 = From 32d3d353fa9cccc8bbe563abef3d8461f4aa3ee1 Mon Sep 17 00:00:00 2001 From: Praful Makani Date: Fri, 16 Oct 2020 17:39:00 +0530 Subject: [PATCH 2/3] fix: modify code and test case --- .../com/google/cloud/bigquery/BigQueryImpl.java | 9 +++++---- .../google/cloud/bigquery/it/ITBigQueryTest.java | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) 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 590849bd2..6a5bd2392 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 @@ -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 = @@ -1269,9 +1270,9 @@ public com.google.api.services.bigquery.model.QueryResponse call() { Schema schema; if (results.getSchema() == null && results.getJobComplete()) { JobId jobId = JobId.fromPb(results.getJobReference()); - QueryResponse result = getQueryResults(jobId, getOptions(), optionMap(options)); - numRows = result.getTotalRows(); - schema = result.getSchema(); + Job job = getJob(jobId, options); + TableResult tableResult = job.getQueryResults(); + return tableResult; } else { schema = Schema.fromPb(results.getSchema()); if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) { 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 9bde02e62..36ad12d68 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 @@ -1679,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 From 3982ea57f067d2cefdfd435d2ddecea41e44b4dd Mon Sep 17 00:00:00 2001 From: Praful Makani Date: Fri, 16 Oct 2020 21:39:02 +0530 Subject: [PATCH 3/3] fix: address comment --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 6a5bd2392..a19c65e63 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 @@ -1274,7 +1274,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() { TableResult tableResult = job.getQueryResults(); return tableResult; } else { - schema = Schema.fromPb(results.getSchema()); + schema = results.getSchema() == null ? null : Schema.fromPb(results.getSchema()); if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) { numRows = 0L; } else if (results.getNumDmlAffectedRows() != null) {