自定义实现Structured Streaming的Sink（以MySQLSink为例）

背景

在做流计算时，想把每个批次的数据流式写入到MySQL但是发现并不支持jdbc的format。看了下，可以使用foreach并实现继承ForeachWriter的一个子类来实现这个目标。

//具体实现省略
class MySQLSink(userName:String,password:String,jdbcUrl:String) extends ForeachWriter[Row](){
   def open(partitionId: Long, version: Long): Boolean
   def process(value: T): Unit
   def close(errorOrNull: Throwable): Unit
}

    //foreach传入一个MySQL Sink,并传入相应参数
    val query = df.writeStream
      .foreach(new MySQLSink(userName,password,jdbcUrl))
    query.start()

但是，参考Spark自带的实现format(“kafka”)或者format(“console”)，发现可以更优雅的实现这种自定义的Sink。

！！！于是考虑自己实现一个Structured Streaming的Sink

步骤如下：

1. 创建一个自定义的SinkCreating Custom Sink — DemoSink
2. 创建一个对应的Sink Provider：Creating StreamSinkProvider — DemoSinkProvider
3. 使用ServiceLoader配置接口实现。Optional Sink Registration using META-INF/services
4. 使用！！！

1.创建Sink

jdbc的sink目前是只有batch才支持，不支持流式写入，因此改写sink支持jdbc方式的流式写入mysql。

package org.apache.spark.sql.execution.streaming.sinks.jdbc

import java.util.Properties

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.execution.streaming.Sink
import org.apache.spark.sql.streaming.OutputMode


class MySQLSink(options:Map[String,String],outputMode: OutputMode) extends Sink{

  override def addBatch(batchId: Long, data: DataFrame): Unit = {
    /**从option参数中获取需要的参数**/
    val userName = options.get("userName").orNull
    val password = options.get("password").orNull
    val table = options.get("table").orNull
    val jdbcUrl = options.get("jdbcUrl").orNull
    //data.show()
    //println(userName+"---"+password+"---"+jdbcUrl)
    val pro = new Properties
    pro.setProperty("user",userName)
    pro.setProperty("password",password)
    data.write.mode(outputMode.toString).jdbc(jdbcUrl,table,pro)
  }
}

2.创建Sink Provider

package org.apache.spark.sql.execution.streaming.sinks.jdbc

import org.apache.spark.sql.SQLContext
import org.apache.spark.sql.execution.streaming.Sink
import org.apache.spark.sql.sources.{DataSourceRegister, StreamSinkProvider}
import org.apache.spark.sql.streaming.OutputMode

class MySQLStreamSinkProvider extends StreamSinkProvider with DataSourceRegister{

  override def createSink(sqlContext: SQLContext,
                          parameters: Map[String, String],
                          partitionColumns: Seq[String],
                          outputMode: OutputMode): Sink = {
    new MySQLSink(parameters,outputMode);
  }

  //此名称可以在.format中使用。
  override def shortName(): String = "mysql"
}

3.创建文件

在/META-INFO /services 目录下 创建 名为：

org.apache.spark.sql.sources.DataSourceRegister的文件

并在里面加入

org.apache.spark.sql.execution.streaming.sinks.jdbc.MySQLStreamSinkProvider

对应Spark中的源码 DataSource[519-530]：

def lookupDataSource(provider: String): Class[_] = {
    val provider1 = backwardCompatibilityMap.getOrElse(provider, provider)
    val provider2 = s"$provider1.DefaultSource"
    val loader = Utils.getContextOrSparkClassLoader
    val serviceLoader = ServiceLoader.load(classOf[DataSourceRegister], loader)

    try {
      serviceLoader.asScala.filter(_.shortName().equalsIgnoreCase(provider1)).toList match {
        // the provider format did not match any given registered aliases
        case Nil =>
          try {
              /**根据这个provider1去load实现类**/
            Try(loader.loadClass(provider1)).orElse(Try(loader.loadClass(provider2))) match {

这样使用format的时候就可以找到这个实现了。

4.使用：

在本地mysql创建相应的库和表。并设置format为mysql。

即可流式写入mysql。

package com.dtwave.dipper.dubhe.streaming.example

import java.util.UUID

import org.apache.spark.sql.SparkSession

object KafkaExample extends App {

  val checkpointLocation = "/tmp/temporary-" + UUID.randomUUID.toString
  val spark = SparkSession
    .builder()
    .master("local[3]")
    .appName("kafka")
    .getOrCreate()

  spark.sparkContext.setLogLevel("ERROR")

//生成模拟数据
  val rateData = spark
    .readStream
    .format("rate")
    .option("rowsPerSecond",10)
    .load

  /**format指定为mysql**/
  rateData.writeStream.format("mysql")
    .option("checkpointLocation", "/tmp/demo-checkpoint")
    .option("username","root")
    .option("password","root")
    .option("table","mysql_sink")
    .option("jdbcUrl","jdbc:mysql://localhost:3306/test")
    .start

  spark.streams.awaitAnyTermination()
}

5.查看MySQL数据

6.总结

同理，例如hbase、elasticserach的sink，也可以这样的方式去定制化开发。