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.Collections;
024import java.util.Comparator;
025import java.util.HashMap;
026import java.util.List;
027import java.util.Map;
028
029import org.elasticsearch.action.search.SearchResponse;
030import org.elasticsearch.common.text.Text;
031import org.elasticsearch.search.SearchHit;
032import org.elasticsearch.search.highlight.HighlightField;
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                    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.get(repoName);
088            if (docIds == null) {
089                docIds = new ArrayList<>();
090                ret.put(repoName, docIds);
091            }
092            docIds.add(hit.getId());
093        }
094        return ret;
095    }
096
097    private List<DocumentModel> fetchFromVcs(List<String> ids, CoreSession session) {
098        List<DocumentModel> ret = null;
099        int size = ids.size();
100        int start = 0;
101        int end = Math.min(CHUNK_SIZE, size);
102        boolean done = false;
103        while (!done) {
104            List<DocumentModel> docs = fetchFromVcsChunk(ids.subList(start, end), session);
105            if (ret == null) {
106                ret = docs;
107            } else {
108                ret.addAll(docs);
109            }
110            if (end >= ids.size()) {
111                done = true;
112            } else {
113                start = end;
114                end = Math.min(start + CHUNK_SIZE, size);
115            }
116        }
117        return ret;
118    }
119
120    private List<DocumentModel> fetchFromVcsChunk(final List<String> ids, CoreSession session)
121    {
122        StringBuilder sb = new StringBuilder();
123        sb.append("SELECT * FROM Document, Relation WHERE ecm:uuid IN (");
124        for (int i = 0; i < ids.size(); i++) {
125            sb.append(NXQL.escapeString(ids.get(i)));
126            if (i < ids.size() - 1) {
127                sb.append(", ");
128            }
129        }
130        sb.append(")");
131        return session.query(sb.toString());
132    }
133
134    private void addHighlights(List<DocumentModel> docs) {
135        for (SearchHit hit : getResponse().getHits()) {
136            for (DocumentModel doc : docs) {
137                String docId = doc.getRepositoryName() + doc.getId();
138                String hitId = getRepoForIndex(hit.getIndex()) + hit.getId();
139                if (docId.equals(hitId)) {
140                    // Add highlight if it exists
141                    Map<String, HighlightField> esHighlights = hit.highlightFields();
142                    if (!esHighlights.isEmpty()) {
143                        Map<String, List<String>> fields = new HashMap<>();
144                        for (Map.Entry<String, HighlightField> entry : esHighlights.entrySet()) {
145                            String field = entry.getKey();
146                            List<String> list = new ArrayList<>();
147                            for (Text fragment : entry.getValue().getFragments()) {
148                                list.add(fragment.toString());
149                            }
150                            fields.put(field, list);
151                        }
152                        doc.putContextData(PageProvider.HIGHLIGHT_CTX_DATA, (Serializable) fields);
153                    }
154                    break;
155                }
156            }
157        }
158    }
159
160    private void sortResults(List<DocumentModel> docs) {
161        final List<String> ids = new ArrayList<>();
162        for (SearchHit hit : getResponse().getHits()) {
163            ids.add(getRepoForIndex(hit.getIndex()) + hit.getId());
164        }
165
166        Collections.sort(docs, new Comparator<DocumentModel>() {
167            @Override
168            public int compare(DocumentModel a, DocumentModel b) {
169                return ids.indexOf(a.getRepositoryName() + a.getId()) - ids.indexOf(b.getRepositoryName() + b.getId());
170            }
171        });
172
173    }
174
175}