001/*
002 * (C) Copyright 2019 Nuxeo (http://nuxeo.com/) and others.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 *
016 * Contributors:
017 *     bdelbosc
018 */
019package org.nuxeo.elasticsearch.scroll;
020
021import java.util.Arrays;
022import java.util.List;
023import java.util.NoSuchElementException;
024import java.util.stream.Collectors;
025
026import org.elasticsearch.search.SearchHit;
027import org.nuxeo.ecm.core.scroll.RepositoryScroll;
028import org.nuxeo.elasticsearch.api.ElasticSearchService;
029import org.nuxeo.elasticsearch.api.EsScrollResult;
030import org.nuxeo.elasticsearch.query.NxQueryBuilder;
031import org.nuxeo.runtime.api.Framework;
032
033/**
034 * @since 11.1
035 */
036public class ElasticSearchScroll extends RepositoryScroll {
037
038    protected EsScrollResult esScroll;
039
040    @Override
041    public boolean hasNext() {
042        if (hasNextResult == null) {
043            hasNextResult = fetch();
044        }
045        return hasNextResult;
046    }
047
048    @Override
049    protected boolean fetch() {
050        ElasticSearchService ess = Framework.getService(ElasticSearchService.class);
051        if (esScroll == null) {
052            esScroll = ess.scroll(new NxQueryBuilder(session).nxql(request.getQuery())
053                                                             .limit(request.getSize())
054                                                             .onlyElasticsearchResponse(),
055                    request.getTimeout().toSeconds());
056        } else {
057            esScroll = ess.scroll(esScroll);
058        }
059        SearchHit[] hits = esScroll.getElasticsearchResponse().getHits().getHits();
060        return hits != null && hits.length > 0;
061    }
062
063    @Override
064    public List<String> next() {
065        if (hasNextResult == null) {
066            hasNextResult = fetch();
067        }
068        if (!hasNextResult) {
069            throw new NoSuchElementException();
070        }
071        hasNextResult = null;
072        SearchHit[] hits = esScroll.getElasticsearchResponse().getHits().getHits();
073        return Arrays.stream(hits).map(SearchHit::getId).collect(Collectors.toList());
074    }
075
076    @Override
077    public void close() {
078        super.close();
079        esScroll = null;
080    }
081}