001/*
002 * (C) Copyright 2014 Nuxeo SA (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.fetcher;
020
021import java.io.Serializable;
022import java.util.ArrayList;
023import java.util.Comparator;
024import java.util.HashMap;
025import java.util.List;
026import java.util.Map;
027
028import org.elasticsearch.action.search.SearchResponse;
029import org.elasticsearch.common.text.Text;
030import org.elasticsearch.search.SearchHit;
031import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
032import org.nuxeo.ecm.core.api.CoreInstance;
033import org.nuxeo.ecm.core.api.CoreSession;
034import org.nuxeo.ecm.core.api.DocumentModel;
035import org.nuxeo.ecm.core.api.impl.DocumentModelListImpl;
036import org.nuxeo.ecm.core.query.sql.NXQL;
037import org.nuxeo.ecm.platform.query.api.PageProvider;
038
039/**
040 * @since 6.0
041 */
042public class VcsFetcher extends Fetcher {
043
044    private static final int CHUNK_SIZE = 100;
045
046    public VcsFetcher(CoreSession session, SearchResponse response, Map<String, String> repoNames) {
047        super(session, response, repoNames);
048    }
049
050    @Override
051    public DocumentModelListImpl fetchDocuments() {
052        Map<String, List<String>> repoHits = getHitsPerRepository();
053        List<DocumentModel> docs = fetchFromVcs(repoHits);
054        sortResults(docs);
055        addHighlights(docs);
056        DocumentModelListImpl ret = new DocumentModelListImpl(docs.size());
057        if (!docs.isEmpty()) {
058            ret.addAll(docs);
059        }
060        return ret;
061    }
062
063    @SuppressWarnings("resource") // session closed only if we opened it
064    protected List<DocumentModel> fetchFromVcs(Map<String, List<String>> repoHits) {
065        List<DocumentModel> docs = new ArrayList<>();
066        String openSessionRepository = getSession().getRepositoryName();
067        for (String repo : repoHits.keySet()) {
068            CoreSession session;
069            if (openSessionRepository.equals(repo)) {
070                session = getSession();
071            } else {
072                session = CoreInstance.getCoreSession(repo);
073            }
074            docs.addAll(fetchFromVcs(repoHits.get(repo), session));
075        }
076        return docs;
077    }
078
079    private Map<String, List<String>> getHitsPerRepository() {
080        Map<String, List<String>> ret = new HashMap<>();
081        for (SearchHit hit : getResponse().getHits()) {
082            String repoName = getRepoForIndex(hit.getIndex());
083            List<String> docIds = ret.computeIfAbsent(repoName, k -> new ArrayList<>());
084            docIds.add(hit.getId());
085        }
086        return ret;
087    }
088
089    private List<DocumentModel> fetchFromVcs(List<String> ids, CoreSession session) {
090        List<DocumentModel> ret = null;
091        int size = ids.size();
092        int start = 0;
093        int end = Math.min(CHUNK_SIZE, size);
094        boolean done = false;
095        while (!done) {
096            List<DocumentModel> docs = fetchFromVcsChunk(ids.subList(start, end), session);
097            if (ret == null) {
098                ret = docs;
099            } else {
100                ret.addAll(docs);
101            }
102            if (end >= ids.size()) {
103                done = true;
104            } else {
105                start = end;
106                end = Math.min(start + CHUNK_SIZE, size);
107            }
108        }
109        return ret;
110    }
111
112    private List<DocumentModel> fetchFromVcsChunk(final List<String> ids, CoreSession session) {
113        StringBuilder sb = new StringBuilder();
114        sb.append("SELECT * FROM Document, Relation WHERE ecm:uuid IN (");
115        for (int i = 0; i < ids.size(); i++) {
116            sb.append(NXQL.escapeString(ids.get(i)));
117            if (i < ids.size() - 1) {
118                sb.append(", ");
119            }
120        }
121        sb.append(")");
122        return session.query(sb.toString());
123    }
124
125    private void addHighlights(List<DocumentModel> docs) {
126        for (SearchHit hit : getResponse().getHits()) {
127            for (DocumentModel doc : docs) {
128                String docId = doc.getRepositoryName() + doc.getId();
129                String hitId = getRepoForIndex(hit.getIndex()) + hit.getId();
130                if (docId.equals(hitId)) {
131                    // Add highlight if it exists
132                    Map<String, HighlightField> esHighlights = hit.getHighlightFields();
133                    if (!esHighlights.isEmpty()) {
134                        Map<String, List<String>> fields = new HashMap<>();
135                        for (Map.Entry<String, HighlightField> entry : esHighlights.entrySet()) {
136                            String field = entry.getKey();
137                            List<String> list = new ArrayList<>();
138                            for (Text fragment : entry.getValue().getFragments()) {
139                                list.add(fragment.toString());
140                            }
141                            fields.put(field, list);
142                        }
143                        doc.putContextData(PageProvider.HIGHLIGHT_CTX_DATA, (Serializable) fields);
144                    }
145                    break;
146                }
147            }
148        }
149    }
150
151    private void sortResults(List<DocumentModel> docs) {
152        final List<String> ids = new ArrayList<>();
153        for (SearchHit hit : getResponse().getHits()) {
154            ids.add(getRepoForIndex(hit.getIndex()) + hit.getId());
155        }
156
157        docs.sort(new Comparator<DocumentModel>() {
158            @Override
159            public int compare(DocumentModel a, DocumentModel b) {
160                return ids.indexOf(a.getRepositoryName() + a.getId()) - ids.indexOf(b.getRepositoryName() + b.getId());
161            }
162        });
163
164    }
165
166}