HeatingPaymentsInvoker/conf/script/1000/commons/ESMap.groovy

import com.dySweetFishPlugin.elasticsearch.ESClient
import com.sweetfish.convert.json.JsonConvert
import com.yinjie.heating.common.api.QueryBuilderExecutor
import com.yinjie.heating.common.api.SortBuilderExecutor
import org.elasticsearch.action.search.SearchRequestBuilder
import org.elasticsearch.action.search.SearchResponse
import org.elasticsearch.common.unit.TimeValue
import org.elasticsearch.index.query.QueryBuilder

import java.util.function.ToLongFunction

/**
 * Created by jlutt on 2022-04-20
 *
 * @author jlutt
 */

class ESMap<T> {

    private long[] src

    private ESClient esClient

    private JsonConvert jsonConvert

    private Class<T> clazz

    private QueryBuilderExecutor qbe

    private String index

    private ToLongFunction<T> tl

    private List<SortBuilderExecutor> sbeList = new ArrayList<>()

    static <T> ESMap<T> getESMap() {
        return new ESMap<>()
    }

    ESMap<T> srcIds(long[] value) {
        this.src = value
        return this
    }

    ESMap<T> esClient(ESClient value) {
        this.esClient = value
        return this
    }

    ESMap<T> jsonConvert(JsonConvert value) {
        this.jsonConvert = value
        return this
    }

    ESMap<T> clazz(Class<T> value) {
        this.clazz = value
        return this
    }

    ESMap<T> queryBuilder(QueryBuilderExecutor value) {
        this.qbe = value
        return this
    }

    ESMap<T> sortBuilder(SortBuilderExecutor value) {
        this.sbeList.add(value)
        return this
    }

    ESMap<T> index(String value) {
        this.index = value
        return this
    }

    ESMap<T> convertLongExecutor(ToLongFunction<T> value) {
        this.tl = value
        return this
    }

    Map<Long, List<T>> execute() {
        QueryBuilder qb = qbe.execute(null)

        SearchRequestBuilder requestBuilder = esClient.getClient().prepareSearch(index).setQuery(qb)

        if ((sbeList != null) && (!sbeList.isEmpty())) {
            for (SortBuilderExecutor sortBuilderExecutor : sbeList) {
                if (sortBuilderExecutor != null) {
                    requestBuilder = requestBuilder.addSort(sortBuilderExecutor.execute())
                }
            }
        }

        SearchResponse scrollResp = requestBuilder.setSize(100)
                .setScroll(new TimeValue(60000))
                .execute()
                .actionGet()

        HashMap<Long, List<T>> dataMap = new HashMap<>()

        while (scrollResp.getHits().getHits().length != 0) {
            scrollResp.getHits().getHits().each { x ->
                String json = x.getSourceAsString()
                T result = jsonConvert.convertFrom(clazz, json)
                long iValue = tl.applyAsLong(result)
                if (!dataMap.containsKey(iValue)) {
                    dataMap.put(iValue, new ArrayList<>())
                }
                dataMap.get(iValue).add(result)
            }

            scrollResp = esClient.getClient().prepareSearchScroll(scrollResp.getScrollId()).setScroll(new TimeValue(60000)).execute().actionGet()
        }

        src.each { l ->
            if (!dataMap.containsKey(l)) {
                dataMap.put(l, new ArrayList<>())
            }
        }

        return dataMap
    }
}