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

Add support for Search After API

src/Database/V5/Bloodhound/Client.hs

@@ -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(..)
@@ -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
 

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