{-# LANGUAGE CPP, PatternGuards, MultiParamTypeClasses, FunctionalDependencies, TypeSynonymInstances, FlexibleInstances, FlexibleContexts, GADTs, DataKinds, TemplateHaskell #-}
module Database.PostgreSQL.Typed.Query
  ( PGQuery(..)
  , PGSimpleQuery
  , PGPreparedQuery
  , rawPGSimpleQuery
  , rawPGPreparedQuery
  , QueryFlags(..)
  , simpleQueryFlags
  , parseQueryFlags
  , makePGQuery
  , pgSQL
  , pgExecute
  , pgQuery
  , pgLazyQuery
  ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative ((<$>))
#endif
import Control.Arrow ((***), first, second)
import Control.Exception (try)
import Control.Monad (void, when, mapAndUnzipM)
import Data.Array (listArray, (!), inRange)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Lazy.UTF8 as BSLU
import qualified Data.ByteString.UTF8 as BSU
import Data.Char (isSpace, isAlphaNum)
import qualified Data.Foldable as Fold
import Data.List (dropWhileEnd)
import Data.Maybe (fromMaybe, isNothing)
import Data.String (IsString(..))
import Data.Word (Word32)
import Language.Haskell.Meta.Parse (parseExp)
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Quote (QuasiQuoter(..))

import Database.PostgreSQL.Typed.Types
import Database.PostgreSQL.Typed.Dynamic
import Database.PostgreSQL.Typed.Protocol
import Database.PostgreSQL.Typed.TH
import Database.PostgreSQL.Typed.SQLToken

class PGQuery q a | q -> a where
  -- |Execute a query and return the number of rows affected (or -1 if not known) and a list of results.
  pgRunQuery :: PGConnection -> q -> IO (Int, [a])
  -- |Change the raw SQL query stored within this query.
  -- This is unsafe because the query has already been type-checked, so any change must not change the number or type of results or placeholders (so adding additional static WHERE or ORDER BY clauses is generally safe).
  -- This is useful in cases where you need to construct some part of the query dynamically, but still want to infer the result types.
  -- If you want to add dynamic values to the query, it's best to use 'Database.PostgreSQL.Typed.Dynamic.pgSafeLiteral'.
  -- For example:
  --
  -- > [pgSQL|SELECT a FROM t|] `unsafeModifyQuery` (<> (" WHERE a = " <> pgSafeLiteral x))
  unsafeModifyQuery :: q -> (BS.ByteString -> BS.ByteString) -> q
  getQueryString :: PGTypeEnv -> q -> BS.ByteString
class PGQuery q PGValues => PGRawQuery q

-- |Execute a query that does not return results.
-- Return the number of rows affected (or -1 if not known).
pgExecute :: PGQuery q () => PGConnection -> q -> IO Int
pgExecute :: PGConnection -> q -> IO Int
pgExecute PGConnection
c q
q = (Int, [()]) -> Int
forall a b. (a, b) -> a
fst ((Int, [()]) -> Int) -> IO (Int, [()]) -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGConnection -> q -> IO (Int, [()])
forall q a. PGQuery q a => PGConnection -> q -> IO (Int, [a])
pgRunQuery PGConnection
c q
q

-- |Run a query and return a list of row results.
pgQuery :: PGQuery q a => PGConnection -> q -> IO [a]
pgQuery :: PGConnection -> q -> IO [a]
pgQuery PGConnection
c q
q = (Int, [a]) -> [a]
forall a b. (a, b) -> b
snd ((Int, [a]) -> [a]) -> IO (Int, [a]) -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGConnection -> q -> IO (Int, [a])
forall q a. PGQuery q a => PGConnection -> q -> IO (Int, [a])
pgRunQuery PGConnection
c q
q

instance PGQuery BS.ByteString PGValues where
  pgRunQuery :: PGConnection -> ByteString -> IO (Int, [PGValues])
pgRunQuery PGConnection
c ByteString
sql = PGConnection -> ByteString -> IO (Int, [PGValues])
pgSimpleQuery PGConnection
c (ByteString -> ByteString
BSL.fromStrict ByteString
sql)
  unsafeModifyQuery :: ByteString -> (ByteString -> ByteString) -> ByteString
unsafeModifyQuery ByteString
q ByteString -> ByteString
f = ByteString -> ByteString
f ByteString
q
  getQueryString :: PGTypeEnv -> ByteString -> ByteString
getQueryString PGTypeEnv
_ = ByteString -> ByteString
forall a. a -> a
id

newtype SimpleQuery = SimpleQuery BS.ByteString
  deriving (Int -> SimpleQuery -> ShowS
[SimpleQuery] -> ShowS
SimpleQuery -> String
(Int -> SimpleQuery -> ShowS)
-> (SimpleQuery -> String)
-> ([SimpleQuery] -> ShowS)
-> Show SimpleQuery
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SimpleQuery] -> ShowS
$cshowList :: [SimpleQuery] -> ShowS
show :: SimpleQuery -> String
$cshow :: SimpleQuery -> String
showsPrec :: Int -> SimpleQuery -> ShowS
$cshowsPrec :: Int -> SimpleQuery -> ShowS
Show)
instance PGQuery SimpleQuery PGValues where
  pgRunQuery :: PGConnection -> SimpleQuery -> IO (Int, [PGValues])
pgRunQuery PGConnection
c (SimpleQuery ByteString
sql) = PGConnection -> ByteString -> IO (Int, [PGValues])
pgSimpleQuery PGConnection
c (ByteString -> ByteString
BSL.fromStrict ByteString
sql)
  unsafeModifyQuery :: SimpleQuery -> (ByteString -> ByteString) -> SimpleQuery
unsafeModifyQuery (SimpleQuery ByteString
sql) ByteString -> ByteString
f = ByteString -> SimpleQuery
SimpleQuery (ByteString -> SimpleQuery) -> ByteString -> SimpleQuery
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
f ByteString
sql
  getQueryString :: PGTypeEnv -> SimpleQuery -> ByteString
getQueryString PGTypeEnv
_ (SimpleQuery ByteString
q) = ByteString
q
instance PGRawQuery SimpleQuery

data PreparedQuery = PreparedQuery BS.ByteString [OID] PGValues [Bool]
  deriving (Int -> PreparedQuery -> ShowS
[PreparedQuery] -> ShowS
PreparedQuery -> String
(Int -> PreparedQuery -> ShowS)
-> (PreparedQuery -> String)
-> ([PreparedQuery] -> ShowS)
-> Show PreparedQuery
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PreparedQuery] -> ShowS
$cshowList :: [PreparedQuery] -> ShowS
show :: PreparedQuery -> String
$cshow :: PreparedQuery -> String
showsPrec :: Int -> PreparedQuery -> ShowS
$cshowsPrec :: Int -> PreparedQuery -> ShowS
Show)
instance PGQuery PreparedQuery PGValues where
  pgRunQuery :: PGConnection -> PreparedQuery -> IO (Int, [PGValues])
pgRunQuery PGConnection
c (PreparedQuery ByteString
sql [OID]
types PGValues
bind [Bool]
bc) = PGConnection
-> ByteString
-> [OID]
-> PGValues
-> [Bool]
-> IO (Int, [PGValues])
pgPreparedQuery PGConnection
c ByteString
sql [OID]
types PGValues
bind [Bool]
bc
  unsafeModifyQuery :: PreparedQuery -> (ByteString -> ByteString) -> PreparedQuery
unsafeModifyQuery (PreparedQuery ByteString
sql [OID]
types PGValues
bind [Bool]
bc) ByteString -> ByteString
f = ByteString -> [OID] -> PGValues -> [Bool] -> PreparedQuery
PreparedQuery (ByteString -> ByteString
f ByteString
sql) [OID]
types PGValues
bind [Bool]
bc
  getQueryString :: PGTypeEnv -> PreparedQuery -> ByteString
getQueryString PGTypeEnv
_ (PreparedQuery ByteString
q [OID]
_ PGValues
_ [Bool]
_) = ByteString
q
instance PGRawQuery PreparedQuery


data QueryParser q a = QueryParser (PGTypeEnv -> q) (PGTypeEnv -> PGValues -> a)
instance PGRawQuery q => PGQuery (QueryParser q a) a where
  pgRunQuery :: PGConnection -> QueryParser q a -> IO (Int, [a])
pgRunQuery PGConnection
c (QueryParser PGTypeEnv -> q
q PGTypeEnv -> PGValues -> a
p) = ([PGValues] -> [a]) -> (Int, [PGValues]) -> (Int, [a])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((PGValues -> a) -> [PGValues] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((PGValues -> a) -> [PGValues] -> [a])
-> (PGValues -> a) -> [PGValues] -> [a]
forall a b. (a -> b) -> a -> b
$ PGTypeEnv -> PGValues -> a
p PGTypeEnv
e) ((Int, [PGValues]) -> (Int, [a]))
-> IO (Int, [PGValues]) -> IO (Int, [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGConnection -> q -> IO (Int, [PGValues])
forall q a. PGQuery q a => PGConnection -> q -> IO (Int, [a])
pgRunQuery PGConnection
c (PGTypeEnv -> q
q PGTypeEnv
e) where e :: PGTypeEnv
e = PGConnection -> PGTypeEnv
pgTypeEnv PGConnection
c
  unsafeModifyQuery :: QueryParser q a -> (ByteString -> ByteString) -> QueryParser q a
unsafeModifyQuery (QueryParser PGTypeEnv -> q
q PGTypeEnv -> PGValues -> a
p) ByteString -> ByteString
f = (PGTypeEnv -> q) -> (PGTypeEnv -> PGValues -> a) -> QueryParser q a
forall q a.
(PGTypeEnv -> q) -> (PGTypeEnv -> PGValues -> a) -> QueryParser q a
QueryParser (\PGTypeEnv
e -> q -> (ByteString -> ByteString) -> q
forall q a. PGQuery q a => q -> (ByteString -> ByteString) -> q
unsafeModifyQuery (PGTypeEnv -> q
q PGTypeEnv
e) ByteString -> ByteString
f) PGTypeEnv -> PGValues -> a
p
  getQueryString :: PGTypeEnv -> QueryParser q a -> ByteString
getQueryString PGTypeEnv
e (QueryParser PGTypeEnv -> q
q PGTypeEnv -> PGValues -> a
_) = PGTypeEnv -> q -> ByteString
forall q a. PGQuery q a => PGTypeEnv -> q -> ByteString
getQueryString PGTypeEnv
e (q -> ByteString) -> q -> ByteString
forall a b. (a -> b) -> a -> b
$ PGTypeEnv -> q
q PGTypeEnv
e

instance Functor (QueryParser q) where
  fmap :: (a -> b) -> QueryParser q a -> QueryParser q b
fmap a -> b
f (QueryParser PGTypeEnv -> q
q PGTypeEnv -> PGValues -> a
p) = (PGTypeEnv -> q) -> (PGTypeEnv -> PGValues -> b) -> QueryParser q b
forall q a.
(PGTypeEnv -> q) -> (PGTypeEnv -> PGValues -> a) -> QueryParser q a
QueryParser PGTypeEnv -> q
q (\PGTypeEnv
e -> a -> b
f (a -> b) -> (PGValues -> a) -> PGValues -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PGTypeEnv -> PGValues -> a
p PGTypeEnv
e)

instance Show q => Show (QueryParser q a) where
  showsPrec :: Int -> QueryParser q a -> ShowS
showsPrec Int
p (QueryParser PGTypeEnv -> q
q PGTypeEnv -> PGValues -> a
_) = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
    String -> ShowS
showString String
"QueryParser " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> q -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (PGTypeEnv -> q
q PGTypeEnv
unknownPGTypeEnv)

rawParser :: q -> QueryParser q PGValues
rawParser :: q -> QueryParser q PGValues
rawParser q
q = (PGTypeEnv -> q)
-> (PGTypeEnv -> PGValues -> PGValues) -> QueryParser q PGValues
forall q a.
(PGTypeEnv -> q) -> (PGTypeEnv -> PGValues -> a) -> QueryParser q a
QueryParser (q -> PGTypeEnv -> q
forall a b. a -> b -> a
const q
q) ((PGValues -> PGValues) -> PGTypeEnv -> PGValues -> PGValues
forall a b. a -> b -> a
const PGValues -> PGValues
forall a. a -> a
id)

-- |A simple one-shot query that simply substitutes literal representations of parameters for placeholders.
type PGSimpleQuery = QueryParser SimpleQuery
-- |A prepared query that automatically is prepared in the database the first time it is run and bound with new parameters each subsequent time.
type PGPreparedQuery = QueryParser PreparedQuery

-- |Make a simple query directly from a query string, with no type inference
rawPGSimpleQuery :: BS.ByteString -> PGSimpleQuery PGValues
rawPGSimpleQuery :: ByteString -> PGSimpleQuery PGValues
rawPGSimpleQuery = SimpleQuery -> PGSimpleQuery PGValues
forall q. q -> QueryParser q PGValues
rawParser (SimpleQuery -> PGSimpleQuery PGValues)
-> (ByteString -> SimpleQuery)
-> ByteString
-> PGSimpleQuery PGValues
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> SimpleQuery
SimpleQuery

instance IsString (PGSimpleQuery PGValues) where
  fromString :: String -> PGSimpleQuery PGValues
fromString = ByteString -> PGSimpleQuery PGValues
rawPGSimpleQuery (ByteString -> PGSimpleQuery PGValues)
-> (String -> ByteString) -> String -> PGSimpleQuery PGValues
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
forall a. IsString a => String -> a
fromString
instance IsString (PGSimpleQuery ()) where
  fromString :: String -> PGSimpleQuery ()
fromString = PGSimpleQuery PGValues -> PGSimpleQuery ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (PGSimpleQuery PGValues -> PGSimpleQuery ())
-> (String -> PGSimpleQuery PGValues) -> String -> PGSimpleQuery ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> PGSimpleQuery PGValues
rawPGSimpleQuery (ByteString -> PGSimpleQuery PGValues)
-> (String -> ByteString) -> String -> PGSimpleQuery PGValues
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
forall a. IsString a => String -> a
fromString

-- |Make a prepared query directly from a query string and bind parameters, with no type inference
rawPGPreparedQuery :: BS.ByteString -> PGValues -> PGPreparedQuery PGValues
rawPGPreparedQuery :: ByteString -> PGValues -> PGPreparedQuery PGValues
rawPGPreparedQuery ByteString
sql PGValues
bind = PreparedQuery -> PGPreparedQuery PGValues
forall q. q -> QueryParser q PGValues
rawParser (PreparedQuery -> PGPreparedQuery PGValues)
-> PreparedQuery -> PGPreparedQuery PGValues
forall a b. (a -> b) -> a -> b
$ ByteString -> [OID] -> PGValues -> [Bool] -> PreparedQuery
PreparedQuery ByteString
sql [] PGValues
bind []

-- |Run a prepared query in lazy mode, where only chunk size rows are requested at a time.
-- If you eventually retrieve all the rows this way, it will be far less efficient than using @pgQuery@, since every chunk requires an additional round-trip.
-- Although you may safely stop consuming rows early, currently you may not interleave any other database operation while reading rows.  (This limitation could theoretically be lifted if required.)
pgLazyQuery :: PGConnection -> PGPreparedQuery a -> Word32 -- ^ Chunk size (1 is common, 0 is all-or-nothing)
  -> IO [a]
pgLazyQuery :: PGConnection -> PGPreparedQuery a -> OID -> IO [a]
pgLazyQuery PGConnection
c (QueryParser PGTypeEnv -> PreparedQuery
q PGTypeEnv -> PGValues -> a
p) OID
count =
  (PGValues -> a) -> [PGValues] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (PGTypeEnv -> PGValues -> a
p PGTypeEnv
e) ([PGValues] -> [a]) -> IO [PGValues] -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGConnection
-> ByteString
-> [OID]
-> PGValues
-> [Bool]
-> OID
-> IO [PGValues]
pgPreparedLazyQuery PGConnection
c ByteString
sql [OID]
types PGValues
bind [Bool]
bc OID
count where
  e :: PGTypeEnv
e = PGConnection -> PGTypeEnv
pgTypeEnv PGConnection
c
  PreparedQuery ByteString
sql [OID]
types PGValues
bind [Bool]
bc = PGTypeEnv -> PreparedQuery
q PGTypeEnv
e

-- |Given a SQL statement with placeholders of the form @${expr}@, return a (hopefully) valid SQL statement with @$N@ placeholders and the list of expressions.
-- This does its best to understand SQL syntax, so placeholders are only interpreted in places postgres would understand them (i.e., not in quoted strings).  Since this is not valid SQL otherwise, there is never reason to escape a literal @${@.
-- You can use @$N@ placeholders in the query otherwise to refer to the N-th index placeholder expression.
sqlPlaceholders :: String -> (String, [String])
sqlPlaceholders :: String -> (String, [String])
sqlPlaceholders = Int -> [SQLToken] -> (String, [String])
forall t. (Show t, Enum t) => t -> [SQLToken] -> (String, [String])
sst (Int
1 :: Int) ([SQLToken] -> (String, [String]))
-> (String -> [SQLToken]) -> String -> (String, [String])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [SQLToken]
sqlTokens where
  sst :: t -> [SQLToken] -> (String, [String])
sst t
n (SQLExpr String
e : [SQLToken]
l) = ((Char
'$'Char -> ShowS
forall a. a -> [a] -> [a]
:t -> String
forall a. Show a => a -> String
show t
n) String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS
-> ([String] -> [String])
-> (String, [String])
-> (String, [String])
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** (String
e String -> [String] -> [String]
forall a. a -> [a] -> [a]
:) ((String, [String]) -> (String, [String]))
-> (String, [String]) -> (String, [String])
forall a b. (a -> b) -> a -> b
$ t -> [SQLToken] -> (String, [String])
sst (t -> t
forall a. Enum a => a -> a
succ t
n) [SQLToken]
l
  sst t
n (SQLToken
t : [SQLToken]
l) = ShowS -> (String, [String]) -> (String, [String])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (SQLToken -> String
forall a. Show a => a -> String
show SQLToken
t String -> ShowS
forall a. [a] -> [a] -> [a]
++) ((String, [String]) -> (String, [String]))
-> (String, [String]) -> (String, [String])
forall a b. (a -> b) -> a -> b
$ t -> [SQLToken] -> (String, [String])
sst t
n [SQLToken]
l
  sst t
_ [] = (String
"", [])

-- |Given a SQL statement with placeholders of the form @$N@ and a list of TH 'ByteString' expressions, return a new 'ByteString' expression that substitutes the expressions for the placeholders.
sqlSubstitute :: String -> [TH.Exp] -> TH.Exp
sqlSubstitute :: String -> [Exp] -> Exp
sqlSubstitute String
sql [Exp]
exprl = Exp -> Exp -> Exp
TH.AppE (Name -> Exp
TH.VarE 'BS.concat) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
TH.ListE ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (SQLToken -> Exp) -> [SQLToken] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map SQLToken -> Exp
sst ([SQLToken] -> [Exp]) -> [SQLToken] -> [Exp]
forall a b. (a -> b) -> a -> b
$ String -> [SQLToken]
sqlTokens String
sql where
  bnds :: (Int, Int)
bnds = (Int
1, [Exp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
exprl)
  exprs :: Array Int Exp
exprs = (Int, Int) -> [Exp] -> Array Int Exp
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int, Int)
bnds [Exp]
exprl
  expr :: Int -> Exp
expr Int
n
    | (Int, Int) -> Int -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (Int, Int)
bnds Int
n = Array Int Exp
exprs Array Int Exp -> Int -> Exp
forall i e. Ix i => Array i e -> i -> e
! Int
n
    | Bool
otherwise = String -> Exp
forall a. HasCallStack => String -> a
error (String -> Exp) -> String -> Exp
forall a b. (a -> b) -> a -> b
$ String
"SQL placeholder '$" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"' out of range (not recognized by PostgreSQL)"
  sst :: SQLToken -> Exp
sst (SQLParam Int
n) = Int -> Exp
expr Int
n
  sst SQLToken
t = Name -> Exp
TH.VarE 'BSU.fromString Exp -> Exp -> Exp
`TH.AppE` Lit -> Exp
TH.LitE (String -> Lit
TH.StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ SQLToken -> String
forall a. Show a => a -> String
show SQLToken
t)

splitCommas :: String -> [String]
splitCommas :: String -> [String]
splitCommas = String -> [String]
spl where
  spl :: String -> [String]
spl [] = []
  spl [Char
c] = [[Char
c]]
  spl (Char
',':String
s) = String
""String -> [String] -> [String]
forall a. a -> [a] -> [a]
:String -> [String]
spl String
s
  spl (Char
c:String
s) = (Char
cChar -> ShowS
forall a. a -> [a] -> [a]
:String
h)String -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
t where String
h:[String]
t = String -> [String]
spl String
s

trim :: String -> String
trim :: ShowS
trim = (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhileEnd Char -> Bool
isSpace ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace

-- |Flags affecting how and what type of query to build with 'makePGQuery'.
data QueryFlags = QueryFlags
  { QueryFlags -> Bool
flagQuery :: Bool -- ^ Create a query -- otherwise just call 'pgSubstituteLiterals' to create a string (SQL fragment).
  , QueryFlags -> Maybe Bool
flagNullable :: Maybe Bool -- ^ Disable nullability inference, treating all values as nullable (if 'True') or not (if 'False').
  , QueryFlags -> Maybe [String]
flagPrepare :: Maybe [String] -- ^ Prepare and re-use query, binding parameters of the given types (inferring the rest, like PREPARE).
  }

-- |'QueryFlags' for a default (simple) query.
simpleQueryFlags :: QueryFlags
simpleQueryFlags :: QueryFlags
simpleQueryFlags = Bool -> Maybe Bool -> Maybe [String] -> QueryFlags
QueryFlags Bool
True Maybe Bool
forall a. Maybe a
Nothing Maybe [String]
forall a. Maybe a
Nothing

newName :: Char -> BS.ByteString -> TH.Q TH.Name
newName :: Char -> ByteString -> Q Name
newName Char
pre = String -> Q Name
TH.newName (String -> Q Name)
-> (ByteString -> String) -> ByteString -> Q Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
'_'Char -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> (ByteString -> String) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
preChar -> ShowS
forall a. a -> [a] -> [a]
:) ShowS -> (ByteString -> String) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
filter (\Char
c -> Char -> Bool
isAlphaNum Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'_') ShowS -> (ByteString -> String) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BSC.unpack

-- |Construct a 'PGQuery' from a SQL string.
-- This is the underlying template function for 'pgSQL' which you can use in largely the same way when you want to construct query strings from other variables.
-- For example:
--
-- > selectQuery = "SELECT * FROM"
-- > selectFoo = $(makePGQuery simpleQueryFlags (selectQuery ++ " foo"))
--
-- The only caveat is that variables or functions like @selectQuery@ need to be defined in a different module (due to TH stage restrictions).
makePGQuery :: QueryFlags -> String -> TH.ExpQ
makePGQuery :: QueryFlags -> String -> ExpQ
makePGQuery QueryFlags{ flagQuery :: QueryFlags -> Bool
flagQuery = Bool
False } String
sqle = String -> ExpQ
pgSubstituteLiterals String
sqle
makePGQuery QueryFlags{ flagNullable :: QueryFlags -> Maybe Bool
flagNullable = Maybe Bool
nulls, flagPrepare :: QueryFlags -> Maybe [String]
flagPrepare = Maybe [String]
prep } String
sqle = do
  ([TPGValueInfo]
pt, [TPGValueInfo]
rt) <- IO ([TPGValueInfo], [TPGValueInfo])
-> Q ([TPGValueInfo], [TPGValueInfo])
forall a. IO a -> Q a
TH.runIO (IO ([TPGValueInfo], [TPGValueInfo])
 -> Q ([TPGValueInfo], [TPGValueInfo]))
-> IO ([TPGValueInfo], [TPGValueInfo])
-> Q ([TPGValueInfo], [TPGValueInfo])
forall a b. (a -> b) -> a -> b
$ ByteString
-> [String] -> Bool -> IO ([TPGValueInfo], [TPGValueInfo])
tpgDescribe (String -> ByteString
BSU.fromString String
sqlp) ([String] -> Maybe [String] -> [String]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [String]
prep) (Maybe Bool -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Bool
nulls)
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([TPGValueInfo] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TPGValueInfo]
pt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
exprs) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Not all expression placeholders were recognized by PostgreSQL"

  Name
e <- String -> Q Name
TH.newName String
"_tenv"
  Name
l <- String -> Q Name
TH.newName String
"l"
  ([Pat]
vars, [Exp]
vals) <- (TPGValueInfo -> Q (Pat, Exp))
-> [TPGValueInfo] -> Q ([Pat], [Exp])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM (\TPGValueInfo
t -> do
    Name
v <- Char -> ByteString -> Q Name
newName Char
'p' (ByteString -> Q Name) -> ByteString -> Q Name
forall a b. (a -> b) -> a -> b
$ TPGValueInfo -> ByteString
tpgValueName TPGValueInfo
t
    (Pat, Exp) -> Q (Pat, Exp)
forall (m :: * -> *) a. Monad m => a -> m a
return 
      ( Name -> Pat
TH.VarP Name
v
      , Bool -> TPGValueInfo -> Name -> Exp
tpgTypeEncoder (Maybe [String] -> Bool
forall a. Maybe a -> Bool
isNothing Maybe [String]
prep) TPGValueInfo
t Name
e Exp -> Exp -> Exp
`TH.AppE` Name -> Exp
TH.VarE Name
v
      )) [TPGValueInfo]
pt
  ([Pat]
pats, [Exp]
conv, [Exp]
bins) <- [(Pat, Exp, Exp)] -> ([Pat], [Exp], [Exp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(Pat, Exp, Exp)] -> ([Pat], [Exp], [Exp]))
-> Q [(Pat, Exp, Exp)] -> Q ([Pat], [Exp], [Exp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TPGValueInfo -> Q (Pat, Exp, Exp))
-> [TPGValueInfo] -> Q [(Pat, Exp, Exp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\TPGValueInfo
t -> do
    Name
v <- Char -> ByteString -> Q Name
newName Char
'c' (ByteString -> Q Name) -> ByteString -> Q Name
forall a b. (a -> b) -> a -> b
$ TPGValueInfo -> ByteString
tpgValueName TPGValueInfo
t
    (Pat, Exp, Exp) -> Q (Pat, Exp, Exp)
forall (m :: * -> *) a. Monad m => a -> m a
return 
      ( Name -> Pat
TH.VarP Name
v
      , Bool -> TPGValueInfo -> Name -> Exp
tpgTypeDecoder (Maybe Bool -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
Fold.and Maybe Bool
nulls) TPGValueInfo
t Name
e Exp -> Exp -> Exp
`TH.AppE` Name -> Exp
TH.VarE Name
v
      , TPGValueInfo -> Name -> Exp
tpgTypeBinary TPGValueInfo
t Name
e
      )) [TPGValueInfo]
rt
  (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
TH.AppE ([Pat] -> Exp -> Exp
TH.LamE [Pat]
vars (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
TH.ConE 'QueryParser
    Exp -> Exp -> Exp
`TH.AppE` [Pat] -> Exp -> Exp
TH.LamE [Name -> Pat
TH.VarP Name
e] (Exp -> ([String] -> Exp) -> Maybe [String] -> Exp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
      (Name -> Exp
TH.ConE 'SimpleQuery
        Exp -> Exp -> Exp
`TH.AppE` String -> [Exp] -> Exp
sqlSubstitute String
sqlp [Exp]
vals)
      (\[String]
p -> Name -> Exp
TH.ConE 'PreparedQuery
        Exp -> Exp -> Exp
`TH.AppE` (Name -> Exp
TH.VarE 'BSU.fromString Exp -> Exp -> Exp
`TH.AppE` Lit -> Exp
TH.LitE (String -> Lit
TH.StringL String
sqlp))
        Exp -> Exp -> Exp
`TH.AppE` [Exp] -> Exp
TH.ListE (((String, TPGValueInfo) -> Exp)
-> [(String, TPGValueInfo)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Lit -> Exp
TH.LitE (Lit -> Exp)
-> ((String, TPGValueInfo) -> Lit) -> (String, TPGValueInfo) -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Lit
TH.IntegerL (Integer -> Lit)
-> ((String, TPGValueInfo) -> Integer)
-> (String, TPGValueInfo)
-> Lit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OID -> Integer
forall a. Integral a => a -> Integer
toInteger (OID -> Integer)
-> ((String, TPGValueInfo) -> OID)
-> (String, TPGValueInfo)
-> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TPGValueInfo -> OID
tpgValueTypeOID (TPGValueInfo -> OID)
-> ((String, TPGValueInfo) -> TPGValueInfo)
-> (String, TPGValueInfo)
-> OID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, TPGValueInfo) -> TPGValueInfo
forall a b. (a, b) -> b
snd) ([(String, TPGValueInfo)] -> [Exp])
-> [(String, TPGValueInfo)] -> [Exp]
forall a b. (a -> b) -> a -> b
$ [String] -> [TPGValueInfo] -> [(String, TPGValueInfo)]
forall a b. [a] -> [b] -> [(a, b)]
zip [String]
p [TPGValueInfo]
pt)
        Exp -> Exp -> Exp
`TH.AppE` [Exp] -> Exp
TH.ListE [Exp]
vals 
        Exp -> Exp -> Exp
`TH.AppE` [Exp] -> Exp
TH.ListE 
#ifdef VERSION_postgresql_binary
          [Exp]
bins
#else
          []
#endif
      )
      Maybe [String]
prep)
    Exp -> Exp -> Exp
`TH.AppE` [Pat] -> Exp -> Exp
TH.LamE [Name -> Pat
TH.VarP Name
e, Name -> Pat
TH.VarP Name
l] (Exp -> [Match] -> Exp
TH.CaseE (Name -> Exp
TH.VarE Name
l)
      [ Pat -> Body -> [Dec] -> Match
TH.Match ([Pat] -> Pat
TH.ListP [Pat]
pats) (Exp -> Body
TH.NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ case [Exp]
conv of
          [Exp
x] -> Exp
x
          [Exp]
_ -> [Maybe Exp] -> Exp
TH.TupE
#if MIN_VERSION_template_haskell(2,16,0)
            ([Maybe Exp] -> Exp) -> [Maybe Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Maybe Exp) -> [Exp] -> [Maybe Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Maybe Exp
forall a. a -> Maybe a
Just
#endif
            [Exp]
conv) []
      , Pat -> Body -> [Dec] -> Match
TH.Match Pat
TH.WildP (Exp -> Body
TH.NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
TH.VarE 'error Exp -> Exp -> Exp
`TH.AppE` Lit -> Exp
TH.LitE (String -> Lit
TH.StringL String
"pgSQL: result arity mismatch")) []
      ]))
    ([Exp] -> Exp) -> Q [Exp] -> ExpQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (String -> ExpQ) -> [String] -> Q [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM String -> ExpQ
forall (m :: * -> *). MonadFail m => String -> m Exp
parse [String]
exprs
  where
  (String
sqlp, [String]
exprs) = String -> (String, [String])
sqlPlaceholders String
sqle
  parse :: String -> m Exp
parse String
e = (String -> m Exp) -> (Exp -> m Exp) -> Either String Exp -> m Exp
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> m Exp
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m Exp) -> ShowS -> String -> m Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
forall a. [a] -> [a] -> [a]
(++) (String
"Failed to parse expression {" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
e String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"}: ")) Exp -> m Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Exp -> m Exp) -> Either String Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ String -> Either String Exp
parseExp String
e

-- |Parse flags off the beginning of a query string, returning the flags and the remaining string.
parseQueryFlags :: String -> (QueryFlags, String)
parseQueryFlags :: String -> (QueryFlags, String)
parseQueryFlags = QueryFlags -> String -> (QueryFlags, String)
pqf QueryFlags
simpleQueryFlags where
  pqf :: QueryFlags -> String -> (QueryFlags, String)
pqf f :: QueryFlags
f@QueryFlags{ flagQuery :: QueryFlags -> Bool
flagQuery = Bool
True, flagPrepare :: QueryFlags -> Maybe [String]
flagPrepare = Maybe [String]
Nothing } (Char
'#':String
q) = QueryFlags -> String -> (QueryFlags, String)
pqf QueryFlags
f{ flagQuery :: Bool
flagQuery = Bool
False } String
q
  pqf f :: QueryFlags
f@QueryFlags{ flagQuery :: QueryFlags -> Bool
flagQuery = Bool
True, flagNullable :: QueryFlags -> Maybe Bool
flagNullable = Maybe Bool
Nothing } (Char
'?':String
q) = QueryFlags -> String -> (QueryFlags, String)
pqf QueryFlags
f{ flagNullable :: Maybe Bool
flagNullable = Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True } String
q
  pqf f :: QueryFlags
f@QueryFlags{ flagQuery :: QueryFlags -> Bool
flagQuery = Bool
True, flagNullable :: QueryFlags -> Maybe Bool
flagNullable = Maybe Bool
Nothing } (Char
'!':String
q) = QueryFlags -> String -> (QueryFlags, String)
pqf QueryFlags
f{ flagNullable :: Maybe Bool
flagNullable = Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False } String
q
  pqf f :: QueryFlags
f@QueryFlags{ flagQuery :: QueryFlags -> Bool
flagQuery = Bool
True, flagPrepare :: QueryFlags -> Maybe [String]
flagPrepare = Maybe [String]
Nothing } (Char
'$':String
q) = QueryFlags -> String -> (QueryFlags, String)
pqf QueryFlags
f{ flagPrepare :: Maybe [String]
flagPrepare = [String] -> Maybe [String]
forall a. a -> Maybe a
Just [] } String
q
  pqf f :: QueryFlags
f@QueryFlags{ flagQuery :: QueryFlags -> Bool
flagQuery = Bool
True, flagPrepare :: QueryFlags -> Maybe [String]
flagPrepare = Just [] } (Char
'(':String
s) = QueryFlags -> String -> (QueryFlags, String)
pqf QueryFlags
f{ flagPrepare :: Maybe [String]
flagPrepare = [String] -> Maybe [String]
forall a. a -> Maybe a
Just [String]
args } (ShowS
sql String
r) where
    args :: [String]
args = ShowS -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ShowS
trim ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ String -> [String]
splitCommas String
arg
    (String
arg, String
r) = (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char
')' Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==) String
s
    sql :: ShowS
sql (Char
')':String
q) = String
q
    sql String
_ = ShowS
forall a. HasCallStack => String -> a
error String
"pgSQL: unterminated argument list" 
  pqf QueryFlags
f String
q = (QueryFlags
f, String
q)

qqQuery :: String -> TH.ExpQ
qqQuery :: String -> ExpQ
qqQuery = (QueryFlags -> String -> ExpQ) -> (QueryFlags, String) -> ExpQ
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry QueryFlags -> String -> ExpQ
makePGQuery ((QueryFlags, String) -> ExpQ)
-> (String -> (QueryFlags, String)) -> String -> ExpQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> (QueryFlags, String)
parseQueryFlags

qqTop :: Bool -> String -> TH.DecsQ
qqTop :: Bool -> String -> DecsQ
qqTop Bool
True (Char
'!':String
sql) = Bool -> String -> DecsQ
qqTop Bool
False String
sql
qqTop Bool
err String
sql = do
  Either PGError (Int, [PGValues])
r <- IO (Either PGError (Int, [PGValues]))
-> Q (Either PGError (Int, [PGValues]))
forall a. IO a -> Q a
TH.runIO (IO (Either PGError (Int, [PGValues]))
 -> Q (Either PGError (Int, [PGValues])))
-> IO (Either PGError (Int, [PGValues]))
-> Q (Either PGError (Int, [PGValues]))
forall a b. (a -> b) -> a -> b
$ IO (Int, [PGValues]) -> IO (Either PGError (Int, [PGValues]))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Int, [PGValues]) -> IO (Either PGError (Int, [PGValues])))
-> IO (Int, [PGValues]) -> IO (Either PGError (Int, [PGValues]))
forall a b. (a -> b) -> a -> b
$ (PGConnection -> IO (Int, [PGValues])) -> IO (Int, [PGValues])
forall a. (PGConnection -> IO a) -> IO a
withTPGConnection ((PGConnection -> IO (Int, [PGValues])) -> IO (Int, [PGValues]))
-> (PGConnection -> IO (Int, [PGValues])) -> IO (Int, [PGValues])
forall a b. (a -> b) -> a -> b
$ \PGConnection
c ->
    PGConnection -> ByteString -> IO (Int, [PGValues])
pgSimpleQuery PGConnection
c (String -> ByteString
BSLU.fromString String
sql)
  (PGError -> Q ())
-> ((Int, [PGValues]) -> Q ())
-> Either PGError (Int, [PGValues])
-> Q ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ((if Bool
err then String -> Q ()
TH.reportError else String -> Q ()
TH.reportWarning) (String -> Q ()) -> (PGError -> String) -> PGError -> Q ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PGError -> String
forall a. Show a => a -> String
show :: PGError -> String)) (Q () -> (Int, [PGValues]) -> Q ()
forall a b. a -> b -> a
const (Q () -> (Int, [PGValues]) -> Q ())
-> Q () -> (Int, [PGValues]) -> Q ()
forall a b. (a -> b) -> a -> b
$ () -> Q ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) Either PGError (Int, [PGValues])
r
  [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return []

-- |A quasi-quoter for PGSQL queries.
--
-- Used in expression context, it may contain any SQL statement @[pgSQL|SELECT ...|]@.
-- The statement may contain PostgreSQL-style placeholders (@$1@, @$2@, ...) or in-line placeholders (@${1+1}@) containing any valid Haskell expression (except @{}@).
-- It will be replaced by a 'PGQuery' object that can be used to perform the SQL statement.
-- If there are more @$N@ placeholders than expressions, it will instead be a function accepting the additional parameters and returning a 'PGQuery'.
-- 
-- Ideally, this mimics postgres' SQL parsing, so that placeholders and expressions will only be expanded when they are in valid positions (i.e., not inside quoted strings).
-- Since @${@ is not valid SQL otherwise, there should be no need to escape it.
--
-- The statement may start with one of more special flags affecting the interpretation:
--
-- [@?@] To disable nullability inference, treating all result values as nullable, thus returning 'Maybe' values regardless of inferred nullability. This makes unexpected NULL errors impossible.
-- [@!@] To disable nullability inference, treating all result values as /not/ nullable, thus only returning 'Maybe' where requested. This is makes unexpected NULL errors more likely.
-- [@$@] To create a 'PGPreparedQuery' (using placeholder parameters) rather than the default 'PGSimpleQuery' (using literal substitution).
-- [@$(type,...)@] To specify specific types for a prepared query (see <http://www.postgresql.org/docs/current/static/sql-prepare.html> for details), rather than inferring parameter types by default.
-- [@#@] Only do literal @${}@ substitution using 'pgSubstituteLiterals' and return a string, not a query.
-- 
-- 'pgSQL' can also be used at the top-level to execute SQL statements at compile-time (without any parameters and ignoring results).
-- Here the query can only be prefixed with @!@ to make errors non-fatal.
--
-- If you want to construct queries out of string variables rather than quasi-quoted strings, you can use the lower-level 'makePGQuery' instead.
pgSQL :: QuasiQuoter
pgSQL :: QuasiQuoter
pgSQL = QuasiQuoter :: (String -> ExpQ)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> DecsQ)
-> QuasiQuoter
QuasiQuoter
  { quoteExp :: String -> ExpQ
quoteExp = String -> ExpQ
qqQuery
  , quoteType :: String -> Q Type
quoteType = Q Type -> String -> Q Type
forall a b. a -> b -> a
const (Q Type -> String -> Q Type) -> Q Type -> String -> Q Type
forall a b. (a -> b) -> a -> b
$ String -> Q Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"pgSQL not supported in types"
  , quotePat :: String -> Q Pat
quotePat = Q Pat -> String -> Q Pat
forall a b. a -> b -> a
const (Q Pat -> String -> Q Pat) -> Q Pat -> String -> Q Pat
forall a b. (a -> b) -> a -> b
$ String -> Q Pat
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"pgSQL not supported in patterns"
  , quoteDec :: String -> DecsQ
quoteDec = Bool -> String -> DecsQ
qqTop Bool
True
  }