JAVA-3166: Enhanced exception handling of C* driver executeAsync()

datastax · Aug 26, 2024 · 585472d · 585472d
1 parent d2d13b2
commit 585472d
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/connector/src/main/scala/com/datastax/spark/connector/writer/AsyncExecutor.scala b/connector/src/main/scala/com/datastax/spark/connector/writer/AsyncExecutor.scala
@@ -1,8 +1,7 @@
 package com.datastax.spark.connector.writer
 
-import java.util.concurrent.{CompletionStage, Semaphore}
+import java.util.concurrent.{CompletableFuture, CompletionStage, Semaphore}
 import java.util.function.BiConsumer
-
 import com.datastax.spark.connector.util.Logging
 
 import scala.jdk.CollectionConverters._
@@ -41,7 +40,12 @@ class AsyncExecutor[T, R](asyncAction: T => CompletionStage[R], maxConcurrentTas
     val executionTimestamp = System.nanoTime()
 
     def tryFuture(): Future[R] = {
-      val value = asyncAction(task)
+      var value : CompletionStage[R] = null
+      try {
+        value = asyncAction(task)
+      } catch {
+        case e: RuntimeException => value = CompletableFuture.failedStage(e);
+      }
 
       value.whenComplete(new BiConsumer[R, Throwable] {
         private def release() {

diff --git a/connector/src/test/scala/com/datastax/spark/connector/writer/AsyncExecutorTest.scala b/connector/src/test/scala/com/datastax/spark/connector/writer/AsyncExecutorTest.scala
@@ -1,5 +1,7 @@
 package com.datastax.spark.connector.writer
 
+import com.datastax.oss.driver.api.core.cql.{AsyncResultSet, SimpleStatement, Statement}
+
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.{Callable, CompletableFuture, CompletionStage}
 
@@ -55,4 +57,20 @@ class AsyncExecutorTest {
     totalFinishedExecutionsCounter.get() shouldBe taskCount
     asyncExecutor.getLatestException() shouldBe None
   }
+
+  @Test
+  def testGracefullyHandleCqlSessionExecuteExceptions() {
+    val executor = new AsyncExecutor[Statement[_], AsyncResultSet](
+      _ => {
+        // simulate exception returned by session.executeAsync() (not future)
+        throw new IllegalStateException("something bad happened")
+      }, 10, None, None
+    )
+    val stmt = SimpleStatement.newInstance("INSERT INTO table1 (key, value) VALUES (1, '100')");
+    val future = executor.executeAsync(stmt)
+    assertTrue(future.isCompleted)
+    val value = future.value.get
+    assertTrue(value.isInstanceOf[Failure[_]])
+    assertTrue(value.asInstanceOf[Failure[_]].exception.isInstanceOf[IllegalStateException])
+  }
 }