io.polyglotted.elastic.index.Indexer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of elastic-common Show documentation
Elastic Common Utils
There is a newer version: 6.8.5
package io.polyglotted.elastic.index;

import io.polyglotted.common.model.AuthHeader;
import io.polyglotted.common.model.Pair;
import io.polyglotted.elastic.client.ElasticClient;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.DocWriteRequest;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.lucene.uid.Versions;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;

import java.util.function.BiConsumer;

import static io.polyglotted.common.model.Pair.pair;
import static io.polyglotted.common.util.MapBuilder.immutableMap;
import static io.polyglotted.common.util.MapBuilder.simpleMap;
import static io.polyglotted.common.util.MapRetriever.longStrVal;
import static io.polyglotted.common.util.NullUtil.nonNull;
import static io.polyglotted.elastic.common.MetaFields.SEQUENCE_FIELD;
import static io.polyglotted.elastic.common.MetaFields.TIMESTAMP_FIELD;
import static io.polyglotted.elastic.index.Validator.STRICT;
import static org.elasticsearch.action.DocWriteRequest.OpType.CREATE;
import static org.elasticsearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
import static org.elasticsearch.rest.RestStatus.CREATED;

@Slf4j @SuppressWarnings({"unused", "WeakerAccess", "UnusedReturnValue"})
@RequiredArgsConstructor
public final class Indexer {
    private final ElasticClient client;

    public void lockTheIndexOrFail(String repo, String keyString) { lockTheIndexOrFail(repo, keyString, false); }

    public void lockTheIndexOrFail(String repo, String keyString, boolean refresh) {
        IndexResponse response = client.index(new IndexRequest(repo, "_doc", keyString)
            .opType(CREATE).source(immutableMap(TIMESTAMP_FIELD, 1)));
        if (response.status() != CREATED) { throw new IndexerException("response failed while locking the keyString " + keyString); }
        if (refresh) { client.forceRefresh(repo); }
    }

    public boolean checkLock(String repo, String key) { return client.exists(new GetRequest(repo, "_doc", key)); }

    public void unlockIndex(String repo, String key) { client.delete(new DeleteRequest(repo, "_doc", key)); client.forceRefresh(repo); }

    public long generateSequence(String repo, String key) { return generateSequence(repo, key, 1)._b; }

    public Pair generateSequence(String repo, String key, int count) {
        GetResponse get = client.get(new GetRequest(repo, "_doc", key));
        long result = get.isExists() ? longStrVal(get.getSourceAsMap(), SEQUENCE_FIELD) + count : count;
        client.index(new IndexRequest(repo, "_doc", key).source(simpleMap(SEQUENCE_FIELD, count))
            .create(!get.isExists()).version(get.isExists() ? get.getVersion() : Versions.MATCH_ANY));
        return pair(result - count + 1, result);
    }

    public boolean bulkSave(AuthHeader auth, BulkRecord bulkRecord, Validator validator) {
        try {
            BulkRequest bulkRequest = validator.validateAll(client, auth, bulkRecord, new BulkRequest().setRefreshPolicy(IMMEDIATE));
            if (bulkRequest.numberOfActions() <= 0) { return bulkRecord.hasNoFailure(); }

            BulkResponse responses = client.bulk(auth, bulkRequest);
            return checkResponse(responses, bulkRecord.ignoreErrors, bulkRecord::actionWith);
        } catch (RuntimeException ex) { throw logError(ex); }
    }

    public boolean strictSave(AuthHeader auth, BulkRecord bulkRecord, Validator validator) {
        lockTheIndexOrFail(bulkRecord.repo, bulkRecord.model);
        try {
            return bulkSave(auth, bulkRecord, validator);
        } finally { unlockIndex(bulkRecord.repo, bulkRecord.model); }
    }

    @SneakyThrows public String bulkSave(AuthHeader auth, IndexRecord record) {
        try {
            XContentBuilder result = XContentFactory.jsonBuilder().startObject();
            save(auth, record, result);
            client.forceRefresh(auth, record.repo);
            return Strings.toString(result.endObject());

        } catch (NoopException nex) {
            return nex.getMessage();
        } catch (RuntimeException ex) { throw logError(ex); }
    }

    public void validateRecord(AuthHeader auth, IndexRecord record, BulkRequest bulkRequest, Validator validator) {
        IndexRequest archiveRequest = validator.validate(client, auth, record);
        bulkRequest.add(record.request());
        if (archiveRequest != null) { bulkRequest.add(archiveRequest); }
    }

    public String strictSave(AuthHeader auth, Pair pair) { return strictSave(auth, pair._a, pair._b, STRICT); }

    public String strictSave(AuthHeader auth, IndexRecord record, Validator validator) { return strictSave(auth, record, null, validator); }

    @SneakyThrows
    private String strictSave(AuthHeader auth, IndexRecord primary, IndexRecord aux, Validator validator) {
        if (checkLock(primary.repo, primary.model)) { throw new IndexerException("index-model locked for write"); }

        String lockString = nonNull(aux, primary).lockString();
        lockTheIndexOrFail(primary.repo, lockString);
        try {
            XContentBuilder result = XContentFactory.jsonBuilder().startObject();
            writeStrict(auth, primary, validator, aux == null ? result : null);
            if (aux != null) { writeStrict(auth, aux, Validator.OVERRIDE, result); }
            return Strings.toString(result.endObject());

        } catch (NoopException nex) {
            return nex.toJson();
        } catch (RuntimeException ex) {
            throw logError(ex);
        } finally { unlockIndex(primary.repo, lockString); }
    }

    private void writeStrict(AuthHeader auth, IndexRecord record, Validator validator, XContentBuilder result) {
        IndexRequest archiveRequest = validator.validate(client, auth, record);
        save(auth, record, result);
        if (archiveRequest != null) { client.index(auth, archiveRequest); }
    }

    @SneakyThrows private void save(AuthHeader auth, IndexRecord record, XContentBuilder result) {
        DocWriteRequest request = record.request();
        DocWriteResponse response = (request instanceof IndexRequest) ? client.index(auth, (IndexRequest) request) :
            ((request instanceof DeleteRequest) ? client.delete(auth, (DeleteRequest) request) : null);
        if (response != null && result != null) {
            record.toAction(result, nonNull(record.result(), response.getResult().getLowercase()));
        }
    }

    private static boolean checkResponse(BulkResponse responses, IgnoreErrors ignore, BiConsumer consumer) {
        boolean noErrors = true; int index = 0;
        for (BulkItemResponse response : responses) {
            if (response.isFailed()) {
                String failureMessage = response.getFailureMessage();
                if (!ignore.ignoreFailure(failureMessage)) {
                    noErrors = false; consumer.accept(index, failureMessage);
                }
                else { consumer.accept(index, "accepted"); }
            }
            else { consumer.accept(index, response.getResponse().getResult().getLowercase()); }
            index++;
        }
        return noErrors;
    }

    private static RuntimeException logError(RuntimeException ex) {
        if (ex instanceof IndexerException) { log.error("two phase commit failed: " + ex.getMessage()); return ex; }
        else { log.error("two phase commit failed: " + ex.getMessage(), ex); return new IndexerException(ex.getMessage(), ex); }
    }
}