Skip to content

Commit 51f31fb

Browse files
pongadpongad
authored
---
yaml --- r: 8921 b: refs/heads/lesv-patch-1 c: c34e949 h: refs/heads/master i: 8919: a67b155
1 parent 5d44eeb commit 51f31fb

3 files changed

Lines changed: 85 additions & 17 deletions

File tree

‎[refs]‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ refs/tags/v0.22.0: 18b298fe4bfe8ec2f20b0e0bf7ffdcce5cc3c5fe
6666
refs/heads/vam-google-patch-1: d0c8fee3a4074d0bf7360ce8c4f7f7223d0ee7b9
6767
refs/heads/vam-google-patch-CODEOWNERS: 2ac1616e25229e51d08a984708ef1918f91a35ee
6868
refs/heads/danoscarmike-patch-1: 7342a9916bce4ed00002c7202e2a16c5d46afaea
69-
refs/heads/lesv-patch-1: 3313bcbde7266205097020726b7c1b74ca5945ca
69+
refs/heads/lesv-patch-1: c34e9490ce6788ded930a85751d79f42713ffa5a
7070
refs/heads/ml-update-branch: 079dd6610017f5c51b9d1938c12d6d55b61513cf
7171
refs/heads/vkedia-patch-2: 7d8241388a9769a5c069334761b06c7012c878e7
7272
refs/heads/vkedia-patch-3: 4d128043acaa7db9160faf439d2ca6104e8a88cb

‎branches/lesv-patch-1/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,13 @@
2525
import com.google.cloud.RetryOption;
2626
import com.google.cloud.bigquery.BigQuery.JobOption;
2727
import com.google.cloud.bigquery.BigQuery.QueryResultsOption;
28+
import com.google.cloud.bigquery.BigQuery.TableDataListOption;
2829
import com.google.cloud.bigquery.JobConfiguration.Type;
2930
import java.io.IOException;
3031
import java.io.ObjectInputStream;
32+
import java.util.ArrayList;
33+
import java.util.Arrays;
34+
import java.util.List;
3135
import java.util.Objects;
3236
import java.util.concurrent.Callable;
3337
import java.util.concurrent.ExecutionException;
@@ -62,11 +66,9 @@ public class Job extends JobInfo {
6266
.setMaxRetryDelay(Duration.ofSeconds(3L))
6367
.build();
6468

65-
static final QueryResultsOption[] DEFAULT_QUERY_WAIT_OPTIONS =
66-
new QueryResultsOption[] {
67-
QueryResultsOption.pageSize(0L),
68-
QueryResultsOption.maxWaitTime(Duration.ofMinutes(1).toMillis())
69-
};
69+
static final QueryResultsOption[] DEFAULT_QUERY_WAIT_OPTIONS = {
70+
QueryResultsOption.pageSize(0L),
71+
};
7072

7173
private final BigQueryOptions options;
7274
private transient BigQuery bigquery;
@@ -272,13 +274,35 @@ public TableResult getQueryResults(QueryResultsOption... options)
272274
}
273275

274276
TableId table = ((QueryJobConfiguration) getConfiguration()).getDestinationTable();
275-
// TODO(pongad): merge options?
277+
278+
List<QueryResultsOption> waitOptions =
279+
new ArrayList<>(Arrays.asList(DEFAULT_QUERY_WAIT_OPTIONS));
280+
List<TableDataListOption> listOptions = new ArrayList<>();
281+
for (QueryResultsOption option : options) {
282+
switch (option.getRpcOption()) {
283+
case MAX_RESULTS:
284+
listOptions.add(TableDataListOption.pageSize((Long) option.getValue()));
285+
break;
286+
case PAGE_TOKEN:
287+
listOptions.add(TableDataListOption.pageToken((String) option.getValue()));
288+
break;
289+
case START_INDEX:
290+
listOptions.add(TableDataListOption.startIndex((Long) option.getValue()));
291+
break;
292+
case TIMEOUT:
293+
waitOptions.add(QueryResultsOption.maxWaitTime((Long) option.getValue()));
294+
break;
295+
}
296+
}
297+
276298
QueryResponse response =
277-
waitForQueryResults(DEFAULT_JOB_WAIT_SETTINGS, DEFAULT_QUERY_WAIT_OPTIONS);
299+
waitForQueryResults(
300+
DEFAULT_JOB_WAIT_SETTINGS, waitOptions.toArray(new QueryResultsOption[0]));
278301
if (response.getSchema() == null) {
279302
throw new JobException(getJobId(), response.getErrors());
280303
}
281-
return bigquery.listTableData(table, response.getSchema());
304+
return bigquery.listTableData(
305+
table, response.getSchema(), listOptions.toArray(new TableDataListOption[0]));
282306
}
283307

284308
private QueryResponse waitForQueryResults(

‎branches/lesv-patch-1/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java‎

Lines changed: 52 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1127,6 +1127,55 @@ public void testQueryRequestCompleted() throws InterruptedException {
11271127
.setTotalRows(BigInteger.valueOf(1L))
11281128
.setSchema(TABLE_SCHEMA.toPb());
11291129

1130+
EasyMock.expect(
1131+
bigqueryRpcMock.create(
1132+
JOB_INFO.toPb(), Collections.<BigQueryRpc.Option, Object>emptyMap()))
1133+
.andReturn(jobResponsePb);
1134+
EasyMock.expect(
1135+
bigqueryRpcMock.getQueryResults(
1136+
PROJECT, JOB, BigQueryImpl.optionMap(Job.DEFAULT_QUERY_WAIT_OPTIONS)))
1137+
.andReturn(responsePb);
1138+
EasyMock.expect(
1139+
bigqueryRpcMock.listTableData(
1140+
PROJECT, DATASET, TABLE, Collections.<BigQueryRpc.Option, Object>emptyMap()))
1141+
.andReturn(
1142+
new TableDataList()
1143+
.setPageToken("")
1144+
.setRows(ImmutableList.of(TABLE_ROW))
1145+
.setTotalRows(1L));
1146+
1147+
EasyMock.replay(bigqueryRpcMock);
1148+
bigquery = options.getService();
1149+
TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY, queryJob);
1150+
assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA);
1151+
assertThat(result.getTotalRows()).isEqualTo(1);
1152+
for (FieldValueList row : result.getValues()) {
1153+
assertThat(row.get(0).getBooleanValue()).isFalse();
1154+
assertThat(row.get(1).getLongValue()).isEqualTo(1);
1155+
}
1156+
}
1157+
1158+
@Test
1159+
public void testQueryRequestCompletedOptions() throws InterruptedException {
1160+
JobId queryJob = JobId.of(PROJECT, JOB);
1161+
com.google.api.services.bigquery.model.Job jobResponsePb =
1162+
new com.google.api.services.bigquery.model.Job()
1163+
.setConfiguration(QUERY_JOB_CONFIGURATION_FOR_QUERY.toPb())
1164+
.setJobReference(queryJob.toPb())
1165+
.setId(JOB)
1166+
.setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE"));
1167+
jobResponsePb.getConfiguration().getQuery().setDestinationTable(TABLE_ID.toPb());
1168+
GetQueryResultsResponse responsePb =
1169+
new GetQueryResultsResponse()
1170+
.setJobReference(queryJob.toPb())
1171+
.setRows(ImmutableList.of(TABLE_ROW))
1172+
.setJobComplete(true)
1173+
.setCacheHit(false)
1174+
.setPageToken(CURSOR)
1175+
.setTotalBytesProcessed(42L)
1176+
.setTotalRows(BigInteger.valueOf(1L))
1177+
.setSchema(TABLE_SCHEMA.toPb());
1178+
11301179
EasyMock.expect(
11311180
bigqueryRpcMock.create(JOB_INFO.toPb(), Collections.<BigQueryRpc.Option, Object>emptyMap()))
11321181
.andReturn(jobResponsePb);
@@ -1139,9 +1188,7 @@ public void testQueryRequestCompleted() throws InterruptedException {
11391188
bigqueryRpcMock.getQueryResults(
11401189
PROJECT, JOB, BigQueryImpl.optionMap(Job.DEFAULT_QUERY_WAIT_OPTIONS)))
11411190
.andReturn(responsePb);
1142-
EasyMock.expect(
1143-
bigqueryRpcMock.listTableData(
1144-
PROJECT, DATASET, TABLE, Collections.<BigQueryRpc.Option, Object>emptyMap()))
1191+
EasyMock.expect(bigqueryRpcMock.listTableData(PROJECT, DATASET, TABLE, optionMap))
11451192
.andReturn(
11461193
new TableDataList()
11471194
.setPageToken("")
@@ -1151,7 +1198,8 @@ public void testQueryRequestCompleted() throws InterruptedException {
11511198
EasyMock.replay(bigqueryRpcMock);
11521199
bigquery = options.getService();
11531200
// TODO(pongad): pagesize = 42
1154-
TableResult result = bigquery.query(QUERY_JOB_CONFIGURATION_FOR_QUERY, queryJob);
1201+
Job job = bigquery.create(JobInfo.of(queryJob, QUERY_JOB_CONFIGURATION_FOR_QUERY));
1202+
TableResult result = job.getQueryResults(pageSizeOption);
11551203
assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA);
11561204
assertThat(result.getTotalRows()).isEqualTo(1);
11571205
for (FieldValueList row : result.getValues()) {
@@ -1196,10 +1244,6 @@ public void testQueryRequestCompletedOnSecondAttempt() throws InterruptedExcepti
11961244
JOB_INFO.toPb(), Collections.<BigQueryRpc.Option, Object>emptyMap()))
11971245
.andReturn(jobResponsePb1);
11981246

1199-
QueryResultsOption pageSizeOption = QueryResultsOption.pageSize(42L);
1200-
Map<BigQueryRpc.Option, Object> optionMap = Maps.newEnumMap(BigQueryRpc.Option.class);
1201-
optionMap.put(pageSizeOption.getRpcOption(), pageSizeOption.getValue());
1202-
12031247
EasyMock.expect(
12041248
bigqueryRpcMock.getQueryResults(
12051249
PROJECT, JOB, BigQueryImpl.optionMap(Job.DEFAULT_QUERY_WAIT_OPTIONS)))

0 commit comments

Comments
 (0)