@@ -42,6 +42,8 @@ module Database.Bloodhound.Client
4242 , searchByIndex
4343 , searchByType
4444 , scanSearch
45+ , getInitialScroll
46+ , advanceScroll
4547 , refreshIndex
4648 , mkSearch
4749 , mkAggregateSearch
@@ -78,6 +80,7 @@ import Data.Monoid
7880import Data.Text (Text)
7981import qualified Data.Text as T
8082import qualified Data.Text.Encoding as T
83+import Data.Time.Clock
8184import qualified Data.Vector as V
8285import Network.HTTP.Client
8386import qualified Network.HTTP.Types.Method as NHTM
@@ -289,7 +292,7 @@ existentialQuery url = do
289292-- failing that tries to parse it as an EsError. All well-formed, JSON
290293-- responses from elasticsearch should fall into these two
291294-- categories. If they don't, a 'StatusCodeException' will be thrown.
292−parseEsResponse :: (MonadBH m, MonadThrow m, FromJSON a) => Reply
295+parseEsResponse :: (MonadThrow m, FromJSON a) => Reply
293296 -> m (Either EsError a)
294297parseEsResponse reply
295298 | respIsTwoHunna reply = case eitherDecode body of
@@ -608,8 +611,12 @@ searchByType (IndexName indexName)
608611 (MappingName mappingName) = bindM2 dispatchSearch url . return
609612 where url = joinPath [indexName, mappingName, "_search"]
610613
611−scanSearch' :: MonadBH m => IndexName -> MappingName -> Search -> m (Maybe ScrollId)
612−scanSearch' (IndexName indexName) (MappingName mappingName) search = do
614+-- | For a given scearch, request a scroll for efficient streaming of
615+-- search results. Note that the search is put into 'SearchTypeScan'
616+-- mode and thus results will not be sorted. Combine this with
617+-- 'advanceScroll' to efficiently stream through the full result set
618+getInitialScroll :: MonadBH m => IndexName -> MappingName -> Search -> m (Maybe ScrollId)
619+getInitialScroll (IndexName indexName) (MappingName mappingName) search = do
613620 let url = joinPath [indexName, mappingName, "_search"]
614621 search' = search { searchType = SearchTypeScan }
615622 resp' <- bindM2 dispatchSearch url (return search')
@@ -617,31 +624,54 @@ scanSearch' (IndexName indexName) (MappingName mappingName) search = do
617624 msid = maybe Nothing scrollId msr
618625 return msid
619626
620−scroll' :: (FromJSON a, MonadBH m) => Maybe ScrollId -> m ([Hit a], Maybe ScrollId)
627+scroll' :: (FromJSON a, MonadBH m, MonadThrow m) => Maybe ScrollId -> m ([Hit a], Maybe ScrollId)
621628scroll' Nothing = return ([], Nothing)
622629scroll' (Just sid) = do
623− url <- joinPath ["_search/scroll?scroll=1m"]
624− resp' <- post url (Just . L.fromStrict $ T.encodeUtf8 sid)
625− let msr = decode' $ responseBody resp' :: FromJSON a => Maybe (SearchResult a)
626− resp = case msr of
627− Just sr -> (hits $ searchHits sr, scrollId sr)
628− _ -> ([], Nothing)
629− return resp
630−
631−simpleAccumilator :: (FromJSON a, MonadBH m) => [Hit a] -> ([Hit a], Maybe ScrollId) -> m ([Hit a], Maybe ScrollId)
632−simpleAccumilator oldHits (newHits, Nothing) = return (oldHits ++ newHits, Nothing)
633−simpleAccumilator oldHits ([], _) = return (oldHits, Nothing)
634−simpleAccumilator oldHits (newHits, msid) = do
630+ res <- advanceScroll sid 60
631+ case res of
632+ Right SearchResult {..} -> return (hits searchHits, scrollId)
633+ Left _ -> return ([], Nothing)
634+
635+-- | Use the given scroll to fetch the next page of documents. If
636+-- there are still further pages, there will be a value in the
637+-- 'scrollId' field of the 'SearchResult'
638+advanceScroll
639+ :: ( FromJSON a
640+ , MonadBH m
641+ , MonadThrow m
642+ )
643+ => ScrollId
644+ -> NominalDiffTime
645+ -- ^ How long should the snapshot of data be kept around? This timeout is updated every time 'advanceScroll' is used, so don't feel the need to set it to the entire duration of your search processing. Note that durations < 1s will be rounded up. Also note that 'NominalDiffTime' is an instance of Num so literals like 60 will be interpreted as seconds. 60s is a reasonable default.
646+ -> m (Either EsError (SearchResult a))
647+advanceScroll (ScrollId sid) scroll = do
648+ url <- joinPath ["_search/scroll?scroll=" <> scrollTime]
649+ parseEsResponse =<< post url (Just . L.fromStrict $ T.encodeUtf8 sid)
650+ where scrollTime = showText secs <> "s"
651+ secs :: Integer
652+ secs = round scroll
653+
654+simpleAccumulator :: (FromJSON a, MonadBH m, MonadThrow m) => [Hit a] -> ([Hit a], Maybe ScrollId) -> m ([Hit a], Maybe ScrollId)
655+simpleAccumulator oldHits (newHits, Nothing) = return (oldHits ++ newHits, Nothing)
656+simpleAccumulator oldHits ([], _) = return (oldHits, Nothing)
657+simpleAccumulator oldHits (newHits, msid) = do
635658 (newHits', msid') <- scroll' msid
636− simpleAccumilator (oldHits ++ newHits) (newHits', msid')
659+ simpleAccumulator (oldHits ++ newHits) (newHits', msid')
637660
638661-- | 'scanSearch' uses the 'scan&scroll' API of elastic,
639−-- for a given 'IndexName' and 'MappingName',
640−scanSearch :: (FromJSON a, MonadBH m) => IndexName -> MappingName -> Search -> m [Hit a]
662+-- for a given 'IndexName' and 'MappingName'. Note that this will
663+-- consume the entire search result set and will be doing O(n) list
664+-- appends so this may not be suitable for large result sets. In that
665+-- case, 'getInitialScroll' and 'advanceScroll' are good low level
666+-- tools. You should be able to hook them up trivially to conduit,
667+-- pipes, or your favorite streaming IO abstraction of choice. Note
668+-- that ordering on the search would destroy performance and thus is
669+-- ignored.
670+scanSearch :: (FromJSON a, MonadBH m, MonadThrow m) => IndexName -> MappingName -> Search -> m [Hit a]
641671scanSearch indexName mappingName search = do
642− msid <- scanSearch' indexName mappingName search
672+ msid <- getInitialScroll indexName mappingName search
643673 (hits, msid') <- scroll' msid
644− (totalHits, _) <- simpleAccumilator [] (hits, msid')
674+ (totalHits, _) <- simpleAccumulator [] (hits, msid')
645675 return totalHits
646676
647677-- | 'mkSearch' is a helper function for defaulting additional fields of a 'Search'