package com.dfire.platform.alchemy.connectors.elasticsearch; import org.apache.flink.api.common.functions.RuntimeContext; import org.apache.flink.api.common.serialization.SerializationSchema; import org.apache.flink.api.java.tuple.Tuple2; import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction; import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchUpsertTableSinkBase; import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer; import org.apache.flink.types.Row; 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.util.Arrays; import java.util.Objects; /** * @author dongbinglin */ public class Elasticsearch6SinkFunction implements ElasticsearchSinkFunction<Tuple2<Boolean, Row>> { private final String index; private final Integer fieldIndex; private final String docType; private final String keyDelimiter; private final String keyNullLiteral; private final SerializationSchema<Row> serializationSchema; private final XContentType contentType; private final ElasticsearchUpsertTableSinkBase.RequestFactory requestFactory; private final int[] keyFieldIndices; public Elasticsearch6SinkFunction(String index, Integer fieldIndex, String docType, String keyDelimiter, String keyNullLiteral, SerializationSchema<Row> serializationSchema, XContentType contentType, ElasticsearchUpsertTableSinkBase.RequestFactory requestFactory, int[] keyFieldIndices) { this.index = index; this.fieldIndex = fieldIndex; this.docType = Preconditions.checkNotNull(docType); this.keyDelimiter = Preconditions.checkNotNull(keyDelimiter); this.serializationSchema = Preconditions.checkNotNull(serializationSchema); this.contentType = Preconditions.checkNotNull(contentType); this.keyFieldIndices = Preconditions.checkNotNull(keyFieldIndices); this.requestFactory = Preconditions.checkNotNull(requestFactory); this.keyNullLiteral = Preconditions.checkNotNull(keyNullLiteral); } @Override public void process(Tuple2<Boolean, Row> element, RuntimeContext ctx, RequestIndexer indexer) { if (element.f0) { processUpsert(element.f1, indexer); } else { processDelete(element.f1, indexer); } } private void processUpsert(Row row, RequestIndexer indexer) { final byte[] document = serializationSchema.serialize(row); if (keyFieldIndices.length == 0) { final IndexRequest indexRequest = requestFactory.createIndexRequest( getIndex(row), docType, contentType, document); indexer.add(indexRequest); } else { final String key = createKey(row); final UpdateRequest updateRequest = requestFactory.createUpdateRequest( getIndex(row), docType, key, contentType, document); indexer.add(updateRequest); } } private void processDelete(Row row, RequestIndexer indexer) { final String key = createKey(row); final DeleteRequest deleteRequest = requestFactory.createDeleteRequest( getIndex(row), docType, key); indexer.add(deleteRequest); } private String createKey(Row row) { final StringBuilder builder = new StringBuilder(); for (int i = 0; i < keyFieldIndices.length; i++) { final int keyFieldIndex = keyFieldIndices[i]; if (i > 0) { builder.append(keyDelimiter); } final Object value = row.getField(keyFieldIndex); if (value == null) { builder.append(keyNullLiteral); } else { builder.append(value.toString()); } } return builder.toString(); } /** * 获取自定义索引名 * * @param row * @return */ private String getIndex(Row row) { if (fieldIndex == null) { return this.index; } return (String) row.getField(fieldIndex); } @Override public boolean equals(Object o) { if (this == o) { return true; } if (o == null || getClass() != o.getClass()) { return false; } Elasticsearch6SinkFunction that = (Elasticsearch6SinkFunction) o; return Objects.equals(index, that.index) && Objects.equals(fieldIndex, that.fieldIndex) && Objects.equals(docType, that.docType) && Objects.equals(keyDelimiter, that.keyDelimiter) && Objects.equals(keyNullLiteral, that.keyNullLiteral) && Objects.equals(serializationSchema, that.serializationSchema) && contentType == that.contentType && Objects.equals(requestFactory, that.requestFactory) && Arrays.equals(keyFieldIndices, that.keyFieldIndices); } @Override public int hashCode() { int result = Objects.hash( index, fieldIndex, docType, keyDelimiter, keyNullLiteral, serializationSchema, contentType, requestFactory); result = 31 * result + Arrays.hashCode(keyFieldIndices); return result; } }