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