com.exasol.spark.ExasolWriter.scala Maven / Gradle / Ivy

package com.exasol.spark.writer

import java.sql.SQLException

import org.apache.spark.SparkContext
import org.apache.spark.TaskContext
import org.apache.spark.scheduler.SparkListener
import org.apache.spark.scheduler.SparkListenerApplicationEnd
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.StructType

import com.exasol.errorreporting.ExaError
import com.exasol.jdbc.EXAConnection
import com.exasol.spark.util.Converter
import com.exasol.spark.util.ExasolConnectionManager
import com.exasol.spark.util.Types

/**
 */
class ExasolWriter(
  @transient val sc: SparkContext,
  tableName: String,
  rddSchema: StructType,
  manager: ExasolConnectionManager
) extends Serializable {

  // scalastyle:off null
  @transient private var mainConnection: EXAConnection = null
  private var hosts: Seq[String] = null
  // scalastyle:on

  def closeMainResources(): Unit =
    if (mainConnection != null && !mainConnection.isClosed) {
      mainConnection.close()
    }

  def startParallel(): Int = {
    mainConnection = manager.writerMainConnection()

    if (mainConnection == null) {
      throw new RuntimeException(
        ExaError
          .messageBuilder("F-SEC-7")
          .message("Could not create main JDBC connection to Exasol cluster.")
          .mitigation("Please make sure that there network connection between Spark and Exasol clusters.")
          .toString()
      )
    }

    val cnt = manager.initParallel(mainConnection)

    // Close Exasol main connection when SparkContext finishes. This is a lifetime of a Spark
    // application.
    sc.addSparkListener(new SparkListener {
      override def onApplicationEnd(appEnd: SparkListenerApplicationEnd): Unit =
        closeMainResources()
    })

    // Populate hosts
    hosts = manager.subConnections(mainConnection)

    cnt
  }

  def insertStmt(): String = {
    val columns = rddSchema.fields.map(_.name).mkString(",")
    val placeholders = rddSchema.fields.map(_ => "?").mkString(",")
    s"INSERT INTO $tableName ($columns) VALUES ($placeholders)"
  }

  def insertPartition(iter: Iterator[Row]): Unit = {
    val partitionId = TaskContext.getPartitionId()
    val subConnectionUrl = hosts(partitionId)
    val subConn = manager.subConnection(subConnectionUrl)

    val stmt = subConn.prepareStatement(insertStmt())

    val setters = rddSchema.fields.map(f => Converter.makeSetter(f.dataType))
    val nullTypes = rddSchema.fields.map(f => Types.jdbcTypeFromSparkDataType(f.dataType))
    val fieldCnt = rddSchema.fields.length

    val batchSize = manager.config.batch_size

    try {
      var rowCnt = 0
      var totalCnt = 0

      while (iter.hasNext) {
        val row = iter.next()
        var i = 0
        while (i < fieldCnt) {
          if (row.isNullAt(i)) {
            stmt.setNull(i + 1, nullTypes(i))
          } else {
            setters(i).apply(stmt, row, i)
          }
          i = i + 1
        }
        stmt.addBatch()
        rowCnt += 1
        if (rowCnt % batchSize == 0) {
          val _ = stmt.executeBatch()
          totalCnt += rowCnt
          rowCnt = 0
        }
      }
      if (rowCnt > 0) {
        val _ = stmt.executeBatch()
        totalCnt += rowCnt
      }

      ()
    } catch {
      case ex: SQLException =>
        throw ex
    } finally {
      stmt.close()
      subConn.commit()
      subConn.close()
    }

  }

}