module Database.PostgreSQL.Typed.Query
( PGQuery(..)
, PGSimpleQuery
, PGPreparedQuery
, rawPGSimpleQuery
, rawPGPreparedQuery
, QueryFlags(..)
, simpleQueryFlags
, makePGQuery
, pgSQL
, pgExecute
, pgQuery
, pgLazyQuery
) where
import Control.Applicative ((<$>))
import Control.Arrow ((***), first, second)
import Control.Exception (try)
import Control.Monad (when, mapAndUnzipM)
import Data.Array (listArray, (!), inRange)
import Data.Char (isDigit, isSpace)
import Data.List (dropWhileEnd)
import Data.Maybe (fromMaybe, isNothing)
import Data.Word (Word32)
import Language.Haskell.Meta.Parse (parseExp)
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Quote (QuasiQuoter(..))
import Numeric (readDec)
import Database.PostgreSQL.Typed.Internal
import Database.PostgreSQL.Typed.Types
import Database.PostgreSQL.Typed.Dynamic
import Database.PostgreSQL.Typed.Protocol
import Database.PostgreSQL.Typed.TH
class PGQuery q a | q -> a where
pgRunQuery :: PGConnection -> q -> IO (Int, [a])
class PGQuery q PGValues => PGRawQuery q
pgExecute :: PGQuery q () => PGConnection -> q -> IO Int
pgExecute c q = fst <$> pgRunQuery c q
pgQuery :: PGQuery q a => PGConnection -> q -> IO [a]
pgQuery c q = snd <$> pgRunQuery c q
data SimpleQuery = SimpleQuery String
instance PGQuery SimpleQuery PGValues where
pgRunQuery c (SimpleQuery sql) = pgSimpleQuery c sql
instance PGRawQuery SimpleQuery where
data PreparedQuery = PreparedQuery String [OID] PGValues [Bool]
instance PGQuery PreparedQuery PGValues where
pgRunQuery c (PreparedQuery sql types bind bc) = pgPreparedQuery c sql types bind bc
instance PGRawQuery PreparedQuery where
data QueryParser q a = QueryParser (PGTypeEnv -> q) (PGTypeEnv -> PGValues -> a)
instance PGRawQuery q => PGQuery (QueryParser q a) a where
pgRunQuery c (QueryParser q p) = second (fmap $ p e) <$> pgRunQuery c (q e) where e = pgTypeEnv c
instance Functor (QueryParser q) where
fmap f (QueryParser q p) = QueryParser q (\e -> f . p e)
rawParser :: q -> QueryParser q PGValues
rawParser q = QueryParser (const q) (const id)
type PGSimpleQuery = QueryParser SimpleQuery
type PGPreparedQuery = QueryParser PreparedQuery
rawPGSimpleQuery :: String -> PGSimpleQuery PGValues
rawPGSimpleQuery = rawParser . SimpleQuery
rawPGPreparedQuery :: String -> PGValues -> PGPreparedQuery PGValues
rawPGPreparedQuery sql bind = rawParser $ PreparedQuery sql [] bind []
pgLazyQuery :: PGConnection -> PGPreparedQuery a -> Word32
-> IO [a]
pgLazyQuery c (QueryParser q p) count =
fmap (p e) <$> pgPreparedLazyQuery c sql types bind bc count where
e = pgTypeEnv c
PreparedQuery sql types bind bc = q e
sqlPlaceholders :: String -> (String, [String])
sqlPlaceholders = sph (1 :: Int) where
sph n ('$':'$':'{':s) = first (('$':) . ('{':)) $ sph n s
sph n ('$':'{':s)
| (e, '}':r) <- break (\c -> c == '{' || c == '}') s =
(('$':show n) ++) *** (e :) $ sph (succ n) r
| otherwise = error $ "Error parsing SQL statement: could not find end of expression: ${" ++ s
sph n (c:s) = first (c:) $ sph n s
sph _ "" = ("", [])
sqlSubstitute :: String -> [TH.Exp] -> TH.Exp
sqlSubstitute sql exprl = ss sql where
bnds = (1, length exprl)
exprs = listArray bnds exprl
expr n
| inRange bnds n = exprs ! n
| otherwise = error $ "SQL placeholder '$" ++ show n ++ "' out of range (not recognized by PostgreSQL); literal occurrences may need to be escaped with '$$'"
ss ('$':'$':d:r) | isDigit d = ['$',d] ++$ ss r
ss ('$':s@(d:_)) | isDigit d, [(n, r)] <- readDec s = expr n $++$ ss r
ss (c:r) = [c] ++$ ss r
ss "" = stringE ""
splitCommas :: String -> [String]
splitCommas = spl where
spl [] = []
spl [c] = [[c]]
spl (',':s) = "":spl s
spl (c:s) = (c:h):t where h:t = spl s
trim :: String -> String
trim = dropWhileEnd isSpace . dropWhile isSpace
data QueryFlags = QueryFlags
{ flagQuery :: Bool
, flagNullable :: Bool
, flagPrepare :: Maybe [String]
}
simpleQueryFlags :: QueryFlags
simpleQueryFlags = QueryFlags True False Nothing
makePGQuery :: QueryFlags -> String -> TH.ExpQ
makePGQuery QueryFlags{ flagQuery = False } sqle = pgSubstituteLiterals sqle
makePGQuery QueryFlags{ flagNullable = nulls, flagPrepare = prep } sqle = do
(pt, rt) <- TH.runIO $ tpgDescribe sqlp (fromMaybe [] prep) (not nulls)
when (length pt < length exprs) $ fail "Not all expression placeholders were recognized by PostgreSQL; literal occurrences of '${' may need to be escaped with '$${'"
e <- TH.newName "_tenv"
(vars, vals) <- mapAndUnzipM (\t -> do
v <- TH.newName $ 'p':tpgValueName t
return
( TH.VarP v
, tpgTypeEncoder (isNothing prep) t e `TH.AppE` TH.VarE v
)) pt
(pats, conv, bins) <- unzip3 <$> mapM (\t -> do
v <- TH.newName $ 'c':tpgValueName t
return
( TH.VarP v
, tpgTypeDecoder t e `TH.AppE` TH.VarE v
, tpgTypeBinary t e
)) rt
foldl TH.AppE (TH.LamE vars $ TH.ConE 'QueryParser
`TH.AppE` TH.LamE [TH.VarP e] (if isNothing prep
then TH.ConE 'SimpleQuery
`TH.AppE` sqlSubstitute sqlp vals
else TH.ConE 'PreparedQuery
`TH.AppE` stringE sqlp
`TH.AppE` TH.ListE (map (TH.LitE . TH.IntegerL . toInteger . tpgValueTypeOID) pt)
`TH.AppE` TH.ListE vals
`TH.AppE` TH.ListE
#ifdef USE_BINARY
bins
#else
[]
#endif
)
`TH.AppE` TH.LamE [TH.VarP e, TH.ListP pats] (TH.TupE conv))
<$> mapM parse exprs
where
(sqlp, exprs) = sqlPlaceholders sqle
parse e = either (fail . (++) ("Failed to parse expression {" ++ e ++ "}: ")) return $ parseExp e
qqQuery :: QueryFlags -> String -> TH.ExpQ
qqQuery f@QueryFlags{ flagQuery = True, flagPrepare = Nothing } ('#':q) = qqQuery f{ flagQuery = False } q
qqQuery f@QueryFlags{ flagQuery = True, flagNullable = False } ('?':q) = qqQuery f{ flagNullable = True } q
qqQuery f@QueryFlags{ flagQuery = True, flagPrepare = Nothing } ('$':q) = qqQuery f{ flagPrepare = Just [] } q
qqQuery f@QueryFlags{ flagQuery = True, flagPrepare = Just [] } ('(':s) = qqQuery f{ flagPrepare = Just args } =<< sql r where
args = map trim $ splitCommas arg
(arg, r) = break (')' ==) s
sql (')':q) = return q
sql _ = fail "pgSQL: unterminated argument list"
qqQuery f q = makePGQuery f q
qqTop :: Bool -> String -> TH.DecsQ
qqTop True ('!':sql) = qqTop False sql
qqTop err sql = do
r <- TH.runIO $ try $ withTPGConnection $ \c ->
pgSimpleQuery c sql
either ((if err then TH.reportError else TH.reportWarning) . (show :: PGError -> String)) (const $ return ()) r
return []
pgSQL :: QuasiQuoter
pgSQL = QuasiQuoter
{ quoteExp = qqQuery simpleQueryFlags
, quoteType = const $ fail "pgSQL not supported in types"
, quotePat = const $ fail "pgSQL not supported in patterns"
, quoteDec = qqTop True
}