From 4967c628435e410371ff56de3774b19ada2b1fec Mon Sep 17 00:00:00 2001 From: Shulhi Sapli Date: Tue, 31 Jul 2018 14:53:33 +0800 Subject: [PATCH 1/2] Add search after functionality --- src/Database/V5/Bloodhound/Client.hs | 32 +++++++++---------- .../V5/Bloodhound/Internal/Aggregation.hs | 4 +++ src/Database/V5/Bloodhound/Types.hs | 7 ++-- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/src/Database/V5/Bloodhound/Client.hs b/src/Database/V5/Bloodhound/Client.hs index 1613f8e..181a2e9 100644 --- a/src/Database/V5/Bloodhound/Client.hs +++ b/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 diff --git a/src/Database/V5/Bloodhound/Internal/Aggregation.hs b/src/Database/V5/Bloodhound/Internal/Aggregation.hs index f6fc1ab..7265915 100644 --- a/src/Database/V5/Bloodhound/Internal/Aggregation.hs +++ b/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 diff --git a/src/Database/V5/Bloodhound/Types.hs b/src/Database/V5/Bloodhound/Types.hs index 5882797..266c022 100644 --- a/src/Database/V5/Bloodhound/Types.hs +++ b/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 From 13df8a8674cc3b8bd7d4d73362826b3db8344679 Mon Sep 17 00:00:00 2001 From: Shulhi Sapli Date: Tue, 31 Jul 2018 15:36:50 +0800 Subject: [PATCH 2/2] Add test for search after API --- tests/V5/Test/Generators.hs | 3 ++- tests/V5/Test/Sorting.hs | 2 +- tests/V5/tests.hs | 17 +++++++++++++++++ 3 files changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/V5/Test/Generators.hs b/tests/V5/Test/Generators.hs index facf42d..a3b9ef0 100644 --- a/tests/V5/Test/Generators.hs +++ b/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 diff --git a/tests/V5/Test/Sorting.hs b/tests/V5/Test/Sorting.hs index 8665368..b9f838c 100644 --- a/tests/V5/Test/Sorting.hs +++ b/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 diff --git a/tests/V5/tests.hs b/tests/V5/tests.hs index 600bca5..e99c4be 100644 --- a/tests/V5/tests.hs +++ b/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