Flink Elasticsearch Sink

看flink源码时看到了Flink对Elasticsearch的连接器，因为平时往es写数据比较频繁，所以参考源码并且精简了一下，有需要的可以试试~

1. 源码

1.1. EsRowElasticsearchSinkFunction

package com.dengdz.Sink;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.table.api.TableException;
import org.apache.flink.util.Preconditions;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.common.xcontent.XContentType;

import java.io.Serializable;


/**
 * @author dengdz
 */
public class EsRowElasticsearchSinkFunction
        implements ElasticsearchSinkFunction<EsRow>, Serializable {


    private final XContentType contentType;
    private final Elasticsearch7RequestFactory requestFactory;

    public EsRowElasticsearchSinkFunction(
            XContentType contentType,
            Elasticsearch7RequestFactory requestFactory) {
        this.contentType = Preconditions.checkNotNull(contentType);
        this.requestFactory = Preconditions.checkNotNull(requestFactory);
    }

    @Override
    public void open() {
    }

    @Override
    public void process(EsRow element, RuntimeContext ctx, RequestIndexer indexer) {

        switch (element.getType()) {
            case INSERT:
                processInsert(element, indexer);
                break;
            case UPDATE:
                processUpdate(element, indexer);
                break;
            case DELETE:
                processDelete(element, indexer);
                break;
            default:
                throw new TableException("Unsupported message kind: " + element.getType());
        }
    }

    private void processInsert(EsRow row, RequestIndexer indexer) {

        // 如果文档ID不为空
        if (null != row && row.getId() != null) {
            final IndexRequest indexRequest =
                    requestFactory.createIndexRequest(
                            row.getIndex(),
                            row.getId(),
                            contentType,
                            row.getData().toString());
            indexer.add(indexRequest);
        }
    }

    private void processUpdate(EsRow row, RequestIndexer indexer) {

        // 如果文档ID不为空
        if (null != row && row.getId() != null) {
            final UpdateRequest updateRequest =
                    requestFactory.createUpdateRequest(
                            row.getIndex(),
                            row.getId(),
                            contentType,
                            row.getData().toString());
            indexer.add(updateRequest);
        }
    }

    private void processDelete(EsRow row, RequestIndexer indexer) {

        final DeleteRequest deleteRequest =
                requestFactory.createDeleteRequest(
                        row.getIndex(),
                        row.getId()
                );
        indexer.add(deleteRequest);
    }

}

1.2. EsRow

package com.dengdz.Sink;

import com.alibaba.fastjson.JSONObject;

/**
 * @author dengdz
 */
public class EsRow {

    private OperationType type;
    private String id;
    private String index;
    private JSONObject data;

    public EsRow(OperationType type, String id, String index, JSONObject data) {
        this.type = type;
        this.id = id;
        this.index = index;
        this.data = data;
    }

    public OperationType getType() {
        return type;
    }

    public void setType(OperationType type) {
        this.type = type;
    }

    public String getId() {
        return id;
    }

    public void setId(String id) {
        this.id = id;
    }

    public String getIndex() {
        return index;
    }

    public void setIndex(String index) {
        this.index = index;
    }

    public JSONObject getData() {
        return data;
    }

    public void setData(JSONObject data) {
        this.data = data;
    }

    public enum OperationType {
        INSERT, UPDATE, DELETE
    }

}

1.3. Elasticsearch7RequestFactory

package com.dengdz.Sink;

import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.common.xcontent.XContentType;

import java.io.Serializable;

/**
 * @author dengdz
 */
public class Elasticsearch7RequestFactory implements Serializable {

    public UpdateRequest createUpdateRequest(
            String index,
            String id,
            XContentType contentType,
            String document) {
        return new UpdateRequest(index, id)
                .doc(document, contentType)
                .upsert(document, contentType);
    }


    public IndexRequest createIndexRequest(
            String index,
            String id,
            XContentType contentType,
            String document) {
        return new IndexRequest(index)
                .id(id)
                .source(document, contentType);
    }


    public DeleteRequest createDeleteRequest(String index, String id) {
        return new DeleteRequest(index, id);
    }
}

1.4. AuthRestClientFactory

package com.dengdz.Sink;

import org.apache.flink.streaming.connectors.elasticsearch7.RestClientFactory;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.client.RestClientBuilder;

import javax.annotation.Nullable;

/**
 * @author dengdz
 */
public class AuthRestClientFactory implements RestClientFactory {

    private final String pathPrefix;
    private final String username;
    private final String password;
    private transient CredentialsProvider credentialsProvider;

    public AuthRestClientFactory(
            @Nullable String pathPrefix, String username, String password) {
        this.pathPrefix = pathPrefix;
        this.password = password;
        this.username = username;
    }

    @Override
    public void configureRestClientBuilder(RestClientBuilder restClientBuilder) {
        if (pathPrefix != null) {
            restClientBuilder.setPathPrefix(pathPrefix);
        }
        if (credentialsProvider == null) {
            credentialsProvider = new BasicCredentialsProvider();
            credentialsProvider.setCredentials(
                    AuthScope.ANY, new UsernamePasswordCredentials(username, password));
        }
        restClientBuilder.setHttpClientConfigCallback(
                httpAsyncClientBuilder ->
                        httpAsyncClientBuilder.setDefaultCredentialsProvider(
                                credentialsProvider));
    }

}