-
Notifications
You must be signed in to change notification settings - Fork 0
/
search_after.clj
56 lines (50 loc) · 1.98 KB
/
search_after.clj
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
(ns scroll.search-after
(:require [clojure.tools.logging :as log]
[scroll.batch :as batch]
[scroll.hits :as hits]
[scroll.request :as request]))
(def default-query {:sort ["_doc"]})
(defn remove-from-param [query]
(dissoc query :from))
(defn start [es-host index-name query opts]
(request/execute-request
{:url (format "%s/%s/_search" es-host (or index-name "*"))
:body (remove-from-param
(batch/set-batch-size
(if (:sort query)
(or query default-query)
(assoc query :sort [:_doc]))
opts))
:opts opts}))
(defn continue [es-host index-name query search-after opts]
(request/execute-request
{:url (format "%s/%s/_search" es-host (or index-name "*"))
:body (remove-from-param
(batch/set-batch-size
(if (:sort query)
(assoc (or query default-query) :search_after search-after)
(assoc query :sort [:_doc]
:search_after search-after)) opts))
:opts opts}))
(defn extract-search-after [batch keywordize?]
(if keywordize?
(-> (get-in batch [:hits :hits])
last
(get :sort))
(-> (get-in batch ["hits" "hits"])
last
(get "sort"))))
(defn fetch [{:keys [es-host index-name query search-after opts] :as req}]
(try
(let [batch (if search-after
(continue es-host index-name query search-after opts)
(start es-host index-name query opts))]
(log/debugf "Fetching a batch took: %s ms" (or (get batch :took) (get batch "took")))
(when-let [current-hits (seq (hits/extract-hits batch (get opts :keywordize?)))]
(lazy-cat current-hits
(when-let [sa (extract-search-after batch (get opts :keywordize?))]
(fetch (assoc req :search-after sa))))))
(catch Exception e
(.printStackTrace e)
(log/errorf "Failed to search_after: %s" e)
[])))