Skip to content

Commit 19e98b3

Browse files
committed
Fix flaky test
1 parent eed4122 commit 19e98b3

File tree

1 file changed

+7
-8
lines changed

1 file changed

+7
-8
lines changed

jdbc/src/main/java/tech/ydb/jdbc/context/StreamQueryResult.java

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,7 @@ public void onNextPart(QueryResultPart part) {
250250
private class LazyResultSet extends BaseYdbResultSet {
251251
private final BlockingQueue<ResultSetReader> readers = new ArrayBlockingQueue<>(5);
252252
private final AtomicLong rowsCount = new AtomicLong();
253-
private volatile boolean isCompleted = false;
253+
private final CompletableFuture<Void> isCompleted = new CompletableFuture<>();
254254
private volatile boolean isClosed = false;
255255

256256
private ResultSetReader current = null;
@@ -300,7 +300,7 @@ public boolean next() throws SQLException {
300300
return true;
301301
}
302302

303-
if (isCompleted && readers.isEmpty()) {
303+
if (isCompleted.isDone() && readers.isEmpty()) {
304304
current = null;
305305
if (rowsCount.get() > 0) {
306306
rowIndex = rowsCount.intValue() + 1;
@@ -317,10 +317,7 @@ public boolean next() throws SQLException {
317317
}
318318

319319
public void complete() {
320-
if (isCompleted) {
321-
return;
322-
}
323-
isCompleted = true;
320+
isCompleted.complete(null);
324321
}
325322

326323
@Override
@@ -351,7 +348,8 @@ public boolean isBeforeFirst() throws SQLException {
351348

352349
@Override
353350
public boolean isAfterLast() throws SQLException {
354-
return isCompleted && rowsCount.get() > 0 && rowIndex > rowsCount.intValue();
351+
isCompleted.join();
352+
return rowsCount.get() > 0 && rowIndex > rowsCount.intValue();
355353
}
356354

357355
@Override
@@ -361,7 +359,8 @@ public boolean isFirst() throws SQLException {
361359

362360
@Override
363361
public boolean isLast() throws SQLException {
364-
return isCompleted && rowsCount.get() > 0 && rowIndex == rowsCount.intValue();
362+
isCompleted.join();
363+
return rowsCount.get() > 0 && rowIndex == rowsCount.intValue();
365364
}
366365

367366
@Override

0 commit comments

Comments
 (0)