refactor(batch): report errors as CompletionExceptions

Author: bevzzz

src/main/java/io/weaviate/client6/v1/api/collections/batch/BatchContext.java

@@ -126,7 +126,7 @@
  * @see StreamObserver
  * @see State
  * @see #shutdownNow
- * @see TaskHandle#result()
+ * @see TaskHandle#done()
  *
  * @author Dyma Solovei
  */

src/main/java/io/weaviate/client6/v1/api/collections/batch/BatchServerException.java

@@ -0,0 +1,10 @@
+package io.weaviate.client6.v1.api.collections.batch;
+
+import io.weaviate.client6.v1.api.WeaviateException;
+
+public class BatchServerException extends WeaviateException {
+
+  public BatchServerException(String message) {
+    super(message);
+  }
+}

src/main/java/io/weaviate/client6/v1/api/collections/batch/TaskHandle.java

@@ -33,17 +33,11 @@ public final class TaskHandle {
   /** Flag indicating the task has been ack'ed. */
   private final CompletableFuture<Void> acked = new CompletableFuture<>();
 
-  public final record Result(Optional<String> error) {
-    public Result {
-      requireNonNull(error, "error is null");
-    }
-  }
-
   /**
    * Task result completes when the client receives {@link Event.Results}
    * containing this handle's {@link #id}.
    */
-  private final CompletableFuture<Result> result = new CompletableFuture<>();
+  private final CompletableFuture<Void> done = new CompletableFuture<>();
 
   /** The number of times this task has been retried. */
   private final int retries;
@@ -84,7 +78,7 @@ private TaskHandle() {
 
   /**
    * Creates a new task containing the same data as this task and {@link #retries}
-   * counter incremented by 1. The {@link #acked} and {@link #result} futures
+   * counter incremented by 1. The {@link #acked} and {@link #done} futures
    * are not copied to the returned task.
    *
    * @return Task handle.
@@ -111,7 +105,7 @@ void setAcked() {
    * {@link #setError} afterwards will have no effect.
    */
   void setSuccess() {
-    setResult(new Result(Optional.empty()));
+    setResult(null);
   }
 
   /**
@@ -124,20 +118,24 @@ void setSuccess() {
    *              status for the task; prefer {@link #setSuccess} in that case.
    */
   void setError(String error) {
-    setResult(new Result(Optional.ofNullable(error)));
+    setResult(error);
   }
 
   /**
    * Set result for this task.
    *
    * @throws IllegalStateException if the task has not been ack'ed.
    */
-  private void setResult(Result result) {
+  private void setResult(String error) {
     if (!acked.isDone()) {
       // TODO(dyma): can this happen due to us?
       throw new IllegalStateException("Result can only be set for an ack'ed task");
     }
-    this.result.complete(result);
+    if (error == null) {
+      this.done.complete(null);
+    } else {
+      this.done.completeExceptionally(new BatchServerException(error));
+    }
   }
 
   /**
@@ -155,8 +153,8 @@ public CompletableFuture<Void> isAcked() {
    * @return A future which completes when the server
    *         has reported the result for this task.
    */
-  public CompletableFuture<Result> result() {
-    return result;
+  public CompletableFuture<Void> done() {
+    return done;
   }
 
   /**

src/test/java/io/weaviate/client6/v1/api/collections/batch/BatchContextTest.java

@@ -220,10 +220,9 @@ public void test_sendOneBatch() throws Exception {
     closeContext();
 
     // By the time context.close() returns all tasks MUST have results set.
-    Assertions.assertThat(tasks).extracting(TaskHandle::result)
+    Assertions.assertThat(tasks).extracting(TaskHandle::done)
         .allMatch(CompletableFuture::isDone)
-        .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
-        .allMatch(Optional::isEmpty);
+        .noneMatch(CompletableFuture::isCompletedExceptionally);
   }
 
   @Test
@@ -256,10 +255,9 @@ public void test_drainOnClose() throws Exception {
 
     closeContext();
 
-    Assertions.assertThat(tasks).extracting(TaskHandle::result)
+    Assertions.assertThat(tasks).extracting(TaskHandle::done)
         .allMatch(CompletableFuture::isDone)
-        .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
-        .allMatch(Optional::isEmpty);
+        .noneMatch(CompletableFuture::isCompletedExceptionally);
   }
 
   @Test
@@ -293,10 +291,9 @@ public void test_backoff() throws Exception {
 
     closeContext();
 
-    Assertions.assertThat(tasks).extracting(TaskHandle::result)
+    Assertions.assertThat(tasks).extracting(TaskHandle::done)
         .allMatch(CompletableFuture::isDone)
-        .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
-        .allMatch(Optional::isEmpty);
+        .noneMatch(CompletableFuture::isCompletedExceptionally);
   }
 
   @Test
@@ -328,10 +325,9 @@ public void test_backoffBacklog() throws Exception {
 
     closeContext();
 
-    Assertions.assertThat(tasks).extracting(TaskHandle::result)
+    Assertions.assertThat(tasks).extracting(TaskHandle::done)
         .allMatch(CompletableFuture::isDone)
-        .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
-        .allMatch(Optional::isEmpty);
+        .noneMatch(CompletableFuture::isCompletedExceptionally);
   }
 
   @Test
@@ -454,10 +450,9 @@ public void test_reconnect_DrainAfterStreamHangup() throws Exception {
     closeContext();
     backgroundAcks.get();
 
-    Assertions.assertThat(tasks).extracting(TaskHandle::result)
+    Assertions.assertThat(tasks).extracting(TaskHandle::done)
         .allMatch(CompletableFuture::isDone)
-        .extracting(CompletableFuture::get).extracting(TaskHandle.Result::error)
-        .allMatch(Optional::isEmpty);
+        .noneMatch(CompletableFuture::isCompletedExceptionally);
   }
 
   @Test

src/test/java/io/weaviate/client6/v1/api/collections/batch/TaskHandleTest.java

@@ -6,7 +6,6 @@
 
 import org.assertj.core.api.AbstractObjectAssert;
 import org.assertj.core.api.Assertions;
-import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.Test;
 
 import com.google.protobuf.GeneratedMessage;
@@ -51,10 +50,8 @@ public void test_newTaskHandle_WeaviateObject_success() {
     taskHandle.setSuccess();
 
     assertAcked(taskHandle, true);
-    assertHasResult(taskHandle, true)
-        .extracting(future -> future.getNow(null)).isNotNull()
-        .extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
-        .isEmpty();
+    Assertions.assertThat(taskHandle.done())
+        .isNotCompletedExceptionally();
   }
 
   @Test
@@ -76,10 +73,9 @@ public void test_newTaskHandle_WeaviateObject_error() {
     taskHandle.setError("Whaam!");
 
     assertAcked(taskHandle, true);
-    assertHasResult(taskHandle, true)
-        .extracting(future -> future.getNow(null)).isNotNull()
-        .extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
-        .get().isEqualTo("Whaam!");
+    Assertions.assertThat(taskHandle.done())
+        .isCompletedExceptionally()
+        .withFailMessage("Whaam!");
   }
 
   @Test
@@ -101,10 +97,8 @@ public void test_newTaskHandle_BatchReference_success() {
     taskHandle.setSuccess();
 
     assertAcked(taskHandle, true);
-    assertHasResult(taskHandle, true)
-        .extracting(future -> future.getNow(null)).isNotNull()
-        .extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
-        .isEmpty();
+    Assertions.assertThat(taskHandle.done())
+        .isNotCompletedExceptionally();
   }
 
   @Test
@@ -126,10 +120,9 @@ public void test_newTaskHandle_BatchReference_error() {
     taskHandle.setError("Whaam!");
 
     assertAcked(taskHandle, true);
-    assertHasResult(taskHandle, true)
-        .extracting(future -> future.getNow(null)).isNotNull()
-        .extracting(TaskHandle.Result::error, InstanceOfAssertFactories.optional(String.class))
-        .get().isEqualTo("Whaam!");
+    Assertions.assertThat(taskHandle.done())
+        .isCompletedExceptionally()
+        .withFailMessage("Whaam!");
   }
 
   @Test
@@ -160,10 +153,10 @@ private void assertAcked(TaskHandle taskHandle, boolean expect) {
         .returns(expect, CompletableFuture::isDone);
   }
 
-  private AbstractObjectAssert<?, CompletableFuture<TaskHandle.Result>> assertHasResult(TaskHandle taskHandle,
+  private AbstractObjectAssert<?, CompletableFuture<Void>> assertHasResult(TaskHandle taskHandle,
       boolean expect) {
     return Assertions.assertThat(taskHandle)
-        .extracting(TaskHandle::result)
+        .extracting(TaskHandle::done)
         .as("expect has result")
         .returns(expect, CompletableFuture::isDone);
   }