Close the client correctly in case of exceptions, improve integration test

akudiyar · Alexey Kuzin · commit b6703b97b940 · 2021-11-15T00:29:27.000+01:00
diff --git a/src/main/scala/io/tarantool/spark/connector/TarantoolSparkException.scala b/src/main/scala/io/tarantool/spark/connector/TarantoolSparkException.scala
@@ -7,10 +7,16 @@ import io.tarantool.driver.exceptions.TarantoolException
   *
   * @author Alexey Kuzin
   */
-case class TarantoolSparkException(message: String) extends TarantoolException(message) {}
+trait TarantoolSparkException extends TarantoolException {}
 
 object TarantoolSparkException {
 
-  def TarantoolSparkException(message: String): TarantoolSparkException =
-    new TarantoolSparkException(message)
+  def apply(message: String): TarantoolSparkException =
+    new TarantoolException(message) with TarantoolSparkException
+
+  def apply(exception: Throwable): TarantoolSparkException =
+    new TarantoolException(exception) with TarantoolSparkException
+
+  def apply(message: String, exception: Throwable): TarantoolSparkException =
+    new TarantoolException(message, exception) with TarantoolSparkException
 }
diff --git a/src/main/scala/io/tarantool/spark/connector/rdd/TarantoolRDD.scala b/src/main/scala/io/tarantool/spark/connector/rdd/TarantoolRDD.scala
@@ -9,7 +9,11 @@ import io.tarantool.spark.connector.config.{ReadConfig, TarantoolConfig}
 import io.tarantool.spark.connector.connection.TarantoolConnection
 import io.tarantool.spark.connector.partition.TarantoolPartition
 import io.tarantool.spark.connector.rdd.converter.{FunctionBasedTupleConverter, TupleConverter}
-import io.tarantool.spark.connector.util.ScalaToJavaHelper.{toJavaConsumer, toJavaFunction}
+import io.tarantool.spark.connector.util.ScalaToJavaHelper.{
+  toJavaBiFunction,
+  toJavaConsumer,
+  toJavaFunction
+}
 import io.tarantool.spark.connector.util.TarantoolCursorIterator
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.tarantool.MapFunctions.rowToTuple
@@ -108,11 +112,12 @@ class TarantoolRDD[R] private[spark] (
             }
             .toArray[CompletableFuture[_]]
 
-        CompletableFuture
-          .allOf(allFutures: _*)
-          .thenAccept(toJavaConsumer {
-            _: Void =>
-              try {
+        var savedException: Throwable = null
+        try {
+          CompletableFuture
+            .allOf(allFutures: _*)
+            .handle(toJavaBiFunction {
+              (_: Void, exception: Throwable) =>
                 if (failedRowsExceptions.nonEmpty) {
                   val sw: StringWriter = new StringWriter()
                   val pw: PrintWriter = new PrintWriter(sw)
@@ -121,19 +126,34 @@ class TarantoolRDD[R] private[spark] (
                       pw.append("\n\n")
                       exception.printStackTrace(pw)
                     }
-                    throw new TarantoolSparkException("Dataset write failed: " + sw.toString)
+                    savedException = TarantoolSparkException("Dataset write failed: " + sw.toString)
+                    logError(savedException.getMessage)
                   } finally {
                     pw.close()
                   }
                 } else {
-                  logInfo(s"Dataset write success, $rowCount rows written")
+                  if (Option(exception).isDefined) {
+                    savedException = exception
+                    logError("Dataset write failed: ", savedException)
+                  } else {
+                    logInfo(s"Dataset write success, $rowCount rows written")
+                  }
                 }
-              } finally {
-                client.close()
-              }
-          })
-          .get()
-          .asInstanceOf[Unit]
+                null
+            })
+            .join()
+        } catch {
+          case throwable: Throwable => savedException = throwable
+        } finally {
+          client.close()
+        }
+
+        if (Option(savedException).isDefined) {
+          savedException match {
+            case e: RuntimeException => throw e
+            case e: Any              => throw TarantoolSparkException(e)
+          }
+        }
       }
     )
 }
diff --git a/src/main/scala/io/tarantool/spark/connector/util/ScalaToJavaHelper.scala b/src/main/scala/io/tarantool/spark/connector/util/ScalaToJavaHelper.scala
@@ -1,8 +1,11 @@
 package io.tarantool.spark.connector.util
 
-import java.util.function.{Consumer => JConsumer}
-import java.util.function.{Function => JFunction}
-import java.util.function.{Supplier => JSupplier}
+import java.util.function.{
+  BiFunction => JBiFunction,
+  Consumer => JConsumer,
+  Function => JFunction,
+  Supplier => JSupplier
+}
 import scala.reflect.ClassTag
 
 /**
@@ -29,6 +32,14 @@ object ScalaToJavaHelper {
     override def apply(t: T1): R = f.apply(t)
   }
 
+  /**
+    * Converts a Scala {@link Function2} to a Java {@link java.util.function.BiFunction}
+    */
+  def toJavaBiFunction[T1, T2, R](f: (T1, T2) => R): JBiFunction[T1, T2, R] =
+    new JBiFunction[T1, T2, R] {
+      override def apply(t1: T1, t2: T2): R = f.apply(t1, t2)
+    }
+
   /**
     * Converts a Scala {@link Function1} to a Java {@link java.util.function.Function}
     */
diff --git a/src/test/resources/test_teardown.lua b/src/test/resources/test_teardown.lua
@@ -2,10 +2,11 @@ local crud = require('crud')
 
 local function truncate_space(space)
     local ok, err
-    ok, err = crud.truncate('test_space')
+    ok, err = crud.truncate(space)
     if (not ok) then
         error("Failed to truncate space '" .. space .. "', error: " .. tostring(err))
     end
 end
 
 truncate_space('test_space')
+truncate_space('orders')
diff --git a/src/test/scala/io/tarantool/spark/connector/integration/TarantoolSparkWriteClusterTest.scala b/src/test/scala/io/tarantool/spark/connector/integration/TarantoolSparkWriteClusterTest.scala
@@ -4,7 +4,8 @@ import io.tarantool.driver.api.conditions.Conditions
 import io.tarantool.driver.api.tuple.{DefaultTarantoolTupleFactory, TarantoolTuple}
 import io.tarantool.driver.mappers.DefaultMessagePackMapperFactory
 import io.tarantool.spark.connector.toSparkContextFunctions
-import org.apache.spark.sql.{Encoders, Row}
+import org.apache.spark.SparkException
+import org.apache.spark.sql.{Encoders, Row, SaveMode}
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
@@ -27,24 +28,25 @@ class TarantoolSparkWriteClusterTest
 
   private val orderSchema = Encoders.product[Order].schema
 
-  test("should write a list of objects to the space") {
+  test("should write a dataset of objects to the specified space with different modes") {
 
     val orders = Range(1, 10).map(i => Order(i))
 
-    val df = spark.createDataFrame(
+    var df = spark.createDataFrame(
       spark.sparkContext.parallelize(orders.map(order => order.asRow())),
       orderSchema
     )
 
+    // Insert, the partition is empty at first
     df.write
       .format("org.apache.spark.sql.tarantool")
-      .mode("overwrite")
+      .mode(SaveMode.Append)
       .option("tarantool.space", SPACE_NAME)
       .save()
 
-    val actual = spark.sparkContext.tarantoolSpace(SPACE_NAME, Conditions.any()).collect()
-
+    var actual = spark.sparkContext.tarantoolSpace(SPACE_NAME, Conditions.any()).collect()
     actual.length should be > 0
+
     val sorted = actual.sorted[TarantoolTuple](new Ordering[TarantoolTuple]() {
       override def compare(x: TarantoolTuple, y: TarantoolTuple): Int =
         x.getInteger("id").compareTo(y.getInteger("id"))
@@ -70,20 +72,139 @@ class TarantoolSparkWriteClusterTest
       )
       actualItem.getBoolean("cleared") should equal(expectedItem.getBoolean(6))
     }
+
+    // Replace
+    df = spark.createDataFrame(
+      spark.sparkContext.parallelize(
+        orders
+          .map(order => order.changeOrderType(order.orderType + "222"))
+          .map(order => order.asRow())
+      ),
+      orderSchema
+    )
+
+    df.write
+      .format("org.apache.spark.sql.tarantool")
+      .mode(SaveMode.Overwrite)
+      .option("tarantool.space", SPACE_NAME)
+      .save()
+
+    actual = spark.sparkContext.tarantoolSpace(SPACE_NAME, Conditions.any()).collect()
+    actual.length should be > 0
+
+    actual.foreach(item => item.getString("order_type") should endWith("222"))
+
+    // Second insert with the same IDs produces an exception
+    var thrownException: Throwable = the[SparkException] thrownBy {
+      df.write
+        .format("org.apache.spark.sql.tarantool")
+        .mode(SaveMode.Append)
+        .option("tarantool.space", SPACE_NAME)
+        .save()
+    }
+    thrownException.getMessage should include("Duplicate key exists")
+
+    // ErrorIfExists mode checks that partition is empty and provides an exception if it is not
+    thrownException = the[IllegalStateException] thrownBy {
+      df.write
+        .format("org.apache.spark.sql.tarantool")
+        .mode(SaveMode.ErrorIfExists)
+        .option("tarantool.space", SPACE_NAME)
+        .save()
+    }
+    thrownException.getMessage should include("already exists in Tarantool")
+
+    // Clear the data and check that they are written in ErrorIfExists mode
+    container.executeScript("test_teardown.lua").get()
+
+    df = spark.createDataFrame(
+      spark.sparkContext.parallelize(
+        orders
+          .map(order => order.changeOrderType(order.orderType + "333"))
+          .map(order => order.asRow())
+      ),
+      orderSchema
+    )
+
+    df.write
+      .format("org.apache.spark.sql.tarantool")
+      .mode(SaveMode.ErrorIfExists)
+      .option("tarantool.space", SPACE_NAME)
+      .save()
+
+    actual = spark.sparkContext.tarantoolSpace(SPACE_NAME, Conditions.any()).collect()
+    actual.length should be > 0
+
+    actual.foreach(item => item.getString("order_type") should endWith("333"))
+
+    // Check that new data are not written in Ignore mode if the partition is not empty
+    df = spark.createDataFrame(
+      spark.sparkContext.parallelize(
+        orders
+          .map(order => order.changeOrderType(order.orderType + "444"))
+          .map(order => order.asRow())
+      ),
+      orderSchema
+    )
+
+    df.write
+      .format("org.apache.spark.sql.tarantool")
+      .mode(SaveMode.Ignore)
+      .option("tarantool.space", SPACE_NAME)
+      .save()
+
+    actual = spark.sparkContext.tarantoolSpace(SPACE_NAME, Conditions.any()).collect()
+    actual.length should be > 0
+
+    actual.foreach(item => item.getString("order_type") should endWith("333"))
+
+    // Clear the data and check if they are written in Ignore mode
+    container.executeScript("test_teardown.lua").get()
+
+    df.write
+      .format("org.apache.spark.sql.tarantool")
+      .mode(SaveMode.Ignore)
+      .option("tarantool.space", SPACE_NAME)
+      .save()
+
+    actual = spark.sparkContext.tarantoolSpace(SPACE_NAME, Conditions.any()).collect()
+    actual.length should be > 0
+
+    actual.foreach(item => item.getString("order_type") should endWith("444"))
   }
 
+  test("should throw an exception if the space name is not specified") {
+    assertThrows[IllegalArgumentException] {
+      val orders = Range(1, 10).map(i => Order(i))
+
+      val df = spark.createDataFrame(
+        spark.sparkContext.parallelize(orders.map(order => order.asRow())),
+        orderSchema
+      )
+
+      df.write
+        .format("org.apache.spark.sql.tarantool")
+        .mode(SaveMode.Overwrite)
+        .save()
+    }
+  }
 }
 
 case class Order(
   id: Int,
   bucketId: Int,
-  orderType: String,
+  var orderType: String,
   orderValue: BigDecimal,
   orderItems: List[Int],
   options: Map[String, String],
   cleared: Boolean
 ) {
 
+  def changeOrderType(newOrderType: String): Order = {
+    orderType = newOrderType
+    this
+  }
+
   def asRow(): Row =
     Row(id, bucketId, orderType, orderValue, orderItems, options, cleared)
 
