Fix flaky test

alex268 · Sep 25, 2024 · 19e98b3 · 19e98b3
1 parent eed4122
commit 19e98b3
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/StreamQueryResult.java b/jdbc/src/main/java/tech/ydb/jdbc/context/StreamQueryResult.java
@@ -250,7 +250,7 @@ public void onNextPart(QueryResultPart part) {
  private class LazyResultSet extends BaseYdbResultSet {
  private final BlockingQueue<ResultSetReader> readers = new ArrayBlockingQueue<>(5);
  private final AtomicLong rowsCount = new AtomicLong();
- private volatile boolean isCompleted = false;
+ private final CompletableFuture<Void> isCompleted = new CompletableFuture<>();
  private volatile boolean isClosed = false;
 
  private ResultSetReader current = null;
@@ -300,7 +300,7 @@ public boolean next() throws SQLException {
  return true;
  }
 
- if (isCompleted && readers.isEmpty()) {
+ if (isCompleted.isDone() && readers.isEmpty()) {
  current = null;
  if (rowsCount.get() > 0) {
  rowIndex = rowsCount.intValue() + 1;
@@ -317,10 +317,7 @@ public boolean next() throws SQLException {
  }
 
  public void complete() {
- if (isCompleted) {
- return;
- }
- isCompleted = true;
+ isCompleted.complete(null);
  }
 
  @Override
@@ -351,7 +348,8 @@ public boolean isBeforeFirst() throws SQLException {
 
  @Override
  public boolean isAfterLast() throws SQLException {
- return isCompleted && rowsCount.get() > 0 && rowIndex > rowsCount.intValue();
+ isCompleted.join();
+ return rowsCount.get() > 0 && rowIndex > rowsCount.intValue();
  }
 
  @Override
@@ -361,7 +359,8 @@ public boolean isFirst() throws SQLException {
 
  @Override
  public boolean isLast() throws SQLException {
- return isCompleted && rowsCount.get() > 0 && rowIndex == rowsCount.intValue();
+ isCompleted.join();
+ return rowsCount.get() > 0 && rowIndex == rowsCount.intValue();
  }
 
  @Override