package com.anryg.bigdata.clickhouse;

import com.clickhouse.jdbc.ClickHouseConnection;
import com.clickhouse.jdbc.ClickHouseDataSource;
import org.apache.spark.sql.ForeachWriter;
import org.apache.spark.sql.Row;

import java.sql.PreparedStatement;
import java.sql.SQLException;

/**
 * @DESC: 自定义structured streaming的外部sink，通过jdbc写数据到clickhouse中
 * @Auther: Anryg
 * @Date: 2023/7/3 11:29
 */
public class CKSink extends ForeachWriter<Row> {
    private static final String jdbcUrl = "jdbc:ch://192.168.211.107:8123,192.168.211.108:8123,192.168.211.109:8123/local_db"; //为了防止找不到本地表，把整个集群的配置都写上
    //private static final Properties properties = new Properties();
    private static ClickHouseDataSource ckDtaSource; 
    private static ClickHouseConnection connection;

    private static final String user = "default"; //用CK的默认用户
    private static final String pwd = "";     //默认用户没有设置密码
    private static final String tableName = "dns_logs_from_spark"; //写入的CK目标表


    /**
     * @DESC: 执行数据处理之前的准备工作，创建数据库连接，并确保单例，其中open会以partition为单位执行
     * */
    @Override
    public boolean open(long partitionId, long epochId){
        if (ckDtaSource == null || connection == null) {
            synchronized (CKSink.class){
                if (ckDtaSource == null || connection == null) {
                    try {
                        ckDtaSource = new ClickHouseDataSource(jdbcUrl);
                        connection = ckDtaSource.getConnection(user, pwd);
                    } catch (SQLException e) {
                        e.printStackTrace();
                        System.exit(-1); //捕获到异常后进程退出
                    }
                }
            }
        }

        if (connection == null) return false;
        else return true;
    }
    
    
    /**
     * @DESC: 当open函数返回为true之后，会针对partition中的每个ROW进行调用
     * */
    @Override
    public void process(Row value) {
        try {
            PreparedStatement preparedStatement = connection.prepareStatement("insert into " + tableName + " values(?,?,?,?,?,?,?,?,?)");
            preparedStatement.setString(1,value.getString(0));
            preparedStatement.setString(2,value.getString(1));
            preparedStatement.setString(3,value.getString(2));
            preparedStatement.setString(4,value.getString(3));
            preparedStatement.setString(5,value.getString(4));
            preparedStatement.setString(6,value.getString(5));
            preparedStatement.setString(7,value.getString(6));
            preparedStatement.setString(8,value.getString(7));
            preparedStatement.setString(9,value.getString(8));
            preparedStatement.addBatch();
            preparedStatement.executeBatch();
        } catch (SQLException e) {
            e.printStackTrace();
            System.exit(-1); //捕获到异常后进程退出
        }

    }

    /**
     * @DESC: 上两个函数执行完后，开始调用，一般用于关闭连接
     * */
    @Override
    public void close(Throwable errorOrNull) {
        //长连接，不关闭
    }
}

package com.anryg.bigdata.streaming.clickhouse


import com.anryg.bigdata.clickhouse.CKSink
import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.streaming.{OutputMode, Trigger}

/**
  * @DESC: 通过spark structured streaming消费kafka，并通过自定义的ForeachWriter写数据到clickhouse
  * @Auther: Anryg
  * @Date: 2023/7/3 10:18
  */
object Kafka2CK {

    def main(args: Array[String]): Unit = {
        val conf = new SparkConf().setAppName("Kafka2CK")/*.setMaster("local[*]")*/
        val spark = SparkSession.builder().config(conf).getOrCreate()

        val rawDF = spark.readStream //获取数据源
                .format("kafka") //确定数据源的来源格式
                .option("kafka.bootstrap.servers", "192.168.211.107:6667") //指定kafka集群的地址，理论上写一个broker就可以了
                .option("subscribe", ${topic})  //指定topic
                //.option("group.id","test9999") /**不再用该方式来绑定offset，而是每个程序有个唯一的id，该id跟checkpointLocation绑定，虽然group.id属性在运行中依然保留，但是不再跟offset绑定*/
                .option("failOnDataLoss",false)  //如果读取数据源时，发现数据突然缺失，比如被删，则是否马上抛出异常
                .option("fetchOffset.numRetries",3)  //获取消息的偏移量时，最多进行的重试次数
                .option("startingOffsets","latest")  //第一次消费时，读取kafka数据的位置
                .load()

        import spark.implicits._

        val ds = rawDF.selectExpr("CAST(value AS STRING)")  //将kafka中的数据的value转为为string，原始为binary类型
                .map(row => {
                    val line = row.getAs[String]("value") //获取row对象中的field，其实也只有一个field
                    val msgArray = line.split("\\|")  //指定分隔符进行字段切分
                    msgArray
                }).filter(_.length == 9)  //只留字段数为9的数据
                .map(array => (array(0),array(1),array(2),array(3),array(4),array(5),array(6),array(7),array(8))) //将其转化成为元组，为了方便下一步赋予schema
                .toDF("client_ip","domain","time","target_ip","rcode","query_type","authority_record","add_msg","dns_ip") //给裸数据添加字段名


        val query = ds.writeStream
                .outputMode(OutputMode.Append())  //指定数据的写入方式
                .foreach(new CKSink)   //通过自定义的方式写数据到clickhouse          
                .option("checkpointLocation","hdfs://192.168.211.106:8020/tmp/offset/Kafka2CK2") /**用来保存offset，用该目录来绑定对应的offset，如果该目录发生改变则程序运行的id会发生变化，类比group.id的变化*/
                .start()

        query.awaitTermination()
    }

}