a1d2729 · typeable/bloodhound

git.iko.soy/typeable/bloodhound

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

Add support for Search After API

commita1d2729d5c56641d935e65cb15c1fef06a4298e0 authorChris Allen <cma@bitemyapp.com> dateMon 4 Feb 2019, 09:53 (8 years ago) parent0c0d400 13df8a8

6 files changed+45−20

compare with master

unifiedsplit

src/Database/V5/Bloodhound/Client.hs+16−16 src/Database/V5/Bloodhound/Internal/Aggregation.hs+4 src/Database/V5/Bloodhound/Types.hs+5−2 tests/V5/Test/Generators.hs+2−1 tests/V5/Test/Sorting.hs+1−1 tests/V5/tests.hs+17

src/Database/V5/Bloodhound/Client.hs+16−16

@@ -523,13 +523,13 @@ createIndex indexSettings (IndexName indexName) =

523523createIndexWith :: MonadBH m

524524 => [UpdatableIndexSetting]

525525 -> Int -- ^ shard count

526− -> IndexName

526+ -> IndexName

527527 -> m Reply

528528createIndexWith updates shards (IndexName indexName) =

529529 bindM2 put url (return (Just body))

530530 where url = joinPath [indexName]

531531 body = encode $ object

532− ["settings" .= deepMerge

532+ ["settings" .= deepMerge

533533 ( HM.singleton "index.number_of_shards" (toJSON shards) :

534534 [u | Object u <- toJSON <$> updates]

535535 )

@@ -575,8 +575,8 @@ getIndexSettings (IndexName indexName) =

575575 where

576576 url = joinPath [indexName, "_settings"]

577577

578−-- | 'forceMergeIndex'

579−--

578+-- | 'forceMergeIndex'

579+--

580580-- The force merge API allows to force merging of one or more indices through

581581-- an API. The merge relates to the number of segments a Lucene index holds

582582-- within each shard. The force merge operation allows to reduce the number of

@@ -1035,10 +1035,10 @@ searchByType (IndexName indexName)

10351035-- search results. Note that the search is put into 'SearchTypeScan'

10361036-- mode and thus results will not be sorted. Combine this with

10371037-- 'advanceScroll' to efficiently stream through the full result set

1038−getInitialScroll ::

1039− (FromJSON a, MonadThrow m, MonadBH m) => IndexName ->

1040− MappingName ->

1041− Search ->

1038+getInitialScroll ::

1039+ (FromJSON a, MonadThrow m, MonadBH m) => IndexName ->

1040+ MappingName ->

1041+ Search ->

10421042 m (Either EsError (SearchResult a))

10431043getInitialScroll (IndexName indexName) (MappingName mappingName) search' = do

10441044 let url = addQuery params <$> joinPath [indexName, mappingName, "_search"]

@@ -1063,7 +1063,7 @@ getInitialSortedScroll (IndexName indexName) (MappingName mappingName) search =

10631063 resp' <- bindM2 dispatchSearch url (return search)

10641064 parseEsResponse resp'

10651065

1066−scroll' :: (FromJSON a, MonadBH m, MonadThrow m) => Maybe ScrollId ->

1066+scroll' :: (FromJSON a, MonadBH m, MonadThrow m) => Maybe ScrollId ->

10671067 m ([Hit a], Maybe ScrollId)

10681068scroll' Nothing = return ([], Nothing)

10691069scroll' (Just sid) = do

@@ -1090,13 +1090,13 @@ advanceScroll (ScrollId sid) scroll = do

10901090 where scrollTime = showText secs <> "s"

10911091 secs :: Integer

10921092 secs = round scroll

1093−

1093+

10941094 scrollObject = object [ "scroll" .= scrollTime

10951095 , "scroll_id" .= sid

10961096 ]

10971097

1098−simpleAccumulator ::

1099− (FromJSON a, MonadBH m, MonadThrow m) =>

1098+simpleAccumulator ::

1099+ (FromJSON a, MonadBH m, MonadThrow m) =>

11001100 [Hit a] ->

11011101 ([Hit a], Maybe ScrollId) ->

11021102 m ([Hit a], Maybe ScrollId)

@@ -1134,9 +1134,9 @@ scanSearch indexName mappingName search = do

11341134--

11351135-- >>> let query = TermQuery (Term "user" "bitemyapp") Nothing

11361136-- >>> mkSearch (Just query) Nothing

1137−-- 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}

1137+-- 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}

11381138mkSearch :: Maybe Query -> Maybe Filter -> Search

1139−mkSearch query filter = Search query filter Nothing Nothing Nothing False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing

1139+mkSearch query filter = Search query filter Nothing Nothing Nothing False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing Nothing

11401140

11411141-- | 'mkAggregateSearch' is a helper function that defaults everything in a 'Search' except for

11421142-- the 'Query' and the 'Aggregation'.

@@ -1146,7 +1146,7 @@ mkSearch query filter = Search query filter Nothing Nothing Nothing False (From

11461146-- TermsAgg (TermsAggregation {term = Left "user", termInclude = Nothing, termExclude = Nothing, termOrder = Nothing, termMinDocCount = Nothing, termSize = Nothing, termShardSize = Nothing, termCollectMode = Just BreadthFirst, termExecutionHint = Nothing, termAggs = Nothing})

11471147-- >>> let myAggregation = mkAggregateSearch Nothing $ mkAggregations "users" terms

11481148mkAggregateSearch :: Maybe Query -> Aggregations -> Search

1149−mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSearchAggs) Nothing False (From 0) (Size 0) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing

1149+mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSearchAggs) Nothing False (From 0) (Size 0) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing Nothing

11501150

11511151-- | 'mkHighlightSearch' is a helper function that defaults everything in a 'Search' except for

11521152-- the 'Query' and the 'Aggregation'.

@@ -1155,7 +1155,7 @@ mkAggregateSearch query mkSearchAggs = Search query Nothing Nothing (Just mkSear

11551155-- >>> let testHighlight = Highlights Nothing [FieldHighlight (FieldName "message") Nothing]

11561156-- >>> let search = mkHighlightSearch (Just query) testHighlight

11571157mkHighlightSearch :: Maybe Query -> Highlights -> Search

1158−mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing (Just searchHighlights) False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing

1158+mkHighlightSearch query searchHighlights = Search query Nothing Nothing Nothing (Just searchHighlights) False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing Nothing Nothing

11591159

11601160-- | 'pageSearch' is a helper function that takes a search and assigns the from

11611161-- and size fields for the search. The from parameter defines the offset

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

@@ -6,6 +6,7 @@ module Database.V5.Bloodhound.Internal.Aggregation where

77import Bloodhound.Import

9+import qualified Data.Aeson as Aeson

910import qualified Data.HashMap.Strict as HM

1011import qualified Data.Map.Strict as M

1112import qualified Data.Text as T

@@ -445,6 +446,7 @@ instance Monoid (SearchHits a) where

445446 mempty = SearchHits 0 Nothing mempty

446447 mappend = (<>)

447448

449+type SearchAfterKey = [Aeson.Value]

448450

449451data Hit a =

450452 Hit { hitIndex :: IndexName

@@ -452,6 +454,7 @@ data Hit a =

452454 , hitDocId :: DocId

453455 , hitScore :: Score

454456 , hitSource :: Maybe a

457+ , hitSort :: Maybe SearchAfterKey

455458 , hitFields :: Maybe HitFields

456459 , hitHighlight :: Maybe HitHighlight } deriving (Eq, Show)

457460

@@ -462,6 +465,7 @@ instance (FromJSON a) => FromJSON (Hit a) where

462465 v .: "_id" <*>

463466 v .: "_score" <*>

464467 v .:? "_source" <*>

468+ v .:? "sort" <*>

465469 v .:? "fields" <*>

466470 v .:? "highlight"

467471 parseJSON _ = empty

src/Database/V5/Bloodhound/Types.hs+5−2

@@ -382,6 +382,7 @@ module Database.V5.Bloodhound.Types

382382 , DateMathUnit(..)

383383 , TopHitsAggregation(..)

384384 , StatisticsAggregation(..)

385+ , SearchAfterKey

385386

386387 , Highlights(..)

387388 , FieldHighlight(..)

@@ -413,7 +414,7 @@ module Database.V5.Bloodhound.Types

413414 , TokenChar(..)

414415 , Shingle(..)

415416 , Language(..)

416− ) where

417+ ) where

417418

418419import Bloodhound.Import

419420

@@ -444,6 +445,7 @@ data Search = Search { queryBody :: Maybe Query

444445 , from :: From

445446 , size :: Size

446447 , searchType :: SearchType

448+ , searchAfterKey :: Maybe SearchAfterKey

447449 , fields :: Maybe [FieldName]

448450 , scriptFields :: Maybe ScriptFields

449451 , source :: Maybe Source

@@ -453,7 +455,7 @@ data Search = Search { queryBody :: Maybe Query

453455

454456instance ToJSON Search where

455457 toJSON (Search mquery sFilter sort searchAggs

456− highlight sTrackSortScores sFrom sSize _ sFields

458+ highlight sTrackSortScores sFrom sSize _ sAfter sFields

457459 sScriptFields sSource sSuggest) =

458460 omitNulls [ "query" .= query'

459461 , "sort" .= sort

@@ -462,6 +464,7 @@ instance ToJSON Search where

462464 , "from" .= sFrom

463465 , "size" .= sSize

464466 , "track_scores" .= sTrackSortScores

467+ , "search_after" .= sAfter

465468 , "fields" .= sFields

466469 , "script_fields" .= sScriptFields

467470 , "_source" .= sSource

tests/V5/Test/Generators.hs+2−1

@@ -55,6 +55,7 @@ instance (Arbitrary a, Typeable a) => Arbitrary (Hit a) where

5555 <*> arbitrary

5656 <*> arbitraryScore

5757 <*> arbitrary

58+ <*> return Nothing

5859 <*> arbitrary

5960 <*> arbitrary

6061

@@ -183,7 +184,7 @@ instance Arbitrary Query where

183184

184185instance Arbitrary Filter where

185186 arbitrary =

186− Filter <$> arbitrary

187+ Filter <$> arbitrary

187188 shrink (Filter q) =

188189 Filter <$> shrink q

189190

tests/V5/Test/Sorting.hs+1−1

@@ -14,7 +14,7 @@ spec =

1414 let sortSpec = DefaultSortSpec $ mkSort (FieldName "age") Ascending

1515 let search = Search Nothing

1616 Nothing (Just [sortSpec]) Nothing Nothing

17− False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing

17+ False (From 0) (Size 10) SearchTypeQueryThenFetch Nothing Nothing Nothing

1818 Nothing Nothing

1919 result <- searchTweets search

2020 let myTweet = grabFirst result

tests/V5/tests.hs+17

@@ -18,6 +18,7 @@ import Test.Import

1818

1919import Prelude

2020

21+import qualified Data.Aeson as Aeson

2122import qualified Test.Aggregation as Aggregation

2223import qualified Test.BulkAPI as Bulk

2324import qualified Test.Documents as Documents

@@ -117,3 +118,19 @@ main = hspec $ do

117118 regular_search `shouldBe` Right exampleTweet -- Check that the size restrtiction is being honored

118119 liftIO $

119120 scan_search `shouldMatchList` [Just exampleTweet, Just otherTweet]

121+

122+ describe "Search After API" $

123+ it "returns document for search after query" $ withTestEnv $ do

124+ _ <- insertData

125+ _ <- insertOther

126+ let

127+ sortSpec = DefaultSortSpec $ mkSort (FieldName "user") Ascending

128+ searchAfterKey = [Aeson.toJSON ("bitemyapp" :: String)]

129+ search = Search Nothing

130+ Nothing (Just [sortSpec]) Nothing Nothing

131+ False (From 0) (Size 10) SearchTypeQueryThenFetch (Just searchAfterKey)

132+ Nothing Nothing Nothing Nothing

133+ result <- searchTweets search

134+ let myTweet = grabFirst result

135+ liftIO $

136+ myTweet `shouldBe` Right otherTweet