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(); + } + } +}