ElasticSearchWrapper.kt 13.1 KB
Newer Older
Jonas Waeber's avatar
Jonas Waeber committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
/*
 * search-doc-service
 * Copyright (C) 2020 Memoriav
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
18
package ch.memobase.helpers
19

20
21
import com.beust.klaxon.Klaxon
import com.beust.klaxon.KlaxonException
22
23
import org.apache.logging.log4j.LogManager
import org.elasticsearch.ElasticsearchException
Jonas Waeber's avatar
Jonas Waeber committed
24
import org.elasticsearch.action.get.GetRequest
25
26
27
import org.elasticsearch.action.search.ClearScrollRequest
import org.elasticsearch.action.search.SearchRequest
import org.elasticsearch.action.search.SearchScrollRequest
28
29
30
import org.elasticsearch.client.RequestOptions
import org.elasticsearch.client.RestHighLevelClient
import org.elasticsearch.client.core.CountRequest
31
import org.elasticsearch.common.unit.TimeValue
32
33
import org.elasticsearch.index.query.QueryBuilder
import org.elasticsearch.index.query.QueryBuilders.boolQuery
34
35
36
import org.elasticsearch.index.query.QueryBuilders.termQuery
import org.elasticsearch.search.Scroll
import org.elasticsearch.search.builder.SearchSourceBuilder
37
38
import ch.memobase.model.FacetContainer
import ch.memobase.model.LanguageContainer
39

40

Jonas Waeber's avatar
Jonas Waeber committed
41
42
43
44
/**
 * This class facilitates a connection to the Elasticsearch cluster and offers convenience functions to retrieve
 * the necessary data.
 */
45
46
class ElasticSearchWrapper(
    private val client: RestHighLevelClient,
47
48
49
50
    private val translationMappers: TranslationMappers,
    private val documentsIndex: String,
    private val institutionIndex: String,
    private val recordSetIndex: String
51
) {
52
    private val log = LogManager.getLogger("ElasticSearchWrapper")
53
54
    private val klaxon = Klaxon()

55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

    fun countTotalNumberOfDocuments(recordSetIdentifier: String): Int {
        return countNumberOfDocuments(
            recordSetIdentifier,
            termQuery(
                "recordSet.facet", recordSetIdentifier
            )
        )
    }

    fun countNumberOfDocumentsPublished(recordSetIdentifier: String): Int {
        return countNumberOfDocuments(
            recordSetIdentifier,
            boolQuery()
                .must(termQuery("recordSet.facet", recordSetIdentifier))
                .must(termQuery("published", true))
        )
    }

Jonas Waeber's avatar
Jonas Waeber committed
74
75
    /**
     * Counts the number of documents attached to a specific record set.
Jonas Waeber's avatar
Jonas Waeber committed
76
77
78
79
     *
     * @param recordSetIdentifier The memobase identifier of the record set.
     *
     * @return Number of Documents
Jonas Waeber's avatar
Jonas Waeber committed
80
     */
81
    private fun countNumberOfDocuments(recordSetIdentifier: String, query: QueryBuilder): Int {
82
        log.info("Counting documents for record set $recordSetIdentifier from index $documentsIndex.")
83
        val request = CountRequest(documentsIndex)
84
        request.query(query)
85
86
87
88
        return try {
            val response = client.count(
                request, RequestOptions.DEFAULT
            )
Jonas Waeber's avatar
Jonas Waeber committed
89
            val count = response.count.toInt()
90
            log.info("Found $count documents for record set $recordSetIdentifier from index $documentsIndex.")
Jonas Waeber's avatar
Jonas Waeber committed
91
            count
92
93
94
        } catch (ex: ElasticsearchException) {
            log.error(ex.detailedMessage)
            0
95
96
        }
    }
97
98

    /**
Jonas Waeber's avatar
Jonas Waeber committed
99
100
101
102
103
104
     * Collects all the document types present in the index for a specific record set.
     *
     * @param recordSetIdentifier: The memobase identifier of the record set.
     * @param queryField: Name of the field the identifier is stored in.
     *
     * @return A list of unique document type facet containers.
105
106
     */
    fun getDocumentTypesFromRecords(recordSetIdentifier: String, queryField: String): List<FacetContainer> {
107
        return try {
108
            log.info("Attempting to load document type for $recordSetIdentifier in field $queryField from index $documentsIndex.")
109
110
111
112
113
114
115
116
117
118
119
            val resultFacets = mutableListOf<FacetContainer>()
            val typeSet = mutableSetOf<String>()
            val scroll = Scroll(TimeValue.timeValueMinutes(1L))
            val searchRequest = SearchRequest(documentsIndex)
            searchRequest.scroll(scroll)
            val searchSourceBuilder = SearchSourceBuilder()
            searchSourceBuilder.fetchSource(
                arrayOf(
                    "id", "type.filter"
                ), emptyArray<String>()
            )
120

121
122
123
124
125
126
127
128
129
            searchSourceBuilder.query(
                termQuery(
                    queryField, recordSetIdentifier
                )
            )
            searchRequest.source(searchSourceBuilder)
            var searchResponse = client.search(searchRequest, RequestOptions.DEFAULT)
            var scrollId = searchResponse.scrollId
            var searchHits = searchResponse.hits.hits
130

131
132
133
134
135
136
137
138
139
            while (searchHits != null && searchHits.isNotEmpty()) {
                for (hit in searchHits) {
                    val source = hit.sourceAsString
                    try {
                        val document = klaxon.parse<DocumentResponseSource>(source)
                        if (document != null) {
                            if (!typeSet.contains(document.type.filter)) {
                                resultFacets.add(translationMappers.getDocumentType(document.type.filter))
                                typeSet.add(document.type.filter)
140
141
                            }
                        }
142
143
                    } catch (ex: KlaxonException) {
                        log.error("Unable to parse $source from index $documentsIndex.")
144
145
                    }
                }
146
147
148
149
150
                val scrollRequest = SearchScrollRequest(scrollId)
                scrollRequest.scroll(scroll)
                searchResponse = client.scroll(scrollRequest, RequestOptions.DEFAULT)
                scrollId = searchResponse.scrollId
                searchHits = searchResponse.hits.hits
151
            }
152
153
154
            val clearScrollRequest = ClearScrollRequest()
            clearScrollRequest.addScrollId(scrollId)
            client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT)
155
            log.info("Found the following types $typeSet for record set $recordSetIdentifier from index $documentsIndex.")
156
157
158
159
            resultFacets
        } catch (ex: ElasticsearchException) {
            log.error(ex.detailedMessage)
            emptyList()
160
161
        }
    }
Jonas Waeber's avatar
Jonas Waeber committed
162

163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
    /***
     * This function exists because some record sets are managed by multiple institutions. However, the document itself
     * always assumes that there is only one. As a result the institution names must be retrieved from the
     * record set in that case.
     */
    fun getInstitutionNamesFromRecordSet(recordSetIdentifier: String): List<FacetContainer> {
        return try {
            log.debug("Attempting to retrieve record set document.")
            val request = GetRequest(recordSetIndex, recordSetIdentifier)
            val response = client.get(request, RequestOptions.DEFAULT)
            if (response.isExists) {
                val map = response.sourceAsMap["institution"]
                if (map == null) {
                    log.error("Found record set for id $recordSetIdentifier in " +
                            "index $recordSetIndex, but could not extract name.")
                    listOf(FacetContainer(
                        LanguageContainer.EMPTY,
                        recordSetIdentifier,
                        emptyList()
                    ))
                } else {
                    log.info("Retrieval of institution names was successful.")
                    when (map) {
                        is List<*> -> {
                            map.map {
                                FacetContainer.fromMap(it)
                            }
                        }
                        is Map<*, *> -> {
                           listOf(FacetContainer.fromMap(map))
                        }
                        else -> {
                            log.error("The retrieved data is not valid: $map.")
                            listOf(FacetContainer.EMPTY)
                        }
                    }
                }
            } else {
                log.error("Could not find record set id $recordSetIdentifier in index $recordSetIndex.")
                listOf(FacetContainer(
                    LanguageContainer.EMPTY,
                    recordSetIdentifier,
                    emptyList()
                ))
            }
        } catch (ex: ElasticsearchException) {
            log.error(ex.detailedMessage)
            listOf(FacetContainer(
                LanguageContainer.EMPTY,
                recordSetIdentifier,
                emptyList()
            ))
        }
    }

Jonas Waeber's avatar
Jonas Waeber committed
218
219
    fun getInstitutionName(identifier: String): FacetContainer {
        return try {
220
            log.debug("Attempting to retrieve institution record.")
Jonas Waeber's avatar
Jonas Waeber committed
221
222
223
            val request = GetRequest(institutionIndex, identifier)
            val response = client.get(request, RequestOptions.DEFAULT)
            if (response.isExists) {
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
                val map = response.sourceAsMap["name"]
                if (map == null) {
                    log.error("Found institution for id $identifier in index $institutionIndex, but could not extract name.")
                    FacetContainer(
                        LanguageContainer.EMPTY,
                        identifier,
                        emptyList()
                    )
                } else {
                    log.debug("Retrieval of institution names was successful.")
                    FacetContainer(
                        LanguageContainer.fromMap(map),
                        identifier,
                        emptyList()
                    )
                }
Jonas Waeber's avatar
Jonas Waeber committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
            } else {
                log.error("Could not find institution $identifier in index $institutionIndex.")
                FacetContainer(
                    LanguageContainer.EMPTY,
                    identifier,
                    emptyList()
                )
            }
        } catch (ex: ElasticsearchException) {
            log.error(ex.detailedMessage)
            FacetContainer(
                LanguageContainer.EMPTY,
                identifier,
                emptyList()
            )
        }
    }
257
258
259

    fun getRecordSetName(identifier: String): LanguageContainer {
        return try {
260
            log.debug("Attempting to retrieve record set document from $recordSetIndex.")
261
262
263
            val request = GetRequest(recordSetIndex, identifier)
            val response = client.get(request, RequestOptions.DEFAULT)
            if (response.isExists) {
264
265
266
267
268
269
270
271
                val map = response.sourceAsMap["name"]
                if (map == null) {
                    log.error("Found record set for id $identifier in $recordSetIndex, but could not retrieve names in field 'name'.")
                    LanguageContainer.EMPTY
                } else {
                    log.debug("Retrieval of record set names was successful.")
                    LanguageContainer.fromMap(map)
                }
272
273
274
275
276
277
278
279
280
            } else {
                log.error("Could not find record set $identifier in index $recordSetIndex.")
                LanguageContainer.EMPTY
            }
        } catch (ex: ElasticsearchException) {
            log.error(ex.detailedMessage)
            LanguageContainer.EMPTY
        }
    }
281
282
283
284
285
286
287

    fun getExtraInstitutionsFromRecordSet(identifier: String, extraType: String): List<FacetContainer> {
        return try {
            log.debug("Attempting to retrieve record set document from $recordSetIndex.")
            val request = GetRequest(recordSetIndex, identifier)
            val response = client.get(request, RequestOptions.DEFAULT)
            if (response.isExists) {
Jonas Waeber's avatar
Jonas Waeber committed
288
                when (val map = response.sourceAsMap[extraType]) {
289
                    is List<*> -> {
Jonas Waeber's avatar
Jonas Waeber committed
290
                        log.debug("Retrieval of '$extraType' was successful.")
291
292
293
                        map.map { FacetContainer.fromMap(it) }
                    }
                    null -> {
Jonas Waeber's avatar
Jonas Waeber committed
294
                        log.error("Found record set for id $identifier in $recordSetIndex, but could not retrieve names in field '$extraType'.")
295
296
297
                        emptyList()
                    }
                    is Map<*, *> -> {
Jonas Waeber's avatar
Jonas Waeber committed
298
                        log.debug("Retrieval of '$extraType' was successful.")
299
300
301
                        listOf(FacetContainer.fromMap(map))
                    }
                    else -> {
Jonas Waeber's avatar
Jonas Waeber committed
302
                        log.error("Found record set for id $identifier in $recordSetIndex, but could not retrieve names in field '$extraType'.")
303
304
305
306
307
308
309
310
311
312
313
314
                        emptyList()
                    }
                }
            } else {
                log.error("Could not find record set $identifier in index $recordSetIndex.")
                emptyList()
            }
        } catch (ex: ElasticsearchException) {
            log.error(ex.detailedMessage)
            emptyList()
        }
    }
315
}