fr.ouestfrance.querydsl.postgrest.services.BulkExecutorService Maven / Gradle / Ivy

Go to download
package fr.ouestfrance.querydsl.postgrest.services;

import fr.ouestfrance.querydsl.postgrest.annotations.PostgrestConfiguration;
import fr.ouestfrance.querydsl.postgrest.model.*;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

/**
 * Bulk Executor service allow to execute on a wide range of data
 * Options allow you to specify bulk parameters :
 * - countOnly : Do not get object representation
 * - pageSize: Specify chunk size, allow to iterate on postgrest and avoid timeout exceptions
 */
public class BulkExecutorService {

    public  BulkResponse execute(Function> function, BulkRequest request, BulkOptions options) {
        List prefers = request.getHeaders().computeIfAbsent(Prefer.HEADER, x -> new ArrayList<>());
        if (options.isCountOnly()) {
            prefers.stream().filter(x -> x.startsWith("return")).findFirst().ifPresent(prefers::remove);
            prefers.add(Prefer.Return.HEADERS_ONLY);
        }
        // Add default count
        prefers.add("count=" + PostgrestConfiguration.CountType.EXACT.name().toLowerCase());
        // If we want to bulk update (multiple page)
        Pageable pageable = null;
        if (options.getPageSize() > 0) {
            pageable = Pageable.ofSize(options.getPageSize());
            request.getHeaders().put("Range-Unit", List.of("items"));
        }
        // Do first call
        BulkResponse response = new BulkResponse<>(null, 0, 0);
        response.merge(function.apply(request));
        // If pageable is null -> Retrieve default pageable from the first response
        if (pageable == null) {
            pageable = Pageable.ofSize((int) response.getAffectedRows());
        }

        // If everything is not found then do next calls
        while (hasNextPage(response)) {
            // Start on page 1
            pageable = pageable.next();
            request.getHeaders().put("Range", List.of(pageable.toRange()));
            response.merge(function.apply(request));
        }
        return response;
    }

    private static  boolean hasNextPage(BulkResponse response) {
        return response.getTotalElements() > response.getAffectedRows();
    }
}