de.otto.flummi.response.ScrollingSearchHits Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ports-flummi Show documentation
The newest version!
package de.otto.flummi.response;

import java.util.Iterator;
import java.util.List;
import java.util.Spliterator;
import java.util.function.Consumer;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import de.mhus.lib.core.logging.Log;
import de.otto.flummi.request.SearchScrollRequestBuilder;
import de.otto.flummi.util.HttpClientWrapper;

/**
 * Auto-scrolling implementation of SearchHits. Contains a page of search results
 * and automatically fetches more pages from the server as you iterate or stream over the search result.
 */
public class ScrollingSearchHits implements SearchHits {
    private final long totalHits;
    private final Float maxScore;
    private final String scrollId;
    private final String scroll;
    private final HttpClientWrapper client;
    private List hitsCurrentPage;
    private boolean dirty;
    public static final Log LOG = Log.getLog(ScrollingSearchHits.class);


    public ScrollingSearchHits(long totalHits, Float maxScore, String scrollId, String scroll, List hitsCurrentPage, HttpClientWrapper client) {
        this.totalHits = totalHits;
        this.maxScore = maxScore;
        this.scrollId = scrollId;
        this.scroll = scroll;
        this.hitsCurrentPage = hitsCurrentPage;
        this.client = client;
    }


    @Override
    public long getTotalHits() {
        return totalHits;
    }

    @Override
    public Float getMaxScore() {
        return maxScore;
    }

    @Override
    public Iterator iterator() {
        assertNotDirty();
        return new Iterator() {
            int currentPageIdx = 0;

            @Override
            public boolean hasNext() {
                if (currentPageIdx < hitsCurrentPage.size()) {
                    return true;
                }
                if(hitsCurrentPage.isEmpty()) {
                    return false;
                }
                fetchNextPage();
                currentPageIdx = 0;
                return !hitsCurrentPage.isEmpty();
            }

            @Override
            public SearchHit next() {
                if(currentPageIdx>0 && currentPageIdx==hitsCurrentPage.size()) {
                    fetchNextPage();
                    currentPageIdx = 0;
                }
                return hitsCurrentPage.get(currentPageIdx++);
            }
        };
    }

    private void assertNotDirty() {
        if(dirty) {
            throw new IllegalStateException("Result was already iterated / streamed before");
        }
    }

    private void fetchNextPage() {
        dirty = true;
        SearchResponse response = new SearchScrollRequestBuilder(client)
                .setScroll(scroll)
                .setScrollId(scrollId)
                .execute();
        this.hitsCurrentPage = ((SimpleSearchHits)response.getHits()).getHits();
    }

    @Override
    public void forEach(Consumer action) {
        assertNotDirty();
        while(!hitsCurrentPage.isEmpty()) {
            hitsCurrentPage.forEach(action);
            fetchNextPage();
        }
    }

    @Override
    public Spliterator spliterator() {
        assertNotDirty();
        return new Spliterator() {
            Iterator iterator = iterator();
            @Override
            public boolean tryAdvance(Consumer action) {
                if(iterator.hasNext()) {
                    action.accept(iterator.next());
                    return true;
                }
                return false;
            }

            @Override
            public Spliterator trySplit() {
                return null;
            }

            @Override
            public long estimateSize() {
                return totalHits;
            }

            @Override
            public int characteristics() {
                return ORDERED | SIZED | NONNULL | IMMUTABLE;
            }
        };
    }

    @Override
    public Stream stream() {
        return StreamSupport.stream(spliterator(), false);
    }
}