Skip to content

Commit 39902c3

Browse files
olavloiteolavloite
authored
fix: allow getMetadata() calls before calling next() (#3111)
* fix: allow getMetadata() calls before calling next() Calling getMetadata() on a MergedResultSet should be allowed without first calling next() in order to be consistent with other ResultSets that are returned by the Connection API. * fix: remove unnecessary partitions
1 parent 7e7c814 commit 39902c3

File tree

2 files changed

+122
-4
lines changed

2 files changed

+122
-4
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/MergedResultSet.java‎

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import static com.google.common.base.Preconditions.checkState;
2020

21+
import com.google.cloud.spanner.ErrorCode;
2122
import com.google.cloud.spanner.ForwardingStructReader;
2223
import com.google.cloud.spanner.ResultSet;
2324
import com.google.cloud.spanner.SpannerException;
@@ -30,6 +31,7 @@
3031
import com.google.spanner.v1.ResultSetStats;
3132
import java.util.ArrayList;
3233
import java.util.List;
34+
import java.util.concurrent.CountDownLatch;
3335
import java.util.concurrent.ExecutorService;
3436
import java.util.concurrent.Executors;
3537
import java.util.concurrent.LinkedBlockingDeque;
@@ -47,15 +49,18 @@ static class PartitionExecutor implements Runnable {
4749
private final Connection connection;
4850
private final String partitionId;
4951
private final LinkedBlockingDeque<PartitionExecutorResult> queue;
52+
private final CountDownLatch metadataAvailableLatch;
5053
private final AtomicBoolean shouldStop = new AtomicBoolean();
5154

5255
PartitionExecutor(
5356
Connection connection,
5457
String partitionId,
55-
LinkedBlockingDeque<PartitionExecutorResult> queue) {
58+
LinkedBlockingDeque<PartitionExecutorResult> queue,
59+
CountDownLatch metadataAvailableLatch) {
5660
this.connection = Preconditions.checkNotNull(connection);
5761
this.partitionId = Preconditions.checkNotNull(partitionId);
5862
this.queue = queue;
63+
this.metadataAvailableLatch = Preconditions.checkNotNull(metadataAvailableLatch);
5964
}
6065

6166
@Override
@@ -68,6 +73,7 @@ public void run() {
6873
queue.put(
6974
PartitionExecutorResult.dataAndMetadata(
7075
row, resultSet.getType(), resultSet.getMetadata()));
76+
metadataAvailableLatch.countDown();
7177
first = false;
7278
} else {
7379
queue.put(PartitionExecutorResult.data(row));
@@ -82,9 +88,11 @@ public void run() {
8288
queue.put(
8389
PartitionExecutorResult.typeAndMetadata(
8490
resultSet.getType(), resultSet.getMetadata()));
91+
metadataAvailableLatch.countDown();
8592
}
8693
} catch (Throwable exception) {
8794
putWithoutInterruptPropagation(PartitionExecutorResult.exception(exception));
95+
metadataAvailableLatch.countDown();
8896
} finally {
8997
// Emit a special 'finished' result to ensure that the row producer is not blocked on a
9098
// queue that never receives any more results. This ensures that we can safely block on
@@ -215,6 +223,7 @@ private static class RowProducerImpl implements RowProducer {
215223
private final AtomicInteger finishedCounter;
216224
private final LinkedBlockingDeque<PartitionExecutorResult> queue;
217225
private ResultSetMetadata metadata;
226+
private final CountDownLatch metadataAvailableLatch = new CountDownLatch(1);
218227
private Type type;
219228
private Struct currentRow;
220229
private Throwable exception;
@@ -243,7 +252,7 @@ private static class RowProducerImpl implements RowProducer {
243252
this.finishedCounter = new AtomicInteger(partitions.size());
244253
for (String partition : partitions) {
245254
PartitionExecutor partitionExecutor =
246-
new PartitionExecutor(connection, partition, this.queue);
255+
new PartitionExecutor(connection, partition, this.queue, this.metadataAvailableLatch);
247256
this.partitionExecutors.add(partitionExecutor);
248257
this.executor.submit(partitionExecutor);
249258
}
@@ -310,8 +319,27 @@ public Struct get() {
310319
return currentRow;
311320
}
312321

322+
private PartitionExecutorResult getFirstResult() {
323+
try {
324+
metadataAvailableLatch.await();
325+
} catch (InterruptedException interruptedException) {
326+
throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
327+
}
328+
PartitionExecutorResult result = queue.peek();
329+
if (result == null) {
330+
throw SpannerExceptionFactory.newSpannerException(
331+
ErrorCode.FAILED_PRECONDITION, "Thread-unsafe access to ResultSet");
332+
}
333+
if (result.exception != null) {
334+
throw SpannerExceptionFactory.asSpannerException(result.exception);
335+
}
336+
return result;
337+
}
338+
313339
public ResultSetMetadata getMetadata() {
314-
checkState(metadata != null, "next() call required");
340+
if (metadata == null) {
341+
return getFirstResult().metadata;
342+
}
315343
return metadata;
316344
}
317345

@@ -326,7 +354,9 @@ public int getParallelism() {
326354
}
327355

328356
public Type getType() {
329-
checkState(type != null, "next() call required");
357+
if (type == null) {
358+
return getFirstResult().type;
359+
}
330360
return type;
331361
}
332362
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -417,6 +417,94 @@ public void testRunEmptyPartitionedQuery() {
417417
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
418418
}
419419

420+
@Test
421+
public void testGetMetadataWithoutNextCall() {
422+
int generatedRowCount = 1;
423+
RandomResultSetGenerator generator = new RandomResultSetGenerator(generatedRowCount);
424+
Statement statement =
425+
Statement.newBuilder("select * from random_table where active=@active")
426+
.bind("active")
427+
.to(true)
428+
.build();
429+
mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
430+
431+
int maxPartitions = 1;
432+
try (Connection connection = createConnection()) {
433+
connection.setAutocommit(true);
434+
try (PartitionedQueryResultSet resultSet =
435+
connection.runPartitionedQuery(
436+
statement, PartitionOptions.newBuilder().setMaxPartitions(maxPartitions).build())) {
437+
assertNotNull(resultSet.getMetadata());
438+
assertEquals(24, resultSet.getMetadata().getRowType().getFieldsCount());
439+
assertNotNull(resultSet.getType());
440+
assertEquals(24, resultSet.getType().getStructFields().size());
441+
442+
assertTrue(resultSet.next());
443+
assertNotNull(resultSet.getMetadata());
444+
assertEquals(24, resultSet.getMetadata().getRowType().getFieldsCount());
445+
assertNotNull(resultSet.getType());
446+
assertEquals(24, resultSet.getType().getStructFields().size());
447+
448+
assertFalse(resultSet.next());
449+
}
450+
}
451+
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
452+
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
453+
}
454+
455+
@Test
456+
public void testGetMetadataWithoutNextCallOnEmptyResultSet() {
457+
int generatedRowCount = 0;
458+
RandomResultSetGenerator generator = new RandomResultSetGenerator(generatedRowCount);
459+
Statement statement =
460+
Statement.newBuilder("select * from random_table where active=@active")
461+
.bind("active")
462+
.to(true)
463+
.build();
464+
mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate()));
465+
466+
int maxPartitions = 1;
467+
try (Connection connection = createConnection()) {
468+
connection.setAutocommit(true);
469+
try (PartitionedQueryResultSet resultSet =
470+
connection.runPartitionedQuery(
471+
statement, PartitionOptions.newBuilder().setMaxPartitions(maxPartitions).build())) {
472+
assertNotNull(resultSet.getMetadata());
473+
assertEquals(24, resultSet.getMetadata().getRowType().getFieldsCount());
474+
assertNotNull(resultSet.getType());
475+
assertEquals(24, resultSet.getType().getStructFields().size());
476+
477+
assertFalse(resultSet.next());
478+
}
479+
}
480+
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
481+
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
482+
}
483+
484+
@Test
485+
public void testGetMetadataWithoutNextCallOnResultSetWithError() {
486+
Statement statement =
487+
Statement.newBuilder("select * from random_table where active=@active")
488+
.bind("active")
489+
.to(true)
490+
.build();
491+
mockSpanner.putStatementResult(
492+
StatementResult.exception(statement, Status.NOT_FOUND.asRuntimeException()));
493+
494+
int maxPartitions = 1;
495+
try (Connection connection = createConnection()) {
496+
connection.setAutocommit(true);
497+
try (PartitionedQueryResultSet resultSet =
498+
connection.runPartitionedQuery(
499+
statement, PartitionOptions.newBuilder().setMaxPartitions(maxPartitions).build())) {
500+
assertThrows(SpannerException.class, resultSet::getMetadata);
501+
assertThrows(SpannerException.class, resultSet::getType);
502+
}
503+
}
504+
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
505+
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
506+
}
507+
420508
@Test
421509
public void testRunPartitionedQueryUsingSql() {
422510
int generatedRowCount = 20;

0 commit comments

Comments
 (0)