1 {-# LANGUAGE OverloadedStrings, RecordWildCards #-} 2 module Databrary.Solr.Search 3 ( SearchType(..) 4 , SearchQuery(..) 5 , search 6 ) where 7 8 import Control.Arrow (first) 9 import qualified Data.ByteString.Builder as B 10 import qualified Data.ByteString.Char8 as BSC 11 import qualified Data.ByteString.Lazy as BSL 12 import Data.Char (isAlphaNum) 13 import Data.Monoid ((<>)) 14 import qualified Data.Text as T 15 import qualified Data.Text.Encoding as TE 16 import qualified Network.HTTP.Client as HC 17 import Network.HTTP.Types.URI (renderSimpleQuery) 18 19 import Data.ByteString.Builder.Escape (escapeLazyByteStringCharsWith, escapeTextWith) 20 import Databrary.Has 21 import Databrary.Model.Paginate 22 import Databrary.Model.Metric.Types 23 import Databrary.Solr.Service 24 import Databrary.Solr.Document 25 26 data SearchType 27 = SearchVolumes 28 | SearchParties 29 deriving (Eq) 30 31 data SearchQuery = SearchQuery 32 { searchString :: Maybe T.Text 33 , searchFields :: [(T.Text, T.Text)] 34 , searchMetrics :: [(Metric, T.Text)] 35 , searchType :: SearchType 36 , searchPaginate :: !Paginate 37 } 38 39 checkTerm :: T.Text -> Bool 40 checkTerm = cq False [] . T.unpack where 41 cq False [] "" = True 42 cq q g ('"':s) = cq (not q) g s 43 cq q g ('\\':_:s) = cq q g s 44 cq _ _ ['\\'] = False 45 cq False ('(':gs) (')':s) = cq False gs s 46 cq False ('[':gs) (']':s) = cq False gs s 47 cq False ('[':gs) ('}':s) = cq False gs s 48 cq False ('{':gs) (']':s) = cq False gs s 49 cq False ('{':gs) ('}':s) = cq False gs s 50 cq False g (c:s) 51 | c `elem` ['(','[','{'] = cq False (c:g) s 52 | c `elem` [')',']','}'] = False 53 cq q g (_:s) = cq q g s 54 cq _ _ _ = False 55 56 checkField :: T.Text -> Bool 57 checkField = T.all cc where 58 cc '_' = True 59 cc c = isAlphaNum c 60 61 quoteQuery :: (Char -> String -> a -> B.Builder) -> a -> B.Builder 62 quoteQuery e s = B.char8 '"' <> e '\\' "\"\\" s <> B.char8 '"' 63 64 defaultParams :: B.Builder 65 defaultParams = B.string8 "{!dismax qf=\"text_en^0.5 text_gen^0.5 keyword^2 tag_name party_name\" pf=\"keyword^2 tag_name party_name\" ps=3}" 66 67 search :: MonadSolr c m => SearchQuery -> m (HC.Response BSL.ByteString) 68 search SearchQuery{..} = do 69 req <- peeks solrRequest 70 focusIO $ HC.httpLbs req 71 { HC.path = HC.path req <> "search" 72 , HC.queryString = renderSimpleQuery True query 73 , HC.checkResponse = \_ _ -> putStrLn "Search Error" 74 } 75 where 76 query = 77 [ ("q", BSL.toStrict $ B.toLazyByteString $ qp <> uw ql) 78 , ("fq", "content_type:" <> ct) 79 , ("start", BSC.pack $ show $ paginateOffset searchPaginate) 80 , ("rows", BSC.pack $ show $ paginateLimit searchPaginate) 81 , ("q.op", "AND") 82 ] 83 ++ maybe [] (\q -> 84 [ ("spellcheck", "true") 85 , ("spellcheck.q", TE.encodeUtf8 q) 86 ]) searchString 87 (ct, qp, qe) = case searchType of 88 SearchVolumes -> ("volume", mempty, B.string8 "{!join from=volume_id to=volume_id}") 89 SearchParties -> ("party", mempty, mempty) 90 ql = maybe id ((:) . bp . (defaultParams <>) . TE.encodeUtf8Builder) searchString $ 91 map bt (searchFields ++ map (first metricField) searchMetrics) 92 bt (f, v) 93 | checkField f && checkTerm v = bp (TE.encodeUtf8Builder f <> B.string8 ":(" <> (if T.null v then B.char8 '*' else TE.encodeUtf8Builder v) <> B.char8 ')') 94 | otherwise = bp (B.string8 "{!dismax qf=" <> quoteQuery escapeTextWith f <> B.char8 '}' <> TE.encodeUtf8Builder v) 95 bp v = B.string8 "_query_:" <> quoteQuery escapeLazyByteStringCharsWith (B.toLazyByteString $ qe <> v) 96 uw [] = B.string8 "*:*" 97 uw (t:l) = t <> foldMap (B.char8 ' ' <>) l