module Hpgsql.QueryInternal
  ( Query (..),
    SingleQuery (..),
    sql,
    sqlPrep,
    mkQueryInternal,
    breakQueryIntoStatements,
    mkQuery,
    encodeParam,
  )
where

import Data.ByteString (ByteString)
import Data.Either (rights)
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import Hpgsql.Builder (BinaryField)
import Hpgsql.Encoding (FieldEncoder (..), RowEncoder (..), ToPgField (..), ToPgRow (..))
import Hpgsql.InternalTypes (Query (..), SingleQuery (..), SingleQueryFragment (..), breakQueryIntoStatements, renumberParamsFrom)
import Hpgsql.ParsingInternal (BlockOrNotBlock (..), ParsingOpts (..), QQExprKind (..), blockText, flattenBlocks, parseSql)
import Hpgsql.TypeInfo (EncodingContext, Oid)
import Language.Haskell.Meta.Parse (parseExp)
import Language.Haskell.TH
import Language.Haskell.TH.Quote

-- | A useful representation for our quasiquoter parsing.
data SqlFragment
  = NonInterpolatedSqlFragment !Text
  | InterpolatedHaskellExpr !Text
  | EmbeddedQueryExpr !Text
  | SemiColonFragment
  | WhitespaceOrCommentsFragment !Text
  deriving stock (SqlFragment -> SqlFragment -> Bool
(SqlFragment -> SqlFragment -> Bool)
-> (SqlFragment -> SqlFragment -> Bool) -> Eq SqlFragment
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SqlFragment -> SqlFragment -> Bool
== :: SqlFragment -> SqlFragment -> Bool
$c/= :: SqlFragment -> SqlFragment -> Bool
/= :: SqlFragment -> SqlFragment -> Bool
Eq)

-- | Takes in a query string with query arguments as dollar-numbered arguments and returns a Query you
-- can run.
-- Examples:
-- - "SELECT * FROM table WHERE col1=$1 AND col2=$2"
-- - "SELECT * FROM table WHERE (["a", "b", "c"]'::jsonb ? 'b') = $1"
-- Question marks are interpreted literally, i.e. they have no special meaning.
-- Note that if the number of arguments does not match what's in the query string,
-- this will throw an error.
mkQuery :: (ToPgRow a) => ByteString -> a -> Query
mkQuery :: forall a. ToPgRow a => ByteString -> a -> Query
mkQuery ByteString
qryText a
p = [BlockOrNotBlock]
-> [EncodingContext -> (Maybe Oid, BinaryField)] -> Query
mkQueryInternalFromSqlStatements (ParsingOpts -> Text -> [BlockOrNotBlock]
parseSql ParsingOpts
AcceptOnlyDollarNumberedArgs (Text -> [BlockOrNotBlock]) -> Text -> [BlockOrNotBlock]
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
decodeUtf8 ByteString
qryText) (RowEncoder a
forall a. ToPgRow a => RowEncoder a
rowEncoder.toPgParams a
p)
  where
    mkQueryInternalFromSqlStatements :: [BlockOrNotBlock] -> [EncodingContext -> (Maybe Oid, BinaryField)] -> Query
    mkQueryInternalFromSqlStatements :: [BlockOrNotBlock]
-> [EncodingContext -> (Maybe Oid, BinaryField)] -> Query
mkQueryInternalFromSqlStatements [BlockOrNotBlock]
blocks [EncodingContext -> (Maybe Oid, BinaryField)]
allParams =
      let paramsLen :: Int
paramsLen = [EncodingContext -> (Maybe Oid, BinaryField)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [EncodingContext -> (Maybe Oid, BinaryField)]
allParams
          qryTextForError :: [Char]
qryTextForError = Text -> [Char]
forall a. Show a => a -> [Char]
show (Text -> [Char]) -> Text -> [Char]
forall a b. (a -> b) -> a -> b
$ [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (BlockOrNotBlock -> Text) -> [BlockOrNotBlock] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map BlockOrNotBlock -> Text
blockText [BlockOrNotBlock]
blocks
          queryFrags :: [SingleQueryFragment]
queryFrags =
            (BlockOrNotBlock -> SingleQueryFragment)
-> [BlockOrNotBlock] -> [SingleQueryFragment]
forall a b. (a -> b) -> [a] -> [b]
map
              ( \case
                  StaticSql Text
t -> ByteString -> SingleQueryFragment
FragmentOfStaticSql (ByteString -> SingleQueryFragment)
-> ByteString -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
encodeUtf8 Text
t
                  BlockOrNotBlock
SemiColon -> SingleQueryFragment
FragmentWithSemiColon
                  CommentsOrWhitespace Text
t -> ByteString -> SingleQueryFragment
FragmentOfCommentsOrWhitespace (ByteString -> SingleQueryFragment)
-> ByteString -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
encodeUtf8 Text
t
                  DollarNumberedArg Int
n ->
                    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
                      then
                        [Char] -> SingleQueryFragment
forall a. HasCallStack => [Char] -> a
error ([Char] -> SingleQueryFragment) -> [Char] -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ [Char]
"Dollar-numbered query argument placeholders must start from 1. Query: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
qryTextForError
                      else
                        if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
paramsLen
                          then
                            [Char] -> SingleQueryFragment
forall a. HasCallStack => [Char] -> a
error ([Char] -> SingleQueryFragment) -> [Char] -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ [Char]
"Query contains more dollar-numbered query arguments than actually supplied query arguments. Query: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
qryTextForError
                          else Int -> SingleQueryFragment
QueryArgumentPlaceHolder Int
n
                  BlockOrNotBlock
QuestionMarkArg ->
                    [Char] -> SingleQueryFragment
forall a. HasCallStack => [Char] -> a
error ([Char] -> SingleQueryFragment) -> [Char] -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ [Char]
"Bug in Hpgsql: parseSql AcceptOnlyDollarNumberedArgs returned question mark place holders. Query: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
qryTextForError
                  QuasiQuoterExpression QQExprKind
_ Text
_ ->
                    [Char] -> SingleQueryFragment
forall a. HasCallStack => [Char] -> a
error ([Char] -> SingleQueryFragment) -> [Char] -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ [Char]
"Bug in Hpgsql: parseSql AcceptOnlyDollarNumberedArgs returned quasiquoter expression. Query: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
qryTextForError
              )
              [BlockOrNotBlock]
blocks
       in Query {queryString :: [SingleQueryFragment]
queryString = [SingleQueryFragment]
queryFrags, queryParams :: [EncodingContext -> (Maybe Oid, BinaryField)]
queryParams = [EncodingContext -> (Maybe Oid, BinaryField)]
allParams, isPrepared :: Bool
isPrepared = Bool
False}

-- | Takes in a query string with query arguments as question marks (NOT dollar-numbered arguments).
-- Example:
-- - "SELECT * FROM table WHERE col1=? AND col2=?"
-- You should really use `mkQuery` instead of this. This is here only for hpgsql-simple-compat.
mkQueryInternal :: ByteString -> [[Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]] -> Query
mkQueryInternal :: ByteString
-> [[Either
       ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
-> Query
mkQueryInternal ByteString
queryTemplate [[Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
allParams =
  let statements :: [BlockOrNotBlock]
statements = ParsingOpts -> Text -> [BlockOrNotBlock]
parseSql ParsingOpts
AcceptQuestionMarksAsQueryArgs (ByteString -> Text
decodeUtf8 ByteString
queryTemplate)
      paramsByIdx :: Map
  Int
  [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
paramsByIdx = [(Int,
  [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))])]
-> Map
     Int
     [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Int,
   [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))])]
 -> Map
      Int
      [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))])
-> [(Int,
     [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))])]
-> Map
     Int
     [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
forall a b. (a -> b) -> a -> b
$ [Int]
-> [[Either
       ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
-> [(Int,
     [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))])]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int
1 :: Int) ..] [[Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
allParams
      qryTextForError :: Text
qryTextForError = [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (BlockOrNotBlock -> Text) -> [BlockOrNotBlock] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map BlockOrNotBlock -> Text
blockText [BlockOrNotBlock]
statements
      ((Int, Int)
_, [[SingleQueryFragment]]
queryFrags) =
        ((Int, Int)
 -> BlockOrNotBlock -> ((Int, Int), [SingleQueryFragment]))
-> (Int, Int)
-> [BlockOrNotBlock]
-> ((Int, Int), [[SingleQueryFragment]])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL
          ( \(!Int
maxArgSoFar, !Int
maxRealArgSoFar) BlockOrNotBlock
sqlPiece -> case BlockOrNotBlock
sqlPiece of
              StaticSql Text
t -> ((Int
maxArgSoFar, Int
maxRealArgSoFar), [ByteString -> SingleQueryFragment
FragmentOfStaticSql (ByteString -> SingleQueryFragment)
-> ByteString -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
encodeUtf8 Text
t])
              BlockOrNotBlock
SemiColon -> ((Int
maxArgSoFar, Int
maxRealArgSoFar), [SingleQueryFragment
FragmentWithSemiColon])
              CommentsOrWhitespace Text
t -> ((Int
maxArgSoFar, Int
maxRealArgSoFar), [ByteString -> SingleQueryFragment
FragmentOfCommentsOrWhitespace (ByteString -> SingleQueryFragment)
-> ByteString -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
encodeUtf8 Text
t])
              DollarNumberedArg Int
_ ->
                [Char] -> ((Int, Int), [SingleQueryFragment])
forall a. HasCallStack => [Char] -> a
error ([Char] -> ((Int, Int), [SingleQueryFragment]))
-> [Char] -> ((Int, Int), [SingleQueryFragment])
forall a b. (a -> b) -> a -> b
$ [Char]
"Bug in Hpgsql: parsed a DollarNumberedArg in mkQueryInternal. Query: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ByteString, Text) -> [Char]
forall a. Show a => a -> [Char]
show (ByteString
queryTemplate, Text
qryTextForError)
              QuasiQuoterExpression QQExprKind
_ Text
_ ->
                [Char] -> ((Int, Int), [SingleQueryFragment])
forall a. HasCallStack => [Char] -> a
error ([Char] -> ((Int, Int), [SingleQueryFragment]))
-> [Char] -> ((Int, Int), [SingleQueryFragment])
forall a b. (a -> b) -> a -> b
$ [Char]
"Bug in Hpgsql: parsed a QuasiQuoterExpression in mkQueryInternal. Query: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ByteString, Text) -> [Char]
forall a. Show a => a -> [Char]
show (ByteString
queryTemplate, Text
qryTextForError)
              BlockOrNotBlock
QuestionMarkArg ->
                let thisParamNum :: Int
thisParamNum = Int
maxArgSoFar Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                 in case Int
-> Map
     Int
     [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
-> Maybe
     [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int
thisParamNum Map
  Int
  [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
paramsByIdx of
                      Maybe
  [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
Nothing -> [Char] -> ((Int, Int), [SingleQueryFragment])
forall a. HasCallStack => [Char] -> a
error ([Char] -> ((Int, Int), [SingleQueryFragment]))
-> [Char] -> ((Int, Int), [SingleQueryFragment])
forall a b. (a -> b) -> a -> b
$ [Char]
"Could not find query argument of number " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
thisParamNum [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" for query with " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([[Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
-> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
allParams) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" arguments supplied. Did you supply an insufficient amount of query arguments? Query is " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
forall a. Show a => a -> [Char]
show Text
qryTextForError
                      Just [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
args ->
                        let (Int
newMaxRealArg, [SingleQueryFragment]
frags) =
                              (Int
 -> Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))
 -> (Int, SingleQueryFragment))
-> Int
-> [Either
      ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
-> (Int, [SingleQueryFragment])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
List.mapAccumL
                                ( \(!Int
newArgNum) Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))
arg -> case Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))
arg of
                                    Left ByteString
fakeArg ->
                                      (Int
newArgNum, ByteString -> SingleQueryFragment
FragmentOfStaticSql ByteString
fakeArg)
                                    Right EncodingContext -> (Maybe Oid, BinaryField)
_properArg ->
                                      (Int
newArgNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int -> SingleQueryFragment
QueryArgumentPlaceHolder (Int -> SingleQueryFragment) -> Int -> SingleQueryFragment
forall a b. (a -> b) -> a -> b
$ Int
newArgNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                                )
                                Int
maxRealArgSoFar
                                [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
args
                         in ((Int
thisParamNum, Int
newMaxRealArg), [SingleQueryFragment]
frags)
          )
          (Int
0, Int
0)
          [BlockOrNotBlock]
statements
   in Query {queryString :: [SingleQueryFragment]
queryString = [[SingleQueryFragment]] -> [SingleQueryFragment]
forall a. Monoid a => [a] -> a
mconcat [[SingleQueryFragment]]
queryFrags, queryParams :: [EncodingContext -> (Maybe Oid, BinaryField)]
queryParams = ([Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
 -> [EncodingContext -> (Maybe Oid, BinaryField)])
-> [[Either
       ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
-> [EncodingContext -> (Maybe Oid, BinaryField)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]
-> [EncodingContext -> (Maybe Oid, BinaryField)]
forall a b. [Either a b] -> [b]
rights [[Either ByteString (EncodingContext -> (Maybe Oid, BinaryField))]]
allParams, isPrepared :: Bool
isPrepared = Bool
False}

-- A practical quasiquoter for SQL. Interpolate Haskell values with @#{}@:
--
-- > let name = "Alice" :: Text
-- > rows <- query conn [sql|SELECT id, email FROM users WHERE name = #{name}|]
-- > print (rows :: [(Int, Text)])
--
-- Values interpolated with @#{}@ are sent as query parameters, so they are safe from SQL injection.
--
-- To embed a 'Query' fragment (e.g. a dynamic table name or a sub-query), use @^{}@:
--
-- > let tableName = escapeIdentifier "users"
-- > rows <- query conn [sql|SELECT * FROM ^{tableName}|]
sql :: QuasiQuoter
sql :: QuasiQuoter
sql =
  QuasiQuoter
    { quoteExp :: [Char] -> Q Exp
quoteExp = Bool -> [BlockOrNotBlock] -> Q Exp
liftQuery Bool
False ([BlockOrNotBlock] -> Q Exp)
-> ([Char] -> [BlockOrNotBlock]) -> [Char] -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParsingOpts -> Text -> [BlockOrNotBlock]
parseSql ParsingOpts
AcceptQuasiQuoterExpressions (Text -> [BlockOrNotBlock])
-> ([Char] -> Text) -> [Char] -> [BlockOrNotBlock]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
Text.pack,
      quotePat :: [Char] -> Q Pat
quotePat = [Char] -> [Char] -> Q Pat
forall a. HasCallStack => [Char] -> a
error [Char]
"Hpgsql's sql quasiquoter does not implement quotePat",
      quoteType :: [Char] -> Q Type
quoteType = [Char] -> [Char] -> Q Type
forall a. HasCallStack => [Char] -> a
error [Char]
"Hpgsql's sql quasiquoter does not implement quoteType",
      quoteDec :: [Char] -> Q [Dec]
quoteDec = [Char] -> [Char] -> Q [Dec]
forall a. HasCallStack => [Char] -> a
error [Char]
"Hpgsql's sql quasiquoter does not implement quoteDec"
    }

-- | Just like the @sql@ quasiquoter, but every SQL statement inside is
-- a prepared statement.
sqlPrep :: QuasiQuoter
sqlPrep :: QuasiQuoter
sqlPrep =
  QuasiQuoter
    { quoteExp :: [Char] -> Q Exp
quoteExp = Bool -> [BlockOrNotBlock] -> Q Exp
liftQuery Bool
True ([BlockOrNotBlock] -> Q Exp)
-> ([Char] -> [BlockOrNotBlock]) -> [Char] -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParsingOpts -> Text -> [BlockOrNotBlock]
parseSql ParsingOpts
AcceptQuasiQuoterExpressions (Text -> [BlockOrNotBlock])
-> ([Char] -> Text) -> [Char] -> [BlockOrNotBlock]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
Text.pack,
      quotePat :: [Char] -> Q Pat
quotePat = [Char] -> [Char] -> Q Pat
forall a. HasCallStack => [Char] -> a
error [Char]
"Hpgsql's sql quasiquoter does not implement quotePat",
      quoteType :: [Char] -> Q Type
quoteType = [Char] -> [Char] -> Q Type
forall a. HasCallStack => [Char] -> a
error [Char]
"Hpgsql's sql quasiquoter does not implement quoteType",
      quoteDec :: [Char] -> Q [Dec]
quoteDec = [Char] -> [Char] -> Q [Dec]
forall a. HasCallStack => [Char] -> a
error [Char]
"Hpgsql's sql quasiquoter does not implement quoteDec"
    }

liftQuery :: Bool -> [BlockOrNotBlock] -> Q Exp
liftQuery :: Bool -> [BlockOrNotBlock] -> Q Exp
liftQuery Bool
isPrepared ([BlockOrNotBlock] -> [BlockOrNotBlock]
flattenBlocks -> [BlockOrNotBlock]
stmt) = do
  let allFragments :: [SqlFragment]
allFragments = (BlockOrNotBlock -> [SqlFragment])
-> [BlockOrNotBlock] -> [SqlFragment]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap BlockOrNotBlock -> [SqlFragment]
parseBlockQuasiQuoter [BlockOrNotBlock]
stmt
      hasEmbedded :: Bool
hasEmbedded = (SqlFragment -> Bool) -> [SqlFragment] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any SqlFragment -> Bool
isEmbedded [SqlFragment]
allFragments
  if Bool
hasEmbedded
    then Bool -> [SqlFragment] -> Q Exp
liftQueryDynamic Bool
isPrepared [SqlFragment]
allFragments
    else Bool -> [SqlFragment] -> Q Exp
liftQueryStatic Bool
isPrepared [SqlFragment]
allFragments
  where
    isEmbedded :: SqlFragment -> Bool
isEmbedded (EmbeddedQueryExpr Text
_) = Bool
True
    isEmbedded SqlFragment
_ = Bool
False

-- | Static path: no ^{} expressions, SQL string is a compile-time literal.
liftQueryStatic :: Bool -> [SqlFragment] -> Q Exp
liftQueryStatic :: Bool -> [SqlFragment] -> Q Exp
liftQueryStatic Bool
isPrepared [SqlFragment]
allFragments = do
  let ([Q Exp]
fragQExps, [Text]
varNames) = [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars [SqlFragment]
allFragments Int
1
  fragExps <- [Q Exp] -> Q [Exp]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [Q Exp]
fragQExps
  paramExps <- mapM generateParamExp varNames
  [|Query $(pure $ ListE fragExps) $(pure $ ListE paramExps) isPrepared|]

-- | Dynamic path: has ^{} expressions, build query at runtime via buildQueryQQ.
liftQueryDynamic :: Bool -> [SqlFragment] -> Q Exp
liftQueryDynamic :: Bool -> [SqlFragment] -> Q Exp
liftQueryDynamic Bool
isPrepared [SqlFragment]
allFragments = do
  partExps <- (SqlFragment -> Q Exp) -> [SqlFragment] -> Q [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SqlFragment -> Q Exp
fragmentToPartExp [SqlFragment]
allFragments
  [|buildQueryQQ isPrepared $(pure $ ListE partExps)|]

-- | Convert a SqlFragment into a TH expression producing a QueryBuildPart.
fragmentToPartExp :: SqlFragment -> Q Exp
fragmentToPartExp :: SqlFragment -> Q Exp
fragmentToPartExp (NonInterpolatedSqlFragment Text
t) =
  [|StaticSqlPart $(Lit -> Q Exp
forall (m :: * -> *). Quote m => Lit -> m Exp
litE ([Char] -> Lit
stringL (Text -> [Char]
Text.unpack Text
t)))|]
fragmentToPartExp (InterpolatedHaskellExpr Text
haskellExpr) =
  case [Char] -> Either [Char] Exp
parseExp (Text -> [Char]
Text.unpack Text
haskellExpr) of
    Left [Char]
err -> [Char] -> Q Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> Q Exp) -> [Char] -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"Could not parse Haskell expression '" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
haskellExpr [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"': " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err
    Right Exp
expr -> [|ParamPart (encodeParam $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
expr))|]
fragmentToPartExp SqlFragment
SemiColonFragment =
  [|SemiColonPart|]
fragmentToPartExp (WhitespaceOrCommentsFragment Text
t) =
  [|WhitespaceOrCommenstPart $(Lit -> Q Exp
forall (m :: * -> *). Quote m => Lit -> m Exp
litE ([Char] -> Lit
stringL (Text -> [Char]
Text.unpack Text
t)))|]
fragmentToPartExp (EmbeddedQueryExpr Text
haskellExpr) =
  case [Char] -> Either [Char] Exp
parseExp (Text -> [Char]
Text.unpack Text
haskellExpr) of
    Left [Char]
err -> [Char] -> Q Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> Q Exp) -> [Char] -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"Could not parse Haskell expression '" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
Text.unpack Text
haskellExpr [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"': " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err
    Right Exp
expr -> [|EmbeddedQueryPart $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
expr)|]

-- | Walk through fragments in order, building SingleQueryFragment TH expressions
-- and collecting the interpolated Haskell expressions.
buildQueryFragsAndVars :: [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars :: [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars [] Int
_ = ([], [])
buildQueryFragsAndVars (NonInterpolatedSqlFragment Text
t : [SqlFragment]
rest) Int
n =
  let ([Q Exp]
restFrags, [Text]
restVars) = [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars [SqlFragment]
rest Int
n
   in ([|FragmentOfStaticSql (encodeUtf8 $(Lit -> Q Exp
forall (m :: * -> *). Quote m => Lit -> m Exp
litE ([Char] -> Lit
stringL (Text -> [Char]
Text.unpack Text
t))))|] Q Exp -> [Q Exp] -> [Q Exp]
forall a. a -> [a] -> [a]
: [Q Exp]
restFrags, [Text]
restVars)
buildQueryFragsAndVars (WhitespaceOrCommentsFragment Text
t : [SqlFragment]
rest) Int
n =
  let ([Q Exp]
restFrags, [Text]
restVars) = [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars [SqlFragment]
rest Int
n
   in ([|FragmentOfCommentsOrWhitespace (encodeUtf8 $(Lit -> Q Exp
forall (m :: * -> *). Quote m => Lit -> m Exp
litE ([Char] -> Lit
stringL (Text -> [Char]
Text.unpack Text
t))))|] Q Exp -> [Q Exp] -> [Q Exp]
forall a. a -> [a] -> [a]
: [Q Exp]
restFrags, [Text]
restVars)
buildQueryFragsAndVars (InterpolatedHaskellExpr Text
var : [SqlFragment]
rest) Int
n =
  let ([Q Exp]
restFrags, [Text]
restVars) = [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars [SqlFragment]
rest (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
   in ([|QueryArgumentPlaceHolder $(Lit -> Q Exp
forall (m :: * -> *). Quote m => Lit -> m Exp
litE (Integer -> Lit
integerL (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)))|] Q Exp -> [Q Exp] -> [Q Exp]
forall a. a -> [a] -> [a]
: [Q Exp]
restFrags, Text
var Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: [Text]
restVars)
buildQueryFragsAndVars (SqlFragment
SemiColonFragment : [SqlFragment]
rest) Int
n =
  let ([Q Exp]
restFrags, [Text]
restVars) = [SqlFragment] -> Int -> ([Q Exp], [Text])
buildQueryFragsAndVars [SqlFragment]
rest Int
n
   in ([|FragmentWithSemiColon|] Q Exp -> [Q Exp] -> [Q Exp]
forall a. a -> [a] -> [a]
: [Q Exp]
restFrags, [Text]
restVars)
buildQueryFragsAndVars (EmbeddedQueryExpr Text
_ : [SqlFragment]
_) Int
_ =
  [Char] -> ([Q Exp], [Text])
forall a. HasCallStack => [Char] -> a
error [Char]
"Bug in Hpgsql: EmbeddedQueryExpr should not appear in static path"

-- | Parts used to build a SingleQuery at runtime when ^{} embedded queries are present.
data QueryBuildPartQQ
  = StaticSqlPart !ByteString
  | ParamPart !(EncodingContext -> (Maybe Oid, BinaryField))
  | EmbeddedQueryPart !Query
  | SemiColonPart
  | WhitespaceOrCommenstPart !ByteString

buildQueryQQ :: Bool -> [QueryBuildPartQQ] -> Query
buildQueryQQ :: Bool -> [QueryBuildPartQQ] -> Query
buildQueryQQ Bool
isPrepared [QueryBuildPartQQ]
parts =
  let ([SingleQueryFragment]
queryFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
allParams) = [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [QueryBuildPartQQ]
parts Int
1
   in Query {queryString :: [SingleQueryFragment]
queryString = [SingleQueryFragment]
queryFrags, queryParams :: [EncodingContext -> (Maybe Oid, BinaryField)]
queryParams = [EncodingContext -> (Maybe Oid, BinaryField)]
allParams, Bool
isPrepared :: Bool
isPrepared :: Bool
isPrepared}
  where
    go :: [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [] Int
_ = ([], [])
    go (QueryBuildPartQQ
part : [QueryBuildPartQQ]
rest) Int
argNum = case QueryBuildPartQQ
part of
      StaticSqlPart ByteString
bs ->
        let ([SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams) = [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [QueryBuildPartQQ]
rest Int
argNum
         in (ByteString -> SingleQueryFragment
FragmentOfStaticSql ByteString
bs SingleQueryFragment
-> [SingleQueryFragment] -> [SingleQueryFragment]
forall a. a -> [a] -> [a]
: [SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams)
      WhitespaceOrCommenstPart ByteString
bs ->
        let ([SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams) = [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [QueryBuildPartQQ]
rest Int
argNum
         in (ByteString -> SingleQueryFragment
FragmentOfCommentsOrWhitespace ByteString
bs SingleQueryFragment
-> [SingleQueryFragment] -> [SingleQueryFragment]
forall a. a -> [a] -> [a]
: [SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams)
      ParamPart EncodingContext -> (Maybe Oid, BinaryField)
p ->
        let ([SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams) = [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [QueryBuildPartQQ]
rest (Int
argNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
         in (Int -> SingleQueryFragment
QueryArgumentPlaceHolder Int
argNum SingleQueryFragment
-> [SingleQueryFragment] -> [SingleQueryFragment]
forall a. a -> [a] -> [a]
: [SingleQueryFragment]
restFrags, EncodingContext -> (Maybe Oid, BinaryField)
p (EncodingContext -> (Maybe Oid, BinaryField))
-> [EncodingContext -> (Maybe Oid, BinaryField)]
-> [EncodingContext -> (Maybe Oid, BinaryField)]
forall a. a -> [a] -> [a]
: [EncodingContext -> (Maybe Oid, BinaryField)]
restParams)
      QueryBuildPartQQ
SemiColonPart ->
        let ([SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams) = [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [QueryBuildPartQQ]
rest Int
argNum
         in (SingleQueryFragment
FragmentWithSemiColon SingleQueryFragment
-> [SingleQueryFragment] -> [SingleQueryFragment]
forall a. a -> [a] -> [a]
: [SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams)
      EmbeddedQueryPart Query
q ->
        let (Int
newMaxArgNum, [SingleQueryFragment]
renumberedFrags) = [SingleQueryFragment] -> Int -> (Int, [SingleQueryFragment])
renumberParamsFrom Query
q.queryString Int
argNum
            ([SingleQueryFragment]
restFrags, [EncodingContext -> (Maybe Oid, BinaryField)]
restParams) = [QueryBuildPartQQ]
-> Int
-> ([SingleQueryFragment],
    [EncodingContext -> (Maybe Oid, BinaryField)])
go [QueryBuildPartQQ]
rest (Int
newMaxArgNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
         in ([SingleQueryFragment]
renumberedFrags [SingleQueryFragment]
-> [SingleQueryFragment] -> [SingleQueryFragment]
forall a. [a] -> [a] -> [a]
++ [SingleQueryFragment]
restFrags, Query
q.queryParams [EncodingContext -> (Maybe Oid, BinaryField)]
-> [EncodingContext -> (Maybe Oid, BinaryField)]
-> [EncodingContext -> (Maybe Oid, BinaryField)]
forall a. [a] -> [a] -> [a]
++ [EncodingContext -> (Maybe Oid, BinaryField)]
restParams)

parseBlockQuasiQuoter :: BlockOrNotBlock -> [SqlFragment]
parseBlockQuasiQuoter :: BlockOrNotBlock -> [SqlFragment]
parseBlockQuasiQuoter (StaticSql Text
text) = [Text -> SqlFragment
NonInterpolatedSqlFragment Text
text]
parseBlockQuasiQuoter BlockOrNotBlock
SemiColon = [SqlFragment
SemiColonFragment]
parseBlockQuasiQuoter (CommentsOrWhitespace Text
text) = [Text -> SqlFragment
WhitespaceOrCommentsFragment Text
text]
parseBlockQuasiQuoter (DollarNumberedArg Int
n) = [Text -> SqlFragment
NonInterpolatedSqlFragment (Text -> SqlFragment) -> Text -> SqlFragment
forall a b. (a -> b) -> a -> b
$ Text
"$" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Char] -> Text
Text.pack (Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n)] -- If someone uses e.g. $2 in a quasiquoter, it should be an error, but let's assume they know what they're doing and treat it as text
parseBlockQuasiQuoter BlockOrNotBlock
QuestionMarkArg = [Text -> SqlFragment
NonInterpolatedSqlFragment Text
"?"] -- If someone uses a question mark in a quasiquoter, it should be an error, but let's assume they know what they're doing and treat it as text
parseBlockQuasiQuoter (QuasiQuoterExpression QQExprKind
QQInterpolation Text
expr) = [Text -> SqlFragment
InterpolatedHaskellExpr Text
expr]
parseBlockQuasiQuoter (QuasiQuoterExpression QQExprKind
QQEmbeddedQuery Text
expr) = [Text -> SqlFragment
EmbeddedQueryExpr Text
expr]

-- | Generate a parameter expression for a captured variable
generateParamExp :: Text -> Q Exp
generateParamExp :: Text -> Q Exp
generateParamExp (Text -> [Char]
Text.unpack -> [Char]
haskellExpr) =
  case [Char] -> Either [Char] Exp
parseExp [Char]
haskellExpr of
    Left [Char]
err -> [Char] -> Q Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> Q Exp) -> [Char] -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"Could not parse Haskell expression '" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
haskellExpr [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"': " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err
    Right Exp
expr ->
      [|encodeParam $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
expr)|]

encodeParam :: (ToPgField a) => a -> EncodingContext -> (Maybe Oid, BinaryField)
encodeParam :: forall a.
ToPgField a =>
a -> EncodingContext -> (Maybe Oid, BinaryField)
encodeParam a
v =
  let fe :: FieldEncoder a
fe = FieldEncoder a
forall a. ToPgField a => FieldEncoder a
fieldEncoder
   in \EncodingContext
encCtx -> (FieldEncoder a
fe.toTypeOid EncodingContext
encCtx, FieldEncoder a
fe.toPgField EncodingContext
encCtx a
v)