From b42b9660c88af8f0792986bebde4150a94fc6055 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Tue, 20 Oct 2026 16:28:55 +0200 Subject: [PATCH] refactor(spanner): buffer chunks and add non-blocking isDataAvailable to GrpcValueIterator Refactor GrpcValueIterator and GrpcResultSet to support non-blocking data availability checks and bounded chunk buffering. This is Step 2 of the asynchronous stream refactoring plan, allowing reactive and async result set consumers to verify whether a full row is available in memory without blocking the calling thread on network I/O. Summary of changes: - Bounded chunk buffering: GrpcValueIterator maintains an internal queue of PartialResultSet chunks, buffering only as many chunks as needed to assemble at least one complete row (including handling chunked string and list merges across multiple chunks). - Non-blocking isDataAvailable(): Implemented isDataAvailable() on GrpcResultSet and GrpcValueIterator. Polling uses non-blocking tryLock() coordination so checking data availability never blocks behind concurrent worker threads reading from the stream. - Zero-allocation hot path: Optimizes row availability checks to avoid heap allocations on the steady-state path when rows are already available in the current chunk. Bug fixes and safety improvements included: - Fix getMetadata() idempotency: Calling getMetadata() multiple times or mid-stream now returns cached metadata immediately without erroneously advancing the stream or discarding chunks. - Prevent state resurrection on cancellation: If close() is invoked while a background stream.next() call is in flight, late-arriving chunks are safely discarded rather than repopulating state on a closed iterator. - Immediate surfacing of malformed metadata: Streams returning an initial chunk without valid metadata or row types immediately signal isDataAvailable() so the internal error is thrown rather than hanging the caller. - Clean up unused error field in GrpcResultSet. - Added comprehensive unit tests covering all chunked merge permutations, 0-column queries, error paths, and deterministic concurrent access. --- .../google/cloud/spanner/GrpcResultSet.java | 17 +- .../cloud/spanner/GrpcValueIterator.java | 319 +++++- .../cloud/spanner/GrpcResultSetTest.java | 75 ++ .../cloud/spanner/GrpcValueIteratorTest.java | 969 ++++++++++++++++++ 4 files changed, 1339 insertions(+), 41 deletions(-) create mode 100644 java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java index 299e68fe86a2..2f3daef47419 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java @@ -28,7 +28,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; @VisibleForTesting class GrpcResultSet extends AbstractResultSet> @@ -39,9 +39,8 @@ class GrpcResultSet extends AbstractResultSet> private ResultSetMetadata metadata; private GrpcStruct currRow; private List rowData; - private SpannerException error; private ResultSetStats statistics; - private boolean closed; + private volatile boolean closed; GrpcResultSet(CloseableIterator iterator, Listener listener) { this(iterator, listener, DecodeMode.DIRECT); @@ -75,9 +74,6 @@ protected GrpcStruct currRow() { @Override public boolean next() throws SpannerException { - if (error != null) { - throw asSpannerException(error); - } try { if (currRow == null) { metadata = iterator.getMetadata(); @@ -132,6 +128,15 @@ public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMess return iterator.initiateStreaming(streamMessageListener); } + @Override + @InternalApi + public boolean isDataAvailable() { + // If the result set is already closed, return true so callers (such as async stream + // consumers or event loops) do not block or wait indefinitely expecting future data, + // and instead proceed immediately to observe EOF or fast-fail. + return closed || iterator.isDataAvailable(); + } + @Override public void close() { synchronized (this) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java index 09b850c93f3c..593b574c2212 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java @@ -28,9 +28,12 @@ import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.ResultSetStats; import com.google.spanner.v1.TypeCode; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Deque; import java.util.List; -import javax.annotation.Nullable; +import java.util.concurrent.locks.ReentrantLock; +import org.jspecify.annotations.Nullable; /** Adapts a stream of {@code PartialResultSet} messages into a stream of {@code Value} messages. */ class GrpcValueIterator extends AbstractIterator { @@ -40,11 +43,15 @@ private enum StreamValue { } private final CloseableIterator stream; - private ResultSetMetadata metadata; - private Type type; - private PartialResultSet current; - private int pos; - private ResultSetStats statistics; + private final Deque bufferedChunks = new ArrayDeque<>(); + private final ReentrantLock streamPollLock = new ReentrantLock(); + private volatile boolean closed; + private volatile boolean streamFinished; + private volatile ResultSetMetadata metadata; + private volatile Type type; + private volatile PartialResultSet current; + private volatile int pos; + private volatile ResultSetStats statistics; private final Listener listener; GrpcValueIterator(CloseableIterator stream, Listener listener) { @@ -122,8 +129,7 @@ ResultSetMetadata getMetadata() throws SpannerException { * stream. Any attempt to call this method before the caller has finished consuming the results * will return null. */ - @Nullable - ResultSetStats getStats() { + @Nullable ResultSetStats getStats() { return statistics; } @@ -131,43 +137,184 @@ boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageList return stream.initiateStreaming(streamMessageListener); } + /** + * Returns the number of complete {@link com.google.protobuf.Value}s required to assemble one row + * without blocking. + * + *

Returns at least {@code 1} even when the result set has 0 columns (such as DML statements + * without a {@code THEN RETURN} clause or {@code QueryMode.PLAN} queries). For 0-column + * statements, {@link GrpcResultSet#next()} calls {@link #ensureReady(StreamValue)} with {@link + * StreamValue#RESULT}, which loops until the entire stream has finished so that the final {@link + * ResultSetStats} is collected. Requiring at least 1 value ensures that {@link + * #isDataAvailable()} waits for {@link #streamFinished} rather than returning {@code true} as + * soon as an initial metadata-only chunk arrives. + */ + private int requiredCompleteValues() { + if (type != null) { + return Math.max(1, type.getStructFields().size()); + } + PartialResultSet firstChunk = bufferedChunks.peekFirst(); + if (firstChunk != null && firstChunk.hasMetadata() && firstChunk.getMetadata().hasRowType()) { + return Math.max(1, firstChunk.getMetadata().getRowType().getFieldsCount()); + } + return 1; + } + + /** + * Returns {@code true} if the first chunk of the stream has arrived in {@link #bufferedChunks} + * but is missing the required row-type metadata. In that case, {@link #isDataAvailable()} should + * return {@code true} immediately so {@link #ensureReady(StreamValue)} can surface the metadata + * error without waiting for more chunks. + * + *

This should never happen in reality with a well-behaved Spanner backend or mock, which + * always delivers row-type metadata in the first message. This is a defensive check to guarantee + * that even in an ill-behaved situation, the client returns a reasonable error message instead of + * getting stuck waiting for metadata that will never arrive. + */ + private boolean hasInvalidFirstChunkMetadata() { + if (type != null) { + return false; + } + PartialResultSet firstChunk = bufferedChunks.peekFirst(); + return firstChunk != null + && (!firstChunk.hasMetadata() || !firstChunk.getMetadata().hasRowType()); + } + + /** + * Counts how many complete values are contributed by {@code chunk} starting at {@code + * startIndex}. + * + *

Whenever a chunk has {@code chunked_value = true} (with a mergeable trailing value), its + * last value is incomplete and will be merged with the first value of the following non-empty + * chunk. Thus each chunk with an incomplete trailing value reduces the number of complete values + * by 1 (either because its last value merges with the next chunk's first value, or, if it is the + * final buffered chunk, because its last value is still awaiting continuation). + */ + private int countCompleteValuesInChunk(@Nullable PartialResultSet chunk, int startIndex) { + if (chunk == null) { + return 0; + } + int count = chunk.getValuesCount(); + int remaining = count - startIndex; + if (remaining <= 0) { + return 0; + } + boolean hasIncompleteTrailingValue = + chunk.getChunkedValue() && isMergeable(chunk.getValues(count - 1).getKindCase()); + return hasIncompleteTrailingValue ? remaining - 1 : remaining; + } + + /** + * Returns {@code true} if metadata is known (either already parsed into {@link #type} or present + * on the first buffered chunk) and enough complete values are buffered to assemble a full row. + * + *

Short-circuits immediately without allocating an iterator if {@link #current} already + * contains enough complete values or if {@link #bufferedChunks} is empty. + */ + private boolean hasCompleteRowAvailable() { + if (type == null && bufferedChunks.isEmpty()) { + return false; + } + int requiredCompleteValues = requiredCompleteValues(); + int completeValues = countCompleteValuesInChunk(this.current, this.pos); + if (completeValues >= requiredCompleteValues) { + return true; + } + if (bufferedChunks.isEmpty()) { + return false; + } + for (PartialResultSet chunk : bufferedChunks) { + completeValues += countCompleteValuesInChunk(chunk, /* startIndex= */ 0); + if (completeValues >= requiredCompleteValues) { + return true; + } + } + return false; + } + + /** + * Non-blockingly pulls available chunks from {@link #stream} into {@link #bufferedChunks} until a + * complete row is available, EOF is reached, or no further chunks can be polled without blocking. + * + *

Acquires {@link #streamPollLock} via non-blocking {@link ReentrantLock#tryLock()} so that + * calls to {@link #isDataAvailable()} never block if another thread is currently polling or + * waiting on {@link #stream}. + */ + private void bufferAvailableChunks() { + if (!streamPollLock.tryLock()) { + return; + } + try { + while (!streamFinished + && !hasInvalidFirstChunkMetadata() + && !hasCompleteRowAvailable() + && stream.canPollWithoutErrorOrBlocking()) { + if (stream.hasNext()) { + bufferedChunks.addLast(stream.next()); + } else { + streamFinished = true; + } + } + } finally { + streamPollLock.unlock(); + } + } + + /** + * Returns {@code true} if the next call to {@link GrpcResultSet#next()} is guaranteed not to + * block on network I/O (because a complete row, EOF, or an error/closure is already available). + * + *

All operations in this method are guaranteed to be non-blocking: {@link + * #bufferAvailableChunks()} only polls while {@link + * CloseableIterator#canPollWithoutErrorOrBlocking()} is true (which returns immediately from + * memory without blocking or throwing), and subsequent checks inspect only in-memory buffered + * state. Therefore, it is safe to hold the synchronization monitor for this method without risk + * of blocking on network I/O. + */ + synchronized boolean isDataAvailable() { + bufferAvailableChunks(); + if (streamFinished || hasInvalidFirstChunkMetadata() || hasCompleteRowAvailable()) { + return true; + } + return stream.isDataAvailable(); + } + Type type() { checkState(type != null, "metadata has not been received"); return type; } + /** + * Ensures that the iterator is in a ready state satisfying {@code requiredValue}. + * + *

    + *
  • For {@link StreamValue#METADATA}: Ensures that metadata has been received and parsed into + * {@link #type} (and {@link #metadata}). If metadata is already initialized, returns {@code + * true} immediately without advancing chunks. Otherwise, advances to the first chunk and + * initializes metadata, returning {@code true} even if that chunk contains zero values. + *
  • For {@link StreamValue#RESULT}: Ensures that {@link #current} contains at least one + * unread value (i.e. {@link #hasUnreadValuesInCurrentChunk()} is true). If {@link #current} + * is null or exhausted, advances through chunks until a chunk containing values is reached, + * or returns {@code false} if the stream finishes (EOF). + *
+ * + * @param requiredValue whether metadata or a result value is required + * @return {@code true} if the required item is ready; {@code false} if EOF is reached before it + * could be satisfied + * @throws SpannerException if stream or metadata parsing errors occur + */ private boolean ensureReady(StreamValue requiredValue) throws SpannerException { - while (current == null || pos >= current.getValuesCount()) { - if (!stream.hasNext()) { + if (requiredValue == StreamValue.METADATA && type != null) { + return true; + } + while (!hasUnreadValuesInCurrentChunk()) { + if (!advanceToNextChunk()) { return false; } - current = stream.next(); - pos = 0; if (type == null) { - // This is the first message on the stream. - if (!current.hasMetadata() || !current.getMetadata().hasRowType()) { - throw newSpannerException(ErrorCode.INTERNAL, "Missing type metadata in first message"); - } - metadata = current.getMetadata(); - com.google.spanner.v1.Type typeProto = - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.STRUCT) - .setStructType(metadata.getRowType()) - .build(); - try { - type = Type.fromProto(typeProto); - } catch (IllegalArgumentException e) { - throw newSpannerException( - ErrorCode.INTERNAL, "Invalid type metadata: " + e.getMessage(), e); - } - } - // collect the precommit token from each PartialResultSet - if (current.hasPrecommitToken()) { - listener.onPrecommitToken(current.getPrecommitToken()); - } - if (current.hasStats()) { - statistics = current.getStats(); + initializeMetadata(current); } + updatePrecommitTokenAndStats(current); if (requiredValue == StreamValue.METADATA) { return true; } @@ -175,7 +322,109 @@ private boolean ensureReady(StreamValue requiredValue) throws SpannerException { return true; } + private boolean hasUnreadValuesInCurrentChunk() { + return current != null && pos < current.getValuesCount(); + } + + /** + * Advances the iterator to the next {@link PartialResultSet} chunk, first draining any in-memory + * chunks from {@link #bufferedChunks}, and falling back to pulling from {@link #stream}. + * + * @return {@code true} if a new chunk was successfully assigned to {@link #current}; {@code + * false} if the iterator is closed or the stream reached EOF + */ + private boolean advanceToNextChunk() { + PartialResultSet nextChunk; + + // Step 1: Fast-path. Attempt to retrieve an already-buffered chunk from memory. + synchronized (this) { + if (closed) { + return false; + } + nextChunk = bufferedChunks.pollFirst(); + if (nextChunk == null && streamFinished) { + // All buffered chunks have been consumed and the stream has reached EOF. + return false; + } + } + + // Step 2: Slow-path. No in-memory chunk was immediately available; pull from the underlying + // stream. + if (nextChunk == null) { + // Acquire streamPollLock outside 'synchronized (this)' so that blocking on network I/O + // does not block concurrent callers of isDataAvailable(). + streamPollLock.lock(); + try { + // Double-checked locking: Another thread (such as isDataAvailable()) might have buffered + // chunks or closed the iterator while we were waiting to acquire streamPollLock. + synchronized (this) { + if (closed) { + return false; + } + nextChunk = bufferedChunks.pollFirst(); + if (nextChunk == null && streamFinished) { + return false; + } + } + if (nextChunk == null) { + if (!stream.hasNext()) { + synchronized (this) { + streamFinished = true; + } + return false; + } + nextChunk = stream.next(); + } + } finally { + streamPollLock.unlock(); + } + } + + // Step 3: Assign the new chunk to current. + synchronized (this) { + // Check closed once more in case close() was called while stream.next() was in flight. + // If closed, discard the chunk to avoid resurrecting state on a closed iterator. + if (closed) { + return false; + } + pos = 0; + current = nextChunk; + } + return true; + } + + private void initializeMetadata(PartialResultSet firstChunk) { + if (!firstChunk.hasMetadata() || !firstChunk.getMetadata().hasRowType()) { + throw newSpannerException(ErrorCode.INTERNAL, "Missing type metadata in first message"); + } + metadata = firstChunk.getMetadata(); + com.google.spanner.v1.Type typeProto = + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.STRUCT) + .setStructType(metadata.getRowType()) + .build(); + try { + type = Type.fromProto(typeProto); + } catch (IllegalArgumentException e) { + throw newSpannerException(ErrorCode.INTERNAL, "Invalid type metadata: " + e.getMessage(), e); + } + } + + private void updatePrecommitTokenAndStats(PartialResultSet chunk) { + if (chunk.hasPrecommitToken()) { + listener.onPrecommitToken(chunk.getPrecommitToken()); + } + if (chunk.hasStats()) { + statistics = chunk.getStats(); + } + } + void close(@Nullable String message) { + synchronized (this) { + closed = true; + bufferedChunks.clear(); + streamFinished = true; + } stream.close(message); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java index 85f94ebc6d74..614ad37fa041 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java @@ -22,6 +22,9 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import com.google.api.gax.grpc.GrpcCallContext; import com.google.api.gax.rpc.ApiCallContext; @@ -1302,4 +1305,76 @@ public PartialResultSet next() { resultSet.close(); assertTrue(listenerReceivedWithBeginTransaction.get()); } + + @Test + public void testIsDataAvailable() { + assertFalse(resultSet.isDataAvailable()); + + consumer.onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(makeMetadata(Type.struct(Type.StructField.of("f", Type.string())))) + .addValues(Value.string("val1").toProto()) + .addValues(Value.string("val2").toProto()) + .build()); + + assertTrue(resultSet.isDataAvailable()); + assertTrue(resultSet.next()); + assertEquals("val1", resultSet.getString(0)); + + assertTrue(resultSet.isDataAvailable()); + assertTrue(resultSet.next()); + assertEquals("val2", resultSet.getString(0)); + + assertFalse(resultSet.isDataAvailable()); + consumer.onCompleted(); + assertTrue(resultSet.isDataAvailable()); + assertFalse(resultSet.next()); + } + + @Test + public void testIsDataAvailableWhenClosed() { + assertFalse(resultSet.isDataAvailable()); + resultSet.close(); + assertTrue(resultSet.isDataAvailable()); + assertThrows(SpannerException.class, resultSet::next); + } + + @Test + public void testIsDataAvailableWhenClosedMidStream() { + consumer.onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(makeMetadata(Type.struct(Type.StructField.of("f", Type.string())))) + .addValues(Value.string("val1").toProto()) + .build()); + + assertTrue(resultSet.isDataAvailable()); + assertTrue(resultSet.next()); + assertEquals("val1", resultSet.getString(0)); + + // Mid-stream with no further chunks buffered yet + assertFalse(resultSet.isDataAvailable()); + resultSet.close(); + + // After close, isDataAvailable returns true and next() returns false immediately + assertTrue(resultSet.isDataAvailable()); + assertFalse(resultSet.next()); + assertThrows(IllegalStateException.class, () -> resultSet.getString(0)); + } + + @Test + public void testInitiateStreamingDelegation() { + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator mockStream = + mock(AbstractResultSet.CloseableIterator.class); + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + when(mockStream.initiateStreaming(listener)).thenReturn(true); + + try (GrpcResultSet streamingResultSet = new GrpcResultSet(mockStream, new NoOpListener())) { + assertTrue(streamingResultSet.initiateStreaming(listener)); + verify(mockStream).initiateStreaming(listener); + } + + assertFalse(resultSet.initiateStreaming(listener)); + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java new file mode 100644 index 000000000000..7b0a3cd32c87 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java @@ -0,0 +1,969 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; +import com.google.cloud.spanner.AbstractResultSet.Listener; +import com.google.protobuf.ListValue; +import com.google.protobuf.NullValue; +import com.google.spanner.v1.MultiplexedSessionPrecommitToken; +import com.google.spanner.v1.PartialResultSet; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.ResultSetStats; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.Type; +import com.google.spanner.v1.TypeCode; +import java.util.ArrayList; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class GrpcValueIteratorTest { + @SuppressWarnings("unchecked") + private final CloseableIterator stream = mock(CloseableIterator.class); + + private final Listener listener = mock(Listener.class); + + private static StructType createSingleColumnRowType(String columnName, TypeCode typeCode) { + return StructType.newBuilder() + .addFields( + Field.newBuilder().setName(columnName).setType(Type.newBuilder().setCode(typeCode))) + .build(); + } + + private static StructType createArrayColumnRowType(String columnName, TypeCode elementTypeCode) { + return StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName(columnName) + .setType( + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(elementTypeCode)))) + .build(); + } + + private static StructType createMultiColumnRowType(String[] columnNames, TypeCode[] typeCodes) { + StructType.Builder builder = StructType.newBuilder(); + for (int index = 0; index < columnNames.length; index++) { + builder.addFields( + Field.newBuilder() + .setName(columnNames[index]) + .setType(Type.newBuilder().setCode(typeCodes[index]))); + } + return builder.build(); + } + + private static PartialResultSet createChunkWithMetadata( + StructType rowType, boolean chunkedValue, com.google.protobuf.Value... values) { + PartialResultSet.Builder builder = + PartialResultSet.newBuilder() + .setMetadata(ResultSetMetadata.newBuilder().setRowType(rowType)) + .setChunkedValue(chunkedValue); + for (com.google.protobuf.Value value : values) { + builder.addValues(value); + } + return builder.build(); + } + + private static PartialResultSet createChunkWithoutMetadata( + boolean chunkedValue, com.google.protobuf.Value... values) { + PartialResultSet.Builder builder = PartialResultSet.newBuilder().setChunkedValue(chunkedValue); + for (com.google.protobuf.Value value : values) { + builder.addValues(value); + } + return builder.build(); + } + + private static com.google.protobuf.Value stringValue(String value) { + return com.google.protobuf.Value.newBuilder().setStringValue(value).build(); + } + + private static com.google.protobuf.Value int64Value(long value) { + return com.google.protobuf.Value.newBuilder().setStringValue(String.valueOf(value)).build(); + } + + private static com.google.protobuf.Value listValue(String... elements) { + ListValue.Builder listBuilder = ListValue.newBuilder(); + for (String element : elements) { + listBuilder.addValues(com.google.protobuf.Value.newBuilder().setStringValue(element)); + } + return com.google.protobuf.Value.newBuilder().setListValue(listBuilder).build(); + } + + @Test + public void testIsDataAvailableWhenCurrentIsNull() { + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(false); + when(stream.isDataAvailable()).thenReturn(false); + assertFalse(iterator.isDataAvailable()); + + when(stream.isDataAvailable()).thenReturn(true); + assertTrue(iterator.isDataAvailable()); + } + + @Test + public void testIsDataAvailableWhenValuesRemainInCurrentChunk() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet chunk = + createChunkWithMetadata(rowType, false, stringValue("val1"), stringValue("val2")); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.isDataAvailable()).thenReturn(false); + + // Initial state before reading: current is null, stream has no data available + assertFalse(iterator.isDataAvailable()); + + // Advance iterator: reads first value "val1", current is set, pos is now 1 (1 < 2 values) + assertTrue(iterator.hasNext()); + assertEquals("val1", iterator.next().getStringValue()); + + // Even though stream.isDataAvailable() is false, iterator has 1 value remaining in current + assertFalse(stream.isDataAvailable()); + assertTrue(iterator.isDataAvailable()); + + // Read second value "val2", pos is now 2 (all 2 values consumed from current chunk) + assertTrue(iterator.hasNext()); + assertEquals("val2", iterator.next().getStringValue()); + + // Now all values in current chunk are consumed, so isDataAvailable matches stream + assertFalse(iterator.isDataAvailable()); + when(stream.isDataAvailable()).thenReturn(true); + assertTrue(iterator.isDataAvailable()); + } + + @Test + public void testIsDataAvailableMultiColumnRowAcrossChunks() { + StructType rowType = + createMultiColumnRowType( + new String[] {"col1", "col2", "col3"}, + new TypeCode[] {TypeCode.STRING, TypeCode.STRING, TypeCode.STRING}); + PartialResultSet firstChunk = + createChunkWithMetadata(rowType, false, stringValue("val1"), stringValue("val2")); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, stringValue("val3")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(firstChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // After buffering firstChunk, only 2 complete values are present for a 3-column row. + // Row is incomplete, so isDataAvailable must return false. + assertFalse(iterator.isDataAvailable()); + + // Now secondChunk becomes available to poll without blocking + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(secondChunk); + + // With secondChunk buffered, 3 complete values are present, forming a full row. + assertTrue(iterator.isDataAvailable()); + + assertTrue(iterator.hasNext()); + assertEquals("val1", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("val2", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("val3", iterator.next().getStringValue()); + } + + @Test + public void testIsDataAvailableWhenCurrentEndsWithChunkedValueAndStreamHasNoData() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet firstChunk = + createChunkWithMetadata( + rowType, true, stringValue("complete_val"), stringValue("chunked_part_1_")); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, stringValue("chunked_part_2")); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(false); + when(stream.hasNext()).thenReturn(true, true, false); + when(stream.next()).thenReturn(firstChunk, secondChunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.isDataAvailable()).thenReturn(false); + + // Read the first complete value; pos is now 1 (pointing to the last value in firstChunk, + // which has chunked_value = true). + assertTrue(iterator.hasNext()); + assertEquals("complete_val", iterator.next().getStringValue()); + + // Because the remaining value in firstChunk is chunked and stream.isDataAvailable() is false, + // advancing the iterator would block waiting for secondChunk, so isDataAvailable() must be + // false. + assertFalse(iterator.isDataAvailable()); + + // Once secondChunk is available on the stream, isDataAvailable() becomes true. + when(stream.isDataAvailable()).thenReturn(true); + assertTrue(iterator.isDataAvailable()); + assertTrue(iterator.hasNext()); + assertEquals("chunked_part_1_chunked_part_2", iterator.next().getStringValue()); + } + + @Test + public void testIsDataAvailableChunkedStringAcrossThreeChunks() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet firstChunk = createChunkWithMetadata(rowType, true, stringValue("part1_")); + PartialResultSet secondChunk = createChunkWithoutMetadata(true, stringValue("part2_")); + PartialResultSet thirdChunk = createChunkWithoutMetadata(false, stringValue("part3")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(firstChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // First chunk has chunked_value = true, so 0 complete values are available. + assertFalse(iterator.isDataAvailable()); + + // Second chunk arrives, but it is also chunked (part 2 of 3). + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(secondChunk); + assertFalse(iterator.isDataAvailable()); + + // Third chunk arrives with chunked_value = false, completing the string value. + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(thirdChunk); + assertTrue(iterator.isDataAvailable()); + + assertTrue(iterator.hasNext()); + assertEquals("part1_part2_part3", iterator.next().getStringValue()); + } + + @Test + public void testIsDataAvailableChunkedListAcrossChunks() { + StructType rowType = createArrayColumnRowType("col", TypeCode.STRING); + PartialResultSet firstChunk = + createChunkWithMetadata(rowType, true, listValue("a", "b_part1_")); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, listValue("part2_b", "c")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(firstChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // List is chunked, so 0 complete values are buffered. + assertFalse(iterator.isDataAvailable()); + + // Second chunk arrives with remaining elements. + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(secondChunk); + assertTrue(iterator.isDataAvailable()); + + assertTrue(iterator.hasNext()); + com.google.protobuf.Value mergedValue = iterator.next(); + assertEquals(3, mergedValue.getListValue().getValuesCount()); + assertEquals("a", mergedValue.getListValue().getValues(0).getStringValue()); + assertEquals("b_part1_part2_b", mergedValue.getListValue().getValues(1).getStringValue()); + assertEquals("c", mergedValue.getListValue().getValues(2).getStringValue()); + } + + @Test + public void testIsDataAvailableZeroColumnStatementWaitsForStreamFinished() { + // 0-column statement (e.g. DML or QueryMode.PLAN) + StructType emptyRowType = StructType.newBuilder().build(); + PartialResultSet firstChunk = + PartialResultSet.newBuilder() + .setMetadata(ResultSetMetadata.newBuilder().setRowType(emptyRowType)) + .build(); + ResultSetStats expectedStats = ResultSetStats.newBuilder().setRowCountExact(42L).build(); + PartialResultSet secondChunk = PartialResultSet.newBuilder().setStats(expectedStats).build(); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + // Chunk 1 arrives with metadata, but stream is not finished yet. + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(firstChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // For 0-column statements, requiredCompleteValues is Math.max(1, 0) = 1. + // Since completeValues is 0 and streamFinished is false, isDataAvailable must be false. + assertFalse(iterator.isDataAvailable()); + + // Chunk 2 arrives with stats and finishes the stream (EOF). + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(secondChunk); + + // Now streamFinished becomes true when EOF is hit during bufferAvailableChunks. + assertTrue(iterator.isDataAvailable()); + + // hasNext should be false, and stats must be collected. + assertFalse(iterator.hasNext()); + assertEquals(expectedStats, iterator.getStats()); + } + + @Test + public void testIsDataAvailableSurfacesMissingMetadataImmediately() { + // First chunk has no metadata at all. + PartialResultSet chunkWithoutMetadata = + PartialResultSet.newBuilder().addValues(stringValue("orphan_value")).build(); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunkWithoutMetadata); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + // isDataAvailable should return true immediately to allow ensureReady to fail fast. + assertTrue(iterator.isDataAvailable()); + + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("Missing type metadata in first message")); + } + + @Test + public void testIsDataAvailableSurfacesMissingRowTypeImmediately() { + // First chunk has metadata but no rowType. + PartialResultSet chunkWithoutRowType = + PartialResultSet.newBuilder() + .setMetadata(ResultSetMetadata.newBuilder().build()) + .addValues(stringValue("orphan_value")) + .build(); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunkWithoutRowType); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + // isDataAvailable should return true immediately. + assertTrue(iterator.isDataAvailable()); + + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("Missing type metadata in first message")); + } + + @Test + public void testInvalidTypeMetadataThrowsException() { + // Struct with invalid field type that fails Type.fromProto() + StructType invalidRowType = + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("invalid_col") + .setType(Type.newBuilder().setCode(TypeCode.ARRAY))) + .build(); + PartialResultSet invalidChunk = + PartialResultSet.newBuilder() + .setMetadata(ResultSetMetadata.newBuilder().setRowType(invalidRowType)) + .build(); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(invalidChunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("Invalid type metadata")); + } + + @Test + public void testIsDataAvailableWhenStreamHasError() { + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(false); + when(stream.isDataAvailable()).thenReturn(true); + when(stream.hasNext()) + .thenThrow( + SpannerExceptionFactory.newSpannerException(ErrorCode.UNAVAILABLE, "Transport error")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertTrue(iterator.isDataAvailable()); + + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("Transport error")); + } + + @Test + public void testBufferAvailableChunksPullsMultipleChunks() { + StructType rowType = + createMultiColumnRowType( + new String[] {"col1", "col2"}, new TypeCode[] {TypeCode.STRING, TypeCode.INT64}); + PartialResultSet firstChunk = createChunkWithMetadata(rowType, false, stringValue("first")); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, int64Value(100L)); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, true, false); + when(stream.hasNext()).thenReturn(true, true, false); + when(stream.next()).thenReturn(firstChunk, secondChunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertTrue(iterator.isDataAvailable()); + + assertTrue(iterator.hasNext()); + assertEquals("first", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("100", iterator.next().getStringValue()); + assertNotNull(iterator.type()); + assertEquals(2, iterator.type().getStructFields().size()); + } + + @Test + public void testCloseClearsBufferedChunksAndClosesStream() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet chunk = createChunkWithMetadata(rowType, false, stringValue("val1")); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertTrue(iterator.isDataAvailable()); + + iterator.close("Test close"); + verify(stream).close("Test close"); + + // After close, buffered chunks are cleared and streamFinished is true + assertFalse(iterator.hasNext()); + } + + @Test + public void testGetMetadataWhenMetadataReceived() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + ResultSetMetadata metadata = ResultSetMetadata.newBuilder().setRowType(rowType).build(); + PartialResultSet chunk = PartialResultSet.newBuilder().setMetadata(metadata).build(); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertEquals(metadata, iterator.getMetadata()); + assertNotNull(iterator.type()); + } + + @Test + public void testGetMetadataMultipleTimesDoesNotAdvanceStream() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + ResultSetMetadata metadata = ResultSetMetadata.newBuilder().setRowType(rowType).build(); + PartialResultSet firstChunk = PartialResultSet.newBuilder().setMetadata(metadata).build(); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, stringValue("val")); + + when(stream.hasNext()).thenReturn(true, true, false); + when(stream.next()).thenReturn(firstChunk, secondChunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + // First call to getMetadata advances to firstChunk and initializes metadata + assertEquals(metadata, iterator.getMetadata()); + verify(stream, times(1)).next(); + + // Subsequent calls to getMetadata return the cached metadata and must NOT advance to + // secondChunk + assertEquals(metadata, iterator.getMetadata()); + assertEquals(metadata, iterator.getMetadata()); + verify(stream, times(1)).next(); + + // When iterating values, secondChunk is finally fetched from the stream + assertTrue(iterator.hasNext()); + verify(stream, times(2)).next(); + assertEquals("val", iterator.next().getStringValue()); + assertFalse(iterator.hasNext()); + } + + @Test + public void testGetMetadataAfterConsumingChunkValuesDoesNotAdvanceStream() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + ResultSetMetadata metadata = ResultSetMetadata.newBuilder().setRowType(rowType).build(); + PartialResultSet firstChunk = + PartialResultSet.newBuilder().setMetadata(metadata).addValues(stringValue("val1")).build(); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, stringValue("val2")); + + when(stream.hasNext()).thenReturn(true, true, false); + when(stream.next()).thenReturn(firstChunk, secondChunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + // Consume val1 so pos == 1 (all values in firstChunk consumed; hasUnreadValuesInCurrentChunk() + // is false) + assertTrue(iterator.hasNext()); + assertEquals("val1", iterator.next().getStringValue()); + verify(stream, times(1)).next(); + + // Calling getMetadata when pos == valuesCount must not advance the stream to secondChunk + assertEquals(metadata, iterator.getMetadata()); + verify(stream, times(1)).next(); + + // Iterating then advances to secondChunk + assertTrue(iterator.hasNext()); + verify(stream, times(2)).next(); + assertEquals("val2", iterator.next().getStringValue()); + } + + @Test + public void testGetMetadataWhenStreamClosedWithoutMetadataThrowsException() { + when(stream.hasNext()).thenReturn(false); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + SpannerException spannerException = assertThrows(SpannerException.class, iterator::getMetadata); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("Stream closed without sending metadata")); + } + + @Test + public void testTypeBeforeMetadataThrowsIllegalStateException() { + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertThrows(IllegalStateException.class, iterator::type); + } + + @Test + public void testPrecommitTokenAndStatsCollection() { + MultiplexedSessionPrecommitToken precommitToken = + MultiplexedSessionPrecommitToken.newBuilder().setSeqNum(12345).build(); + ResultSetStats stats = ResultSetStats.newBuilder().setRowCountExact(10L).build(); + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setMetadata(ResultSetMetadata.newBuilder().setRowType(rowType)) + .setPrecommitToken(precommitToken) + .setStats(stats) + .addValues(stringValue("value")) + .build(); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertTrue(iterator.hasNext()); + assertEquals("value", iterator.next().getStringValue()); + + verify(listener).onPrecommitToken(precommitToken); + assertEquals(stats, iterator.getStats()); + } + + @Test + public void testInitiateStreamingDelegation() { + AsyncResultSet.StreamMessageListener streamMessageListener = + mock(AsyncResultSet.StreamMessageListener.class); + when(stream.initiateStreaming(streamMessageListener)).thenReturn(true); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertTrue(iterator.initiateStreaming(streamMessageListener)); + verify(stream).initiateStreaming(streamMessageListener); + } + + @Test + public void testIsWithBeginTransactionAndIsLastStatementDelegation() { + when(stream.isWithBeginTransaction()).thenReturn(true); + when(stream.isLastStatement()).thenReturn(false); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + assertTrue(iterator.isWithBeginTransaction()); + assertFalse(iterator.isLastStatement()); + } + + @Test + public void testUnexpectedChunkedPartialResultSetNonMergeable() { + StructType rowType = createSingleColumnRowType("col", TypeCode.BOOL); + com.google.protobuf.Value boolValue = + com.google.protobuf.Value.newBuilder().setBoolValue(true).build(); + PartialResultSet chunk = createChunkWithMetadata(rowType, true, boolValue); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("Unexpected chunked PartialResultSet.")); + } + + @Test + public void testStreamClosedInMiddleOfChunkedValue() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet chunk = createChunkWithMetadata(rowType, true, stringValue("partial")); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(chunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue( + spannerException.getMessage().contains("Stream closed in the middle of chunked value")); + } + + @Test + public void testUnexpectedTypeInMiddleOfChunkedValue() { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet firstChunk = + createChunkWithMetadata(rowType, true, stringValue("partial_string")); + com.google.protobuf.Value nullValue = + com.google.protobuf.Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build(); + PartialResultSet secondChunk = createChunkWithoutMetadata(false, nullValue); + + when(stream.hasNext()).thenReturn(true, true, false); + when(stream.next()).thenReturn(firstChunk, secondChunk); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + SpannerException spannerException = assertThrows(SpannerException.class, iterator::hasNext); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue( + spannerException.getMessage().contains("Unexpected type in middle of chunked value")); + } + + @Test + public void testIsDataAvailableMetadataOnlyChunkWaitsForDataChunk() { + StructType rowType = + createMultiColumnRowType( + new String[] {"col1", "col2"}, new TypeCode[] {TypeCode.STRING, TypeCode.STRING}); + PartialResultSet metadataChunk = + PartialResultSet.newBuilder() + .setMetadata(ResultSetMetadata.newBuilder().setRowType(rowType)) + .build(); + PartialResultSet dataChunk = + createChunkWithoutMetadata(false, stringValue("val1"), stringValue("val2")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(metadataChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // After buffering metadata-only chunk, completeValues is 0 (< 2 required for 2-column row). + // isDataAvailable must return false so async callers do not prematurely invoke next() and + // block. + assertFalse(iterator.isDataAvailable()); + + // When dataChunk becomes available to poll without blocking: + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(dataChunk); + + assertTrue(iterator.isDataAvailable()); + assertTrue(iterator.hasNext()); + assertEquals("val1", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("val2", iterator.next().getStringValue()); + } + + @Test + public void testIsDataAvailableMultipleRowsInSingleChunk() { + StructType rowType = + createMultiColumnRowType( + new String[] {"col1", "col2"}, new TypeCode[] {TypeCode.STRING, TypeCode.STRING}); + PartialResultSet singleChunk = + createChunkWithMetadata( + rowType, + false, + stringValue("r1c1"), + stringValue("r1c2"), + stringValue("r2c1"), + stringValue("r2c2")); + + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(singleChunk); + when(stream.isDataAvailable()).thenReturn(false); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + + // Initial check: singleChunk has 4 values (2 rows of 2 columns) + assertTrue(iterator.isDataAvailable()); + + // Read row 1 + assertTrue(iterator.hasNext()); + assertEquals("r1c1", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("r1c2", iterator.next().getStringValue()); + + // Even though stream.isDataAvailable() is false, row 2 is buffered in current + assertFalse(stream.isDataAvailable()); + assertTrue(iterator.isDataAvailable()); + + // Read row 2 + assertTrue(iterator.hasNext()); + assertEquals("r2c1", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("r2c2", iterator.next().getStringValue()); + + // All rows in the chunk are consumed + assertFalse(iterator.isDataAvailable()); + } + + @Test + public void testIsDataAvailableEmptyIntermediateChunk() { + StructType rowType = + createMultiColumnRowType( + new String[] {"col1", "col2"}, new TypeCode[] {TypeCode.STRING, TypeCode.STRING}); + PartialResultSet firstChunk = createChunkWithMetadata(rowType, false, stringValue("val1")); + MultiplexedSessionPrecommitToken token = + MultiplexedSessionPrecommitToken.newBuilder().setSeqNum(999).build(); + PartialResultSet emptyIntermediateChunk = + PartialResultSet.newBuilder().setPrecommitToken(token).build(); + PartialResultSet thirdChunk = createChunkWithoutMetadata(false, stringValue("val2")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(firstChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // Chunk 1 has 1 value (row requires 2) + assertFalse(iterator.isDataAvailable()); + + // Intermediate empty chunk arrives (0 values) + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(emptyIntermediateChunk); + assertFalse(iterator.isDataAvailable()); + + // Third chunk arrives with the remaining value + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(thirdChunk); + assertTrue(iterator.isDataAvailable()); + + assertTrue(iterator.hasNext()); + assertEquals("val1", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("val2", iterator.next().getStringValue()); + verify(listener).onPrecommitToken(token); + } + + @Test + public void testIsDataAvailableCompoundMultiColumnChunking() { + StructType rowType = + createMultiColumnRowType( + new String[] {"col1", "col2"}, new TypeCode[] {TypeCode.STRING, TypeCode.STRING}); + PartialResultSet firstChunk = + createChunkWithMetadata(rowType, true, stringValue("val1_part1_")); + PartialResultSet secondChunk = + createChunkWithoutMetadata(true, stringValue("part2"), stringValue("val2_part1_")); + PartialResultSet thirdChunk = createChunkWithoutMetadata(false, stringValue("part2")); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(firstChunk); + when(stream.isDataAvailable()).thenReturn(false); + + // Chunk 1 has 1 incomplete trailing value, so 0 complete values are buffered + assertFalse(iterator.isDataAvailable()); + + // Chunk 2 completes col1 ("val1_part1_part2"), but starts chunked col2 ("val2_part1_") + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(secondChunk); + + // Total complete values is 1 (only col1 is complete, row needs 2). Still false! + assertFalse(iterator.isDataAvailable()); + + // Chunk 3 arrives completing col2 ("val2_part1_part2") + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(true, false); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(thirdChunk); + assertTrue(iterator.isDataAvailable()); + + assertTrue(iterator.hasNext()); + assertEquals("val1_part1_part2", iterator.next().getStringValue()); + assertTrue(iterator.hasNext()); + assertEquals("val2_part1_part2", iterator.next().getStringValue()); + } + + @Test + public void testConcurrentIsDataAvailableAndAdvance() throws Exception { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + int rowCount = 100; + List chunks = new ArrayList<>(rowCount); + for (int index = 0; index < rowCount; index++) { + if (index == 0) { + chunks.add(createChunkWithMetadata(rowType, false, stringValue("val_" + index))); + } else { + chunks.add(createChunkWithoutMetadata(false, stringValue("val_" + index))); + } + } + + AtomicInteger streamIndex = new AtomicInteger(0); + when(stream.canPollWithoutErrorOrBlocking()) + .thenAnswer(invocation -> streamIndex.get() < chunks.size()); + when(stream.hasNext()).thenAnswer(invocation -> streamIndex.get() < chunks.size()); + when(stream.next()) + .thenAnswer( + invocation -> { + int currentIndex = streamIndex.getAndIncrement(); + if (currentIndex < chunks.size()) { + return chunks.get(currentIndex); + } + throw new NoSuchElementException("No more chunks available in stream"); + }); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + AtomicBoolean running = new AtomicBoolean(true); + AtomicReference errorRef = new AtomicReference<>(); + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + // Thread 1: repeatedly checks isDataAvailable + Future future1 = + executor.submit( + () -> { + while (running.get()) { + try { + iterator.isDataAvailable(); + Thread.yield(); + } catch (Throwable t) { + errorRef.set(t); + break; + } + } + }); + + // Thread 2: consumes rows + Future future2 = + executor.submit( + () -> { + try { + for (int index = 0; index < rowCount; index++) { + assertTrue(iterator.hasNext()); + assertEquals("val_" + index, iterator.next().getStringValue()); + } + } catch (Throwable t) { + errorRef.set(t); + } finally { + running.set(false); + } + }); + + future2.get(10, TimeUnit.SECONDS); + future1.get(10, TimeUnit.SECONDS); + if (errorRef.get() != null) { + throw new AssertionError("Concurrent execution failed", errorRef.get()); + } + } finally { + executor.shutdownNow(); + } + } + + @Test + public void testIsDataAvailableDoesNotBlockWhenAdvanceToNextChunkIsPollingStream() + throws Exception { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet chunk = createChunkWithMetadata(rowType, false, stringValue("val1")); + + CountDownLatch streamNextEnteredLatch = new CountDownLatch(1); + CountDownLatch allowStreamNextToCompleteLatch = new CountDownLatch(1); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()) + .thenAnswer( + invocation -> { + streamNextEnteredLatch.countDown(); + allowStreamNextToCompleteLatch.await(); + return chunk; + }); + when(stream.canPollWithoutErrorOrBlocking()).thenReturn(false); + when(stream.isDataAvailable()).thenReturn(false); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + // Thread 1 initiates hasNext(), which enters advanceToNextChunk() -> stream.next() + Future consumerFuture = executor.submit(iterator::hasNext); + + // Deterministically wait until Thread 1 is inside stream.next() holding streamPollLock + assertTrue(streamNextEnteredLatch.await(5, TimeUnit.SECONDS)); + + // At this point, Thread 1 holds streamPollLock and is blocked in stream.next(). + // isDataAvailable() must NOT block waiting for streamPollLock; tryLock() returns false, + // and isDataAvailable() returns immediately without blocking. + assertFalse(iterator.isDataAvailable()); + + // Release Thread 1 to finish stream.next() + allowStreamNextToCompleteLatch.countDown(); + + // Consumer finishes successfully + assertTrue(consumerFuture.get(5, TimeUnit.SECONDS)); + assertEquals("val1", iterator.next().getStringValue()); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void testCloseWhileStreamNextIsInFlightDoesNotResurrectState() throws Exception { + StructType rowType = createSingleColumnRowType("col", TypeCode.STRING); + PartialResultSet chunk = createChunkWithMetadata(rowType, false, stringValue("val1")); + + CountDownLatch streamNextEnteredLatch = new CountDownLatch(1); + CountDownLatch allowStreamNextToCompleteLatch = new CountDownLatch(1); + + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()) + .thenAnswer( + invocation -> { + streamNextEnteredLatch.countDown(); + allowStreamNextToCompleteLatch.await(); + return chunk; + }); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + // Thread 1 initiates hasNext(), which enters stream.next() + Future consumerFuture = executor.submit(iterator::hasNext); + + // Deterministically wait until Thread 1 is in flight inside stream.next() + assertTrue(streamNextEnteredLatch.await(5, TimeUnit.SECONDS)); + + // Close the iterator while stream.next() is in flight + iterator.close("Cancelled"); + + // Now allow stream.next() to finish returning chunk + allowStreamNextToCompleteLatch.countDown(); + + // Because closed is true, advanceToNextChunk() must discard chunk, avoid setting current, + // and return false. hasNext() must return false. + assertFalse(consumerFuture.get(5, TimeUnit.SECONDS)); + assertFalse(iterator.hasNext()); + } finally { + executor.shutdownNow(); + } + } +}