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