{-# LANGUAGE CPP #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}

module Database.Esqueleto.TextSearch.Language
  ( (@@.)
  , prefixAndQuery
  , toSearchTerm
  , SearchTerm
  , to_tsvector
  , to_tsquery
  , plainto_tsquery
  , ts_rank
  , ts_rank_cd
  , setweight
  ) where

import Data.String (IsString)
import Data.Text (Text)
import Database.Esqueleto (SqlExpr, Value, val)
#if MIN_VERSION_esqueleto(3,5,0)
import Database.Esqueleto.Internal.Internal (unsafeSqlBinOp, unsafeSqlFunction)
#else
import Database.Esqueleto.Internal.Sql (unsafeSqlBinOp, unsafeSqlFunction)
#endif
import Database.Esqueleto.TextSearch.Types
import qualified Data.Text as T
import Data.List.NonEmpty(nonEmpty, NonEmpty, toList)


-- | Apply some query to a tsvector document
--   for example:
--
-- @
-- searchCompany :: SqlExpr (Entity CompanySearchIndex) -> SearchTerm -> SqlQuery ()
-- searchCompany company term = do
--   let query = prefixAndQuery term
--       norm = val []
--   where_ $ (company ^. CompanySearchIndexDocument) @@. query
-- @
--
(@@.)
  :: SqlExpr (Value TsVector) -- ^ the document to search in
  -> SqlExpr (Value (TsQuery Lexemes)) -- ^ the query made by 'prefixAndQuery'
  -> SqlExpr (Value Bool)
@@. :: SqlExpr (Value TsVector)
-> SqlExpr (Value (TsQuery Lexemes)) -> SqlExpr (Value Bool)
(@@.) = Builder
-> SqlExpr (Value TsVector)
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value Bool)
forall a b c.
Builder
-> SqlExpr (Value a) -> SqlExpr (Value b) -> SqlExpr (Value c)
unsafeSqlBinOp Builder
"@@"

to_tsvector
  :: IsString a
  => SqlExpr (Value RegConfig)
  -> SqlExpr (Value a)
  -> SqlExpr (Value TsVector)
to_tsvector :: forall a.
IsString a =>
SqlExpr (Value RegConfig)
-> SqlExpr (Value a) -> SqlExpr (Value TsVector)
to_tsvector SqlExpr (Value RegConfig)
a SqlExpr (Value a)
b = Builder
-> (SqlExpr (Value RegConfig), SqlExpr (Value a))
-> SqlExpr (Value TsVector)
forall a b.
UnsafeSqlFunctionArgument a =>
Builder -> a -> SqlExpr (Value b)
unsafeSqlFunction Builder
"to_tsvector" (SqlExpr (Value RegConfig)
a, SqlExpr (Value a)
b)

to_tsquery
  :: SqlExpr (Value RegConfig)
  -> SqlExpr (Value (TsQuery Words))
  -> SqlExpr (Value (TsQuery Lexemes) )
to_tsquery :: SqlExpr (Value RegConfig)
-> SqlExpr (Value (TsQuery Words))
-> SqlExpr (Value (TsQuery Lexemes))
to_tsquery SqlExpr (Value RegConfig)
a SqlExpr (Value (TsQuery Words))
b = Builder
-> (SqlExpr (Value RegConfig), SqlExpr (Value (TsQuery Words)))
-> SqlExpr (Value (TsQuery Lexemes))
forall a b.
UnsafeSqlFunctionArgument a =>
Builder -> a -> SqlExpr (Value b)
unsafeSqlFunction Builder
"to_tsquery" (SqlExpr (Value RegConfig)
a, SqlExpr (Value (TsQuery Words))
b)

plainto_tsquery
  :: SqlExpr (Value RegConfig)
  -> SqlExpr (Value Text)
  -> SqlExpr (Value (TsQuery Lexemes))
plainto_tsquery :: SqlExpr (Value RegConfig)
-> SqlExpr (Value Text) -> SqlExpr (Value (TsQuery Lexemes))
plainto_tsquery SqlExpr (Value RegConfig)
a SqlExpr (Value Text)
b = Builder
-> (SqlExpr (Value RegConfig), SqlExpr (Value Text))
-> SqlExpr (Value (TsQuery Lexemes))
forall a b.
UnsafeSqlFunctionArgument a =>
Builder -> a -> SqlExpr (Value b)
unsafeSqlFunction Builder
"plainto_tsquery" (SqlExpr (Value RegConfig)
a, SqlExpr (Value Text)
b)

-- | Organize search result by weights. This allows you to put better
--   matching results higher.
--   for example:
--
-- @
-- searchCompany :: SqlExpr (Entity CompanySearchIndex) -> SearchTerm -> SqlQuery ()
-- searchCompany company term = do
--   let query = prefixAndQuery term
--       norm = val []
--   where_ $ (company ^. CompanySearchIndexDocument) @@. query
--   orderBy [desc (ts_rank (val defaultWeights)
--                  (company ^. CompanySearchIndexDocument)
--                  query norm)]
-- @
--
ts_rank
  :: SqlExpr (Value Weights) -- ^ relative weighting of a b c and d, see 'defaultWeights'
  -> SqlExpr (Value TsVector) -- ^ the document to search in
  -> SqlExpr (Value (TsQuery Lexemes)) -- ^ the query made by 'prefixAndQuery'
  -> SqlExpr (Value [NormalizationOption]) -- ^ normalization option to indicate how to deal with document length
  -> SqlExpr (Value Double)
ts_rank :: SqlExpr (Value Weights)
-> SqlExpr (Value TsVector)
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value [NormalizationOption])
-> SqlExpr (Value Double)
ts_rank SqlExpr (Value Weights)
a SqlExpr (Value TsVector)
b SqlExpr (Value (TsQuery Lexemes))
c SqlExpr (Value [NormalizationOption])
d = Builder
-> (SqlExpr (Value Weights), SqlExpr (Value TsVector),
    SqlExpr (Value (TsQuery Lexemes)),
    SqlExpr (Value [NormalizationOption]))
-> SqlExpr (Value Double)
forall a b.
UnsafeSqlFunctionArgument a =>
Builder -> a -> SqlExpr (Value b)
unsafeSqlFunction Builder
"ts_rank" (SqlExpr (Value Weights)
a, SqlExpr (Value TsVector)
b, SqlExpr (Value (TsQuery Lexemes))
c, SqlExpr (Value [NormalizationOption])
d)

ts_rank_cd
  :: SqlExpr (Value Weights)
  -> SqlExpr (Value TsVector)
  -> SqlExpr (Value (TsQuery Lexemes))
  -> SqlExpr (Value [NormalizationOption])
  -> SqlExpr (Value Double)
ts_rank_cd :: SqlExpr (Value Weights)
-> SqlExpr (Value TsVector)
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value [NormalizationOption])
-> SqlExpr (Value Double)
ts_rank_cd SqlExpr (Value Weights)
a SqlExpr (Value TsVector)
b SqlExpr (Value (TsQuery Lexemes))
c SqlExpr (Value [NormalizationOption])
d = Builder
-> (SqlExpr (Value Weights), SqlExpr (Value TsVector),
    SqlExpr (Value (TsQuery Lexemes)),
    SqlExpr (Value [NormalizationOption]))
-> SqlExpr (Value Double)
forall a b.
UnsafeSqlFunctionArgument a =>
Builder -> a -> SqlExpr (Value b)
unsafeSqlFunction Builder
"ts_rank_cd" (SqlExpr (Value Weights)
a, SqlExpr (Value TsVector)
b, SqlExpr (Value (TsQuery Lexemes))
c, SqlExpr (Value [NormalizationOption])
d)

setweight
  :: SqlExpr (Value TsVector)
  -> SqlExpr (Value Weight)
  -> SqlExpr (Value TsVector)
setweight :: SqlExpr (Value TsVector)
-> SqlExpr (Value Weight) -> SqlExpr (Value TsVector)
setweight SqlExpr (Value TsVector)
a SqlExpr (Value Weight)
b = Builder
-> (SqlExpr (Value TsVector), SqlExpr (Value Weight))
-> SqlExpr (Value TsVector)
forall a b.
UnsafeSqlFunctionArgument a =>
Builder -> a -> SqlExpr (Value b)
unsafeSqlFunction Builder
"setweight" (SqlExpr (Value TsVector)
a, SqlExpr (Value Weight)
b)

-- | (&&) for tsquery. This function would be called (&&.) but
-- Esqueleto's (&&.) confines that fn to sql boolean expressions.
-- x::tsquery && y::tsquery == to_tsquery('x & y')
tsquery_and :: SqlExpr (Value (TsQuery Lexemes))
      -> SqlExpr (Value (TsQuery Lexemes))
      -> SqlExpr (Value (TsQuery Lexemes))
tsquery_and :: SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value (TsQuery Lexemes))
tsquery_and = Builder
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value (TsQuery Lexemes))
forall a b c.
Builder
-> SqlExpr (Value a) -> SqlExpr (Value b) -> SqlExpr (Value c)
unsafeSqlBinOp Builder
"&&"

-- | format the query into lexemes
--   the result can be used in '@@.' for example:
--
-- @
-- searchCompany :: SqlExpr (Entity CompanySearchIndex) -> SearchTerm -> SqlQuery ()
-- searchCompany company term = do
--   let query = prefixAndQuery term
--       norm = val []
--   where_ $ (company ^. CompanySearchIndexDocument) @@. query
-- @
--
prefixAndQuery :: SearchTerm -> SqlExpr (Value (TsQuery Lexemes))
prefixAndQuery :: SearchTerm -> SqlExpr (Value (TsQuery Lexemes))
prefixAndQuery = RegConfig -> SearchTerm -> SqlExpr (Value (TsQuery Lexemes))
prefixAndQueryLang RegConfig
"english"

-- | specify a language to be used with the query.
prefixAndQueryLang :: RegConfig -> SearchTerm -> SqlExpr (Value (TsQuery Lexemes))
prefixAndQueryLang :: RegConfig -> SearchTerm -> SqlExpr (Value (TsQuery Lexemes))
prefixAndQueryLang RegConfig
language (SearchTerm NonEmpty Text
ts) =
  (SqlExpr (Value (TsQuery Lexemes))
 -> SqlExpr (Value (TsQuery Lexemes))
 -> SqlExpr (Value (TsQuery Lexemes)))
-> [SqlExpr (Value (TsQuery Lexemes))]
-> SqlExpr (Value (TsQuery Lexemes))
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value (TsQuery Lexemes))
-> SqlExpr (Value (TsQuery Lexemes))
tsquery_and
  ([SqlExpr (Value (TsQuery Lexemes))]
 -> SqlExpr (Value (TsQuery Lexemes)))
-> [SqlExpr (Value (TsQuery Lexemes))]
-> SqlExpr (Value (TsQuery Lexemes))
forall a b. (a -> b) -> a -> b
$ (Text -> SqlExpr (Value (TsQuery Lexemes)))
-> [Text] -> [SqlExpr (Value (TsQuery Lexemes))]
forall a b. (a -> b) -> [a] -> [b]
map (SqlExpr (Value RegConfig)
-> SqlExpr (Value (TsQuery Words))
-> SqlExpr (Value (TsQuery Lexemes))
to_tsquery (RegConfig -> SqlExpr (Value RegConfig)
forall typ. PersistField typ => typ -> SqlExpr (Value typ)
val RegConfig
language) (SqlExpr (Value (TsQuery Words))
 -> SqlExpr (Value (TsQuery Lexemes)))
-> (Text -> SqlExpr (Value (TsQuery Words)))
-> Text
-> SqlExpr (Value (TsQuery Lexemes))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TsQuery Words -> SqlExpr (Value (TsQuery Words))
forall typ. PersistField typ => typ -> SqlExpr (Value typ)
val (TsQuery Words -> SqlExpr (Value (TsQuery Words)))
-> (Text -> TsQuery Words)
-> Text
-> SqlExpr (Value (TsQuery Words))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Position -> [Weight] -> Text -> TsQuery Words
Word Position
Prefix []) ([Text] -> [SqlExpr (Value (TsQuery Lexemes))])
-> [Text] -> [SqlExpr (Value (TsQuery Lexemes))]
forall a b. (a -> b) -> a -> b
$ NonEmpty Text -> [Text]
forall a. NonEmpty a -> [a]
toList NonEmpty Text
ts


-- | A valid search term.
--   created with 'toSearchTerm'.
newtype SearchTerm = SearchTerm { SearchTerm -> NonEmpty Text
unQuery :: NonEmpty Text }
  deriving (Int -> SearchTerm -> ShowS
[SearchTerm] -> ShowS
SearchTerm -> String
(Int -> SearchTerm -> ShowS)
-> (SearchTerm -> String)
-> ([SearchTerm] -> ShowS)
-> Show SearchTerm
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SearchTerm -> ShowS
showsPrec :: Int -> SearchTerm -> ShowS
$cshow :: SearchTerm -> String
show :: SearchTerm -> String
$cshowList :: [SearchTerm] -> ShowS
showList :: [SearchTerm] -> ShowS
Show)

-- | Constructs a valid search query, removes a bunch of illegal
--   characters and splits the terms for better results.
--   Also checks if there is anything in the search term.
--
--   using a search term is optional, but it's probably what you want.
--   all underlying primitives are exposed.
toSearchTerm :: Text -> Maybe SearchTerm
toSearchTerm :: Text -> Maybe SearchTerm
toSearchTerm Text
q = NonEmpty Text -> SearchTerm
SearchTerm (NonEmpty Text -> SearchTerm)
-> Maybe (NonEmpty Text) -> Maybe SearchTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text] -> Maybe (NonEmpty Text)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [Text]
qs
  -- We disallow whitespace, \ and ' for the sake of producing a Text
  -- that can fit postgresql's requirements for to_tsquery's text
  -- argument. Note that this is not done nor needed for security reasons
  where qs :: [Text]
qs = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Text -> Bool) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Bool
T.null) ([Text] -> [Text]) -> [Text] -> [Text]
forall a b. (a -> b) -> a -> b
$ Text -> [Text]
T.words
             (Text -> [Text]) -> Text -> [Text]
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.filter (Char -> String -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Char
'\\', Char
'\'']) (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Text -> Text
T.strip Text
q