Merge pull request #242 from shulhi/shulhi-search-after

Add support for Search After API

src/Database/V5/Bloodhound/Client.hs

@@ -523,13 +523,13 @@ createIndex indexSettings (IndexName indexName) =
 createIndexWith :: MonadBH m
   => [UpdatableIndexSetting]
   -> Int -- ^ shard count
-  -> IndexName 
+  -> IndexName
   -> m Reply
 createIndexWith updates shards (IndexName indexName) =
   bindM2 put url (return (Just body))
   where url = joinPath [indexName]
         body = encode $ object
-          ["settings" .= deepMerge 
+          ["settings" .= deepMerge
             ( HM.singleton "index.number_of_shards" (toJSON shards) :
               [u | Object u <- toJSON <$> updates]
             )
@@ -575,8 +575,8 @@ getIndexSettings (IndexName indexName) =
   where
     url = joinPath [indexName, "_settings"]
 
--- | 'forceMergeIndex' 
--- 
+-- | 'forceMergeIndex'
+--
 -- The force merge API allows to force merging of one or more indices through
 -- an API. The merge relates to the number of segments a Lucene index holds
 -- within each shard. The force merge operation allows to reduce the number of
@@ -1035,10 +1035,10 @@ searchByType (IndexName indexName)
 -- search results. Note that the search is put into 'SearchTypeScan'
 -- mode and thus results will not be sorted. Combine this with
 -- 'advanceScroll' to efficiently stream through the full result set
-getInitialScroll :: 
-  (FromJSON a, MonadThrow m, MonadBH m) => IndexName -> 
-                                           MappingName -> 
-                                           Search -> 
+getInitialScroll ::
+  (FromJSON a, MonadThrow m, MonadBH m) => IndexName ->
+                                           MappingName ->
+                                           Search ->
                                            m (Either EsError (SearchResult a))
 getInitialScroll (IndexName indexName) (MappingName mappingName) search' = do
     let url = addQuery params <$> joinPath [indexName, mappingName, "_search"]
@@ -1063,7 +1063,7 @@ getInitialSortedScroll (IndexName indexName) (MappingName mappingName) search =
     resp' <- bindM2 dispatchSearch url (return search)
     parseEsResponse resp'
 
-scroll' :: (FromJSON a, MonadBH m, MonadThrow m) => Maybe ScrollId -> 
+scroll' :: (FromJSON a, MonadBH m, MonadThrow m) => Maybe ScrollId ->
                                                     m ([Hit a], Maybe ScrollId)
 scroll' Nothing = return ([], Nothing)
 scroll' (Just sid) = do
@@ -1090,13 +1090,13 @@ advanceScroll (ScrollId sid) scroll = do
   where scrollTime = showText secs <> "s"
         secs :: Integer
         secs = round scroll
-        
+
         scrollObject = object [ "scroll" .= scrollTime
                               , "scroll_id" .= sid
                               ]
 
-simpleAccumulator :: 
-  (FromJSON a, MonadBH m, MonadThrow m) => 
+simpleAccumulator ::
+  (FromJSON a, MonadBH m, MonadThrow m) =>
                                 [Hit a] ->
                                 ([Hit a], Maybe ScrollId) ->
                                 m ([Hit a], Maybe ScrollId)
@@ -1134,9 +1134,9 @@ scanSearch indexName mappingName search = do
 --
 -- >>> let query = TermQuery (Term "user" "bitemyapp") Nothing
 -- >>> mkSearch (Just query) Nothing
--- 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}
+-- Search {queryBody = Just (TermQuery (Term {termField = "user", termValue = "bitemyapp"}) Nothing), filterBody = Nothing, searchAfterKey = Nothing, sortBody = Nothing, aggBody = Nothing, highlight = Nothing, trackSortScores = False, from = From 0, size = Size 10, searchType = SearchTypeQueryThenFetch, fields = Nothing, source = Nothing}
 mkSearch :: Maybe Query -> Maybe Filter -> Search
-mkSearch query filter = Search query filter Nothing Nothing Nothing False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing
+mkSearch query filter = Search query filter Nothing Nothing Nothing False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing Nothing
 
 -- | 'mkAggregateSearch' is a helper function that defaults everything in a 'Search' except for
 --   the 'Query' and the 'Aggregation'.
@@ -1146,7 +1146,7 @@ mkSearch query filter = Search query filter Nothing Nothing Nothing False (From
 -- TermsAgg (TermsAggregation {term = Left "user", termInclude = Nothing, termExclude = Nothing, termOrder = Nothing, termMinDocCount = Nothing, termSize = Nothing, termShardSize = Nothing, termCollectMode = Just BreadthFirst, termExecutionHint = Nothing, termAggs = Nothing})
 -- >>> let myAggregation = mkAggregateSearch Nothing $ mkAggregations "users" terms
 mkAggregateSearch :: Maybe Query -> Aggregations -> Search
-mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSearchAggs) Nothing False (From 0) (Size 0) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing
+mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSearchAggs) Nothing False (From 0) (Size 0) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing Nothing
 
 -- | 'mkHighlightSearch' is a helper function that defaults everything in a 'Search' except for
 --   the 'Query' and the 'Aggregation'.
@@ -1155,7 +1155,7 @@ mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSear
 -- >>> let testHighlight = Highlights Nothing [FieldHighlight (FieldName "message") Nothing]
 -- >>> let search = mkHighlightSearch (Just query) testHighlight
 mkHighlightSearch :: Maybe Query -> Highlights -> Search
-mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing (Just searchHighlights) False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing
+mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing (Just searchHighlights) False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing Nothing
 
 -- | 'pageSearch' is a helper function that takes a search and assigns the from
 --    and size fields for the search. The from parameter defines the offset

src/Database/V5/Bloodhound/Internal/Aggregation.hs

@@ -6,6 +6,7 @@ module Database.V5.Bloodhound.Internal.Aggregation where
 
 import           Bloodhound.Import
 
+import qualified Data.Aeson as Aeson
 import qualified Data.HashMap.Strict as HM
 import qualified Data.Map.Strict as M
 import qualified Data.Text as T
@@ -445,6 +446,7 @@ instance Monoid (SearchHits a) where
   mempty = SearchHits 0 Nothing mempty
   mappend = (<>)
 
+type SearchAfterKey = [Aeson.Value]
 
 data Hit a =
   Hit { hitIndex     :: IndexName
@@ -452,6 +454,7 @@ data Hit a =
       , hitDocId     :: DocId
       , hitScore     :: Score
       , hitSource    :: Maybe a
+      , hitSort      :: Maybe SearchAfterKey
       , hitFields    :: Maybe HitFields
       , hitHighlight :: Maybe HitHighlight } deriving (Eq, Show)
 
@@ -462,6 +465,7 @@ instance (FromJSON a) => FromJSON (Hit a) where
                          v .:  "_id"      <*>
                          v .:  "_score"   <*>
                          v .:? "_source"  <*>
+                         v .:? "sort"     <*>
                          v .:? "fields"   <*>
                          v .:? "highlight"
   parseJSON _          = empty

src/Database/V5/Bloodhound/Types.hs

@@ -382,6 +382,7 @@ module Database.V5.Bloodhound.Types
        , DateMathUnit(..)
        , TopHitsAggregation(..)
        , StatisticsAggregation(..)
+       , SearchAfterKey
 
        , Highlights(..)
        , FieldHighlight(..)
@@ -413,7 +414,7 @@ module Database.V5.Bloodhound.Types
        , TokenChar(..)
        , Shingle(..)
        , Language(..)
-         ) where
+       ) where
 
 import           Bloodhound.Import
 
@@ -444,6 +445,7 @@ data Search = Search { queryBody       :: Maybe Query
                      , from            :: From
                      , size            :: Size
                      , searchType      :: SearchType
+                     , searchAfterKey  :: Maybe SearchAfterKey
                      , fields          :: Maybe [FieldName]
                      , scriptFields    :: Maybe ScriptFields
                      , source          :: Maybe Source
@@ -453,7 +455,7 @@ data Search = Search { queryBody       :: Maybe Query
 
 instance ToJSON Search where
   toJSON (Search mquery sFilter sort searchAggs
-          highlight sTrackSortScores sFrom sSize _ sFields
+          highlight sTrackSortScores sFrom sSize _ sAfter sFields
           sScriptFields sSource sSuggest) =
     omitNulls [ "query"         .= query'
               , "sort"          .= sort
@@ -462,6 +464,7 @@ instance ToJSON Search where
               , "from"          .= sFrom
               , "size"          .= sSize
               , "track_scores"  .= sTrackSortScores
+              , "search_after"  .= sAfter
               , "fields"        .= sFields
               , "script_fields" .= sScriptFields
               , "_source"       .= sSource

tests/V5/Test/Generators.hs

@@ -55,6 +55,7 @@ instance (Arbitrary a, Typeable a) => Arbitrary (Hit a) where
                   <*> arbitrary
                   <*> arbitraryScore
                   <*> arbitrary
+                  <*> return Nothing
                   <*> arbitrary
                   <*> arbitrary
 
@@ -183,7 +184,7 @@ instance Arbitrary Query where
 
 instance Arbitrary Filter where
   arbitrary =
-    Filter <$> arbitrary 
+    Filter <$> arbitrary
   shrink (Filter q) =
     Filter <$> shrink q
 

tests/V5/Test/Sorting.hs

@@ -14,7 +14,7 @@ spec =
       let sortSpec = DefaultSortSpec $ mkSort (FieldName "age") Ascending
       let search = Search Nothing
                    Nothing (Just [sortSpec]) Nothing Nothing
-                   False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing
+                   False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing
                    Nothing Nothing
       result <- searchTweets search
       let myTweet = grabFirst result

tests/V5/tests.hs

@@ -18,6 +18,7 @@ import           Test.Import
 
 import           Prelude
 
+import qualified Data.Aeson as Aeson
 import qualified Test.Aggregation as Aggregation
 import qualified Test.BulkAPI as Bulk
 import qualified Test.Documents as Documents
@@ -117,3 +118,19 @@ main = hspec $ do
         regular_search `shouldBe` Right exampleTweet -- Check that the size restrtiction is being honored
       liftIO $
         scan_search `shouldMatchList` [Just exampleTweet, Just otherTweet]
+
+  describe "Search After API" $
+    it "returns document for search after query" $ withTestEnv $ do
+      _ <- insertData
+      _ <- insertOther
+      let
+        sortSpec = DefaultSortSpec $ mkSort (FieldName "user") Ascending
+        searchAfterKey = [Aeson.toJSON ("bitemyapp" :: String)]
+        search = Search Nothing
+                 Nothing (Just [sortSpec]) Nothing Nothing
+                 False (From 0) (Size 10) SearchTypeQueryThenFetch (Just searchAfterKey)
+                 Nothing Nothing Nothing Nothing
+      result <- searchTweets search
+      let myTweet = grabFirst result
+      liftIO $
+        myTweet `shouldBe` Right otherTweet