{-# 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 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
pgRunQuery :: PGConnection -> q -> IO (Int, [a])
unsafeModifyQuery :: q -> (BS.ByteString -> BS.ByteString) -> q
getQueryString :: PGTypeEnv -> q -> BS.ByteString
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
instance PGQuery BS.ByteString PGValues where
pgRunQuery c sql = pgSimpleQuery c (BSL.fromStrict sql)
unsafeModifyQuery q f = f q
getQueryString _ = id
newtype SimpleQuery = SimpleQuery BS.ByteString
deriving (Show)
instance PGQuery SimpleQuery PGValues where
pgRunQuery c (SimpleQuery sql) = pgSimpleQuery c (BSL.fromStrict sql)
unsafeModifyQuery (SimpleQuery sql) f = SimpleQuery $ f sql
getQueryString _ (SimpleQuery q) = q
instance PGRawQuery SimpleQuery
data PreparedQuery = PreparedQuery BS.ByteString [OID] PGValues [Bool]
deriving (Show)
instance PGQuery PreparedQuery PGValues where
pgRunQuery c (PreparedQuery sql types bind bc) = pgPreparedQuery c sql types bind bc
unsafeModifyQuery (PreparedQuery sql types bind bc) f = PreparedQuery (f sql) types bind bc
getQueryString _ (PreparedQuery q _ _ _) = q
instance PGRawQuery PreparedQuery
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
unsafeModifyQuery (QueryParser q p) f = QueryParser (\e -> unsafeModifyQuery (q e) f) p
getQueryString e (QueryParser q _) = getQueryString e $ q e
instance Functor (QueryParser q) where
fmap f (QueryParser q p) = QueryParser q (\e -> f . p e)
instance Show q => Show (QueryParser q a) where
showsPrec p (QueryParser q _) = showParen (p > 10) $
showString "QueryParser " . showsPrec 11 (q unknownPGTypeEnv)
rawParser :: q -> QueryParser q PGValues
rawParser q = QueryParser (const q) (const id)
type PGSimpleQuery = QueryParser SimpleQuery
type PGPreparedQuery = QueryParser PreparedQuery
rawPGSimpleQuery :: BS.ByteString -> PGSimpleQuery PGValues
rawPGSimpleQuery = rawParser . SimpleQuery
instance IsString (PGSimpleQuery PGValues) where
fromString = rawPGSimpleQuery . fromString
instance IsString (PGSimpleQuery ()) where
fromString = void . rawPGSimpleQuery . fromString
rawPGPreparedQuery :: BS.ByteString -> 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 = sst (1 :: Int) . sqlTokens where
sst n (SQLExpr e : l) = (('$':show n) ++) *** (e :) $ sst (succ n) l
sst n (t : l) = first (show t ++) $ sst n l
sst _ [] = ("", [])
sqlSubstitute :: String -> [TH.Exp] -> TH.Exp
sqlSubstitute sql exprl = TH.AppE (TH.VarE 'BS.concat) $ TH.ListE $ map sst $ sqlTokens 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)"
sst (SQLParam n) = expr n
sst t = TH.VarE 'fromString `TH.AppE` TH.LitE (TH.StringL $ show t)
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 :: Maybe Bool
, flagPrepare :: Maybe [String]
}
simpleQueryFlags :: QueryFlags
simpleQueryFlags = QueryFlags True Nothing Nothing
newName :: Char -> BS.ByteString -> TH.Q TH.Name
newName pre = TH.newName . ('_':) . (pre:) . filter (\c -> isAlphaNum c || c == '_') . BSC.unpack
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 (fromString sqlp) (fromMaybe [] prep) (isNothing nulls)
when (length pt < length exprs) $ fail "Not all expression placeholders were recognized by PostgreSQL"
e <- TH.newName "_tenv"
l <- TH.newName "l"
(vars, vals) <- mapAndUnzipM (\t -> do
v <- 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 <- newName 'c' $ tpgValueName t
return
( TH.VarP v
, tpgTypeDecoder (Fold.and nulls) 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] (maybe
(TH.ConE 'SimpleQuery
`TH.AppE` sqlSubstitute sqlp vals)
(\p -> TH.ConE 'PreparedQuery
`TH.AppE` (TH.VarE 'fromString `TH.AppE` TH.LitE (TH.StringL sqlp))
`TH.AppE` TH.ListE (map (TH.LitE . TH.IntegerL . toInteger . tpgValueTypeOID . snd) $ zip p pt)
`TH.AppE` TH.ListE vals
`TH.AppE` TH.ListE
#ifdef VERSION_postgresql_binary
bins
#else
[]
#endif
)
prep)
`TH.AppE` TH.LamE [TH.VarP e, TH.VarP l] (TH.CaseE (TH.VarE l)
[ TH.Match (TH.ListP pats) (TH.NormalB $ case conv of
[x] -> x
_ -> TH.TupE
#if MIN_VERSION_template_haskell(2,16,0)
$ map Just
#endif
conv) []
, TH.Match TH.WildP (TH.NormalB $ TH.VarE 'error `TH.AppE` TH.LitE (TH.StringL "pgSQL: result arity mismatch")) []
]))
<$> mapM parse exprs
where
(sqlp, exprs) = sqlPlaceholders sqle
parse e = either (fail . (++) ("Failed to parse expression {" ++ e ++ "}: ")) return $ parseExp e
parseQueryFlags :: String -> (QueryFlags, String)
parseQueryFlags = pqf simpleQueryFlags where
pqf f@QueryFlags{ flagQuery = True, flagPrepare = Nothing } ('#':q) = pqf f{ flagQuery = False } q
pqf f@QueryFlags{ flagQuery = True, flagNullable = Nothing } ('?':q) = pqf f{ flagNullable = Just True } q
pqf f@QueryFlags{ flagQuery = True, flagNullable = Nothing } ('!':q) = pqf f{ flagNullable = Just False } q
pqf f@QueryFlags{ flagQuery = True, flagPrepare = Nothing } ('$':q) = pqf f{ flagPrepare = Just [] } q
pqf f@QueryFlags{ flagQuery = True, flagPrepare = Just [] } ('(':s) = pqf f{ flagPrepare = Just args } (sql r) where
args = map trim $ splitCommas arg
(arg, r) = break (')' ==) s
sql (')':q) = q
sql _ = error "pgSQL: unterminated argument list"
pqf f q = (f, q)
qqQuery :: String -> TH.ExpQ
qqQuery = uncurry makePGQuery . parseQueryFlags
qqTop :: Bool -> String -> TH.DecsQ
qqTop True ('!':sql) = qqTop False sql
qqTop err sql = do
r <- TH.runIO $ try $ withTPGConnection $ \c ->
pgSimpleQuery c (fromString sql)
either ((if err then TH.reportError else TH.reportWarning) . (show :: PGError -> String)) (const $ return ()) r
return []
pgSQL :: QuasiQuoter
pgSQL = QuasiQuoter
{ quoteExp = qqQuery
, quoteType = const $ fail "pgSQL not supported in types"
, quotePat = const $ fail "pgSQL not supported in patterns"
, quoteDec = qqTop True
}