@@ -1,5 +1,5 @@
1−{-# LANGUAGE OverloadedStrings #-}
2−{-# LANGUAGE RecordWildCards #-}
1+{-# LANGUAGE OverloadedStrings #-}
2+{-# LANGUAGE RecordWildCards #-}
33
44-------------------------------------------------------------------------------
55-- |
@@ -39,6 +39,7 @@ module Database.Bloodhound.Client
3939 , searchAll
4040 , searchByIndex
4141 , searchByType
42+ , scanSearch
4243 , refreshIndex
4344 , mkSearch
4445 , mkAggregateSearch
@@ -171,6 +172,18 @@ joinPath ps = do
171172 Server s <- bhServer <$> getBHEnv
172173 return $ joinPath' (s:ps)
173174
175+appendSearchTypeParam :: Text -> SearchType -> Text
176+appendSearchTypeParam originalUrl st = addQuery [(keyEq, Just stParams)] originalUrl
177+ where keyEq = "search_type="
178+ stParams
179+ | st == SearchTypeDfsQueryThenFetch = "dfs_query_then_fetch"
180+ | st == SearchTypeCount = "count"
181+ | st == SearchTypeScan = "scan&scroll=1m"
182+ | st == SearchTypeQueryAndFetch = "query_and_fetch"
183+ | st == SearchTypeDfsQueryAndFetch = "dfs_query_and_fetch"
184+ -- used to catch 'SearchTypeQueryThenFetch', which is also the default
185+ | otherwise = "query_then_fetch"
186+
174187-- | Severely dumbed down query renderer. Assumes your data doesn't
175188-- need any encoding
176189addQuery :: [(Text, Maybe Text)] -> Text -> Text
@@ -180,11 +193,10 @@ addQuery q u = u <> rendered
180193 T.decodeUtf8 $ BB.toByteString $ NHTU.renderQueryText prependQuestionMark q
181194 prependQuestionMark = True
182195
183−
184196bindM2 :: (Applicative m, Monad m) => (a -> b -> m c) -> m a -> m b -> m c
185197bindM2 f ma mb = join (f <$> ma <*> mb)
186198
187−-- | Convenience function that sets up a mananager and BHEnv and runs
199+-- | Convenience function that sets up a manager and BHEnv and runs
188200-- the given set of bloodhound operations. Connections will be
189201-- pipelined automatically in accordance with the given manager
190202-- settings in IO. If you've got your own monad transformer stack, you
@@ -496,7 +508,8 @@ documentExists (IndexName indexName)
496508 where url = joinPath [indexName, mappingName, docId]
497509
498510dispatchSearch :: MonadBH m => Text -> Search -> m Reply
499−dispatchSearch url search = post url (Just (encode search))
511+dispatchSearch url search = post url' (Just (encode search))
512+ where url' = appendSearchTypeParam url (searchType search)
500513
501514-- | 'searchAll', given a 'Search', will perform that search against all indexes
502515-- on an Elasticsearch server. Try to avoid doing this if it can be helped.
@@ -530,6 +543,40 @@ searchByType (IndexName indexName)
530543 (MappingName mappingName) = bindM2 dispatchSearch url . return
531544 where url = joinPath [indexName, mappingName, "_search"]
532545
546+scanSearch' :: MonadBH m => Search -> m (Maybe ScrollId)
547+scanSearch' search = do
548+ let url = joinPath ["_search"]
549+ search' = search { searchType = SearchTypeScan }
550+ resp' <- bindM2 dispatchSearch url (return search')
551+ let msr = decode' $ responseBody resp' :: Maybe (SearchResult ())
552+ msid = maybe Nothing scrollId msr
553+ return msid
554+
555+scroll' :: (FromJSON a, MonadBH m) => Maybe ScrollId -> m ([Hit a], Maybe ScrollId)
556+scroll' Nothing = return ([], Nothing)
557+scroll' (Just sid) = do
558+ url <- joinPath ["_search/scroll?scroll=1m"]
559+ resp' <- post url (Just . L.fromStrict $ T.encodeUtf8 sid)
560+ let msr = decode' $ responseBody resp' :: FromJSON a => Maybe (SearchResult a)
561+ resp = case msr of
562+ Just sr -> (hits $ searchHits sr, scrollId sr)
563+ _ -> ([], Nothing)
564+ return resp
565+
566+simpleAccumilator :: (FromJSON a, MonadBH m) => [Hit a] -> ([Hit a], Maybe ScrollId) -> m ([Hit a], Maybe ScrollId)
567+simpleAccumilator oldHits (newHits, Nothing) = return (oldHits ++ newHits, Nothing)
568+simpleAccumilator oldHits ([], _) = return (oldHits, Nothing)
569+simpleAccumilator oldHits (newHits, msid) = do
570+ (newHits', msid') <- scroll' msid
571+ simpleAccumilator (oldHits ++ newHits) (newHits', msid')
572+
573+scanSearch :: (FromJSON a, MonadBH m) => Search -> m [Hit a]
574+scanSearch search = do
575+ msid <- scanSearch' search
576+ (hits, msid') <- scroll' msid
577+ (totalHits, _) <- simpleAccumilator [] (hits, msid')
578+ return totalHits
579+
533580-- | 'mkSearch' is a helper function for defaulting additional fields of a 'Search'
534581-- to Nothing in case you only care about your 'Query' and 'Filter'. Use record update
535582-- syntax if you want to add things like aggregations or highlights while still using
@@ -537,9 +584,9 @@ searchByType (IndexName indexName)
537584--
538585-- >>> let query = TermQuery (Term "user" "bitemyapp") Nothing
539586-- >>> mkSearch (Just query) Nothing
540−-- Search {queryBody = Just (TermQuery (Term {termField = "user", termValue = "bitemyapp"}) Nothing), filterBody = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 0, size = Size 10, fields = Nothing, source = Nothing}
587+-- Search {queryBody = Just (TermQuery (Term {termField = "user", termValue = "bitemyapp"}) Nothing), filterBody = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 0, size = Size 10, searchType = SearchTypeQueryThenFetch, fields = Nothing, source = Nothing}
541588mkSearch :: Maybe Query -> Maybe Filter -> Search
542−mkSearch query filter = Search query filter Nothing Nothing Nothing False (From 0) (Size 10) Nothing Nothing
589+mkSearch query filter = Search query filter Nothing Nothing Nothing False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing
543590
544591-- | 'mkAggregateSearch' is a helper function that defaults everything in a 'Search' except for
545592-- the 'Query' and the 'Aggregation'.
@@ -549,7 +596,7 @@ mkSearch query filter = Search query filter Nothing Nothing Nothing False (From
549596-- TermsAgg (TermsAggregation {term = Left "user", termInclude = Nothing, termExclude = Nothing, termOrder = Nothing, termMinDocCount = Nothing, termSize = Nothing, termShardSize = Nothing, termCollectMode = Just BreadthFirst, termExecutionHint = Nothing, termAggs = Nothing})
550597-- >>> let myAggregation = mkAggregateSearch Nothing $ mkAggregations "users" terms
551598mkAggregateSearch :: Maybe Query -> Aggregations -> Search
552−mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSearchAggs) Nothing False (From 0) (Size 0) Nothing Nothing
599+mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSearchAggs) Nothing False (From 0) (Size 0) SearchTypeQueryThenFetch Nothing Nothing
553600
554601-- | 'mkHighlightSearch' is a helper function that defaults everything in a 'Search' except for
555602-- the 'Query' and the 'Aggregation'.
@@ -558,7 +605,7 @@ mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSear
558605-- >>> let testHighlight = Highlights Nothing [FieldHighlight (FieldName "message") Nothing]
559606-- >>> let search = mkHighlightSearch (Just query) testHighlight
560607mkHighlightSearch :: Maybe Query -> Highlights -> Search
561−mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing (Just searchHighlights) False (From 0) (Size 10) Nothing Nothing
608+mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing (Just searchHighlights) False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing
562609
563610-- | 'pageSearch' is a helper function that takes a search and assigns the from
564611-- and size fields for the search. The from parameter defines the offset
@@ -568,9 +615,9 @@ mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing
568615-- >>> let query = QueryMatchQuery $ mkMatchQuery (FieldName "_all") (QueryString "haskell")
569616-- >>> let search = mkSearch (Just query) Nothing
570617-- >>> search
571−-- Search {queryBody = Just (QueryMatchQuery (MatchQuery {matchQueryField = FieldName "_all", matchQueryQueryString = QueryString "haskell", matchQueryOperator = Or, matchQueryZeroTerms = ZeroTermsNone, matchQueryCutoffFrequency = Nothing, matchQueryMatchType = Nothing, matchQueryAnalyzer = Nothing, matchQueryMaxExpansions = Nothing, matchQueryLenient = Nothing})), filterBody = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 0, size = Size 10, fields = Nothing, source = Nothing}
618+-- Search {queryBody = Just (QueryMatchQuery (MatchQuery {matchQueryField = FieldName "_all", matchQueryQueryString = QueryString "haskell", matchQueryOperator = Or, matchQueryZeroTerms = ZeroTermsNone, matchQueryCutoffFrequency = Nothing, matchQueryMatchType = Nothing, matchQueryAnalyzer = Nothing, matchQueryMaxExpansions = Nothing, matchQueryLenient = Nothing})), filterBody = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 0, size = Size 10, searchType = SearchTypeQueryThenFetch, fields = Nothing, source = Nothing}
572619-- >>> pageSearch (From 10) (Size 100) search
573−-- Search {queryBody = Just (QueryMatchQuery (MatchQuery {matchQueryField = FieldName "_all", matchQueryQueryString = QueryString "haskell", matchQueryOperator = Or, matchQueryZeroTerms = ZeroTermsNone, matchQueryCutoffFrequency = Nothing, matchQueryMatchType = Nothing, matchQueryAnalyzer = Nothing, matchQueryMaxExpansions = Nothing, matchQueryLenient = Nothing})), filterBody = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 10, size = Size 100, fields = Nothing, source = Nothing}
620+-- Search {queryBody = Just (QueryMatchQuery (MatchQuery {matchQueryField = FieldName "_all", matchQueryQueryString = QueryString "haskell", matchQueryOperator = Or, matchQueryZeroTerms = ZeroTermsNone, matchQueryCutoffFrequency = Nothing, matchQueryMatchType = Nothing, matchQueryAnalyzer = Nothing, matchQueryMaxExpansions = Nothing, matchQueryLenient = Nothing})), filterBody = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 10, size = Size 100, searchType = SearchTypeQueryThenFetch, fields = Nothing, source = Nothing}
574621pageSearch :: From -- ^ The result offset
575622 -> Size -- ^ The number of results to return
576623 -> Search -- ^ The current seach