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 = new ArrayList<>();
054        String openSessionRepository = getSession().getRepositoryName();
055        boolean closeSession;
056        CoreSession session;
057        for (String repo : repoHits.keySet()) {
058            if (openSessionRepository.equals(repo)) {
059                session = getSession();
060                closeSession = false;
061            } else {
062                session = CoreInstance.openCoreSession(repo);
063                closeSession = true;
064            }
065            try {
066                docs.addAll(fetchFromVcs(repoHits.get(repo), session));
067            } finally {
068                if (closeSession) {
069                    session.close();
070                }
071            }
072        }
073        sortResults(docs);
074        addHighlights(docs);
075        DocumentModelListImpl ret = new DocumentModelListImpl(docs.size());
076        if (!docs.isEmpty()) {
077            ret.addAll(docs);
078        }
079        return ret;
080    }
081
082    private Map<String, List<String>> getHitsPerRepository() {
083        Map<String, List<String>> ret = new HashMap<>();
084        for (SearchHit hit : getResponse().getHits()) {
085            String repoName = getRepoForIndex(hit.getIndex());
086            List<String> docIds = ret.computeIfAbsent(repoName, k -> new ArrayList<>());
087            docIds.add(hit.getId());
088        }
089        return ret;
090    }
091
092    private List<DocumentModel> fetchFromVcs(List<String> ids, CoreSession session) {
093        List<DocumentModel> ret = null;
094        int size = ids.size();
095        int start = 0;
096        int end = Math.min(CHUNK_SIZE, size);
097        boolean done = false;
098        while (!done) {
099            List<DocumentModel> docs = fetchFromVcsChunk(ids.subList(start, end), session);
100            if (ret == null) {
101                ret = docs;
102            } else {
103                ret.addAll(docs);
104            }
105            if (end >= ids.size()) {
106                done = true;
107            } else {
108                start = end;
109                end = Math.min(start + CHUNK_SIZE, size);
110            }
111        }
112        return ret;
113    }
114
115    private List<DocumentModel> fetchFromVcsChunk(final List<String> ids, CoreSession session) {
116        StringBuilder sb = new StringBuilder();
117        sb.append("SELECT * FROM Document, Relation WHERE ecm:uuid IN (");
118        for (int i = 0; i < ids.size(); i++) {
119            sb.append(NXQL.escapeString(ids.get(i)));
120            if (i < ids.size() - 1) {
121                sb.append(", ");
122            }
123        }
124        sb.append(")");
125        return session.query(sb.toString());
126    }
127
128    private void addHighlights(List<DocumentModel> docs) {
129        for (SearchHit hit : getResponse().getHits()) {
130            for (DocumentModel doc : docs) {
131                String docId = doc.getRepositoryName() + doc.getId();
132                String hitId = getRepoForIndex(hit.getIndex()) + hit.getId();
133                if (docId.equals(hitId)) {
134                    // Add highlight if it exists
135                    Map<String, HighlightField> esHighlights = hit.getHighlightFields();
136                    if (!esHighlights.isEmpty()) {
137                        Map<String, List<String>> fields = new HashMap<>();
138                        for (Map.Entry<String, HighlightField> entry : esHighlights.entrySet()) {
139                            String field = entry.getKey();
140                            List<String> list = new ArrayList<>();
141                            for (Text fragment : entry.getValue().getFragments()) {
142                                list.add(fragment.toString());
143                            }
144                            fields.put(field, list);
145                        }
146                        doc.putContextData(PageProvider.HIGHLIGHT_CTX_DATA, (Serializable) fields);
147                    }
148                    break;
149                }
150            }
151        }
152    }
153
154    private void sortResults(List<DocumentModel> docs) {
155        final List<String> ids = new ArrayList<>();
156        for (SearchHit hit : getResponse().getHits()) {
157            ids.add(getRepoForIndex(hit.getIndex()) + hit.getId());
158        }
159
160        docs.sort(new Comparator<DocumentModel>() {
161            @Override
162            public int compare(DocumentModel a, DocumentModel b) {
163                return ids.indexOf(a.getRepositoryName() + a.getId()) - ids.indexOf(b.getRepositoryName() + b.getId());
164            }
165        });
166
167    }
168
169}