Update by Query logic을 실행하는 Custom OpenSearch Sink Connector 개발하기

package com.hiworks;

import com.hiworks.util.VersionUtil;
import org.apache.kafka.common.config.ConfigDef;
import org.apache.kafka.connect.connector.Task;
import org.apache.kafka.connect.sink.SinkConnector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class UpdateByQuerySinkConnector extends SinkConnector {
    private static final Logger log = LoggerFactory.getLogger(UpdateByQuerySinkConnector.class);
    private Map<String, String> configProps;

    @Override
    public void start(Map<String, String> props) {
        this.configProps = props;
    }

    @Override
    public Class<? extends Task> taskClass() {
        return UpdateByQuerySinkTask.class;
    }

    @Override
    public List<Map<String, String>> taskConfigs(int maxTasks) {
        final List<Map<String, String>> taskConfigs = new ArrayList<>(maxTasks);
        for (int i = 0; i < maxTasks; i++) {
            taskConfigs.add(configProps);
        }
        return taskConfigs;
    }

    @Override
    public void stop() {
        log.info("UpdateSinkConnector stopped");
    }

    @Override
    public ConfigDef config() {
        return new ConfigDef();
    }

    @Override
    public String version() {
        return VersionUtil.getVersion();
    }
}

package com.hiworks;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.hiworks.util.VersionUtil;
import com.hiworks.vo.UpdateField;
import org.apache.http.HttpHost;
import org.apache.kafka.connect.sink.SinkRecord;
import org.apache.kafka.connect.sink.SinkTask;
import org.opensearch.client.RestClient;
import org.opensearch.client.json.JsonData;
import org.opensearch.client.json.jackson.JacksonJsonpMapper;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch._types.FieldValue;
import org.opensearch.client.opensearch._types.Time;
import org.opensearch.client.opensearch._types.query_dsl.TermQuery;
import org.opensearch.client.opensearch.core.UpdateByQueryRequest;
import org.opensearch.client.opensearch.core.UpdateByQueryResponse;
import org.opensearch.client.transport.rest_client.RestClientTransport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Collection;
import java.util.Map;

public class UpdateByQuerySinkTask extends SinkTask {
    private static final Logger log = LoggerFactory.getLogger(UpdateByQuerySinkTask.class);

    private OpenSearchClient client;
    private RestClient restClient;
    private ObjectMapper objectMapper;

    @Override
    public void start(Map<String, String> props) {
        final String host = props.getOrDefault("opensearch.host", "localhost");
        final int port = Integer.parseInt(props.getOrDefault("opensearch.port", "9200"));
        HttpHost httpHost = new HttpHost(host, port);

        this.restClient = RestClient.builder(httpHost).build();
        this.client = new OpenSearchClient(
                new RestClientTransport(restClient, new JacksonJsonpMapper()));
        this.objectMapper = new ObjectMapper();
    }

    @Override
    public void put(Collection<SinkRecord> records) {
        if (records.isEmpty()) return;
        log.info("record 수신 성공 : {}개", records.size());

        for (SinkRecord record : records) {
            log.info("record = {}", record);
            try {
                update(record);
            } catch (Exception e) {
                log.error("record update 중 오류 발생", e);
            }
        }
    }

    private void update(SinkRecord record) throws IOException {
        final UpdateField updateField = UpdateField.fromRecord(record, objectMapper);
        log.info("updateField = {}", updateField);

        try {
            updateByQuery(updateField);
        } catch (Exception e) {
            log.error("record update by query 중 오류 발생 : updateField = {}", updateField, e);
        }
    }

    private void updateByQuery(UpdateField updateField) throws IOException {
        final UpdateByQueryRequest request = new UpdateByQueryRequest.Builder()
                .index(updateField.indexName())
                .query(q -> q
                        .term(TermQuery.of(t -> t
                                .field(updateField.termKey())
                                .value(FieldValue.of(updateField.termValue())))
                        )
                )
                .script(s -> s
                        .inline(i -> i
                                .source("ctx._source." + updateField.dataKey() + " = params." + updateField.dataKey())
                                .params(Map.of(updateField.dataKey(), JsonData.of(updateField.dataValue())))
                        )
                )
                .timeout(Time.of(t -> t.time("2m")))
                .build();
        log.info("update by query request : {}", request.toJsonString());

        final UpdateByQueryResponse response = client.updateByQuery(request);
        log.info("update by query response : {}", response.toJsonString());

        log.info("document update 완료 : {}개 {} ms 소요", response.updated(), response.took());
    }

    @Override
    public void stop() {
        try {
            if (restClient != null) {
                restClient.close();
                log.info("UpdateSinkTask 및 OpenSearch client 중지");
            }
        } catch (IOException e) {
            log.error("OpenSearch client 중지 중 오류 발생", e);
        }
    }

    @Override
    public String version() {
        return VersionUtil.getVersion();
    }
}

package com.hiworks.vo;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.kafka.connect.sink.SinkRecord;
import com.hiworks.enums.UpdateFieldName;

public record UpdateField(String indexName, String termKey, String termValue, String dataKey, String dataValue) {

    public static UpdateField fromRecord(SinkRecord record, ObjectMapper objectMapper) throws JsonProcessingException {
        final String recordStr = objectMapper.writeValueAsString(record.value());
        final JsonNode recordNode = objectMapper.readTree(recordStr);

        for (UpdateFieldName name : UpdateFieldName.values()) {
            if (!recordNode.has(name.code())) {
                throw new IllegalArgumentException(name.code() + " field 지정 필요");
            }
        }

        return new UpdateField(
                recordNode.get(UpdateFieldName.INDEX_NAME.code()).asText(),
                recordNode.get(UpdateFieldName.TERM_KEY.code()).asText(),
                recordNode.get(UpdateFieldName.TERM_VALUE.code()).asText(),
                recordNode.get(UpdateFieldName.DATA_KEY.code()).asText(),
                recordNode.get(UpdateFieldName.DATA_VALUE.code()).asText()
        );
    }
}

package com.hiworks.enums;

public enum UpdateFieldName {
    INDEX_NAME("index_name"),
    TERM_KEY("term_key"),
    TERM_VALUE("term_value"),
    DATA_KEY("data_key"),
    DATA_VALUE("data_value");

    private final String name;

    UpdateFieldName(String name) {
        this.name = name;
    }

    public String code() {
        return name;
    }
}

package com.hiworks.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.util.Properties;

public class VersionUtil {
    private static final Logger log = LoggerFactory.getLogger(VersionUtil.class);
    private static final String VERSION_FILE = "version.properties";
    private static final String VERSION_PROPERTY = "connector.version";
    private static final String DEFAULT_VERSION = "1.0";

    private static String version = null;

    public static String getVersion() {
        if (version != null) {
            return version;
        }

        try (InputStream input = VersionUtil.class.getClassLoader().getResourceAsStream(VERSION_FILE)) {
            if (input == null) {
                log.warn("{} file 읽기 오류 (기본값 {} 사용)", VERSION_FILE, DEFAULT_VERSION);
                version = DEFAULT_VERSION;
                return version;
            }

            Properties props = new Properties();
            props.load(input);

            version = props.getProperty(VERSION_PROPERTY, DEFAULT_VERSION);
            log.info("Connector Version : {}", version);
            return version;

        } catch (IOException e) {
            log.error("Connector Version 조회 오류 발생", e);
            version = DEFAULT_VERSION;
            return version;
        }
    }
}