module Database.PostgreSQL.Typed.Query
( PGQuery(..)
, PGSimpleQuery
, PGPreparedQuery
, rawPGSimpleQuery
, rawPGPreparedQuery
, QueryFlags(..)
, simpleQueryFlags
, parseQueryFlags
, makePGQuery
, pgSQL
, pgExecute
, pgQuery
, pgLazyQuery
) where
import Control.Applicative ((<$>))
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)
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.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])
unsafeModifyQuery :: q -> (BS.ByteString -> BS.ByteString) -> q
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
newtype SimpleQuery = SimpleQuery BS.ByteString
instance PGQuery SimpleQuery PGValues where
pgRunQuery c (SimpleQuery sql) = pgSimpleQuery c (BSL.fromStrict sql)
unsafeModifyQuery (SimpleQuery sql) f = SimpleQuery $ f sql
instance PGRawQuery SimpleQuery
data PreparedQuery = PreparedQuery BS.ByteString [OID] PGValues [Bool]
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
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
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 :: 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 = ssl 1 . sqlSplitExprs where
ssl :: Int -> SQLSplit String True -> (String, [String])
ssl n (SQLLiteral s l) = first (s ++) $ ssp n l
ssl _ SQLSplitEnd = ("", [])
ssp :: Int -> SQLSplit String False -> (String, [String])
ssp n (SQLPlaceholder e l) = (('$':show n) ++) *** (e :) $ ssl (succ n) l
ssp _ SQLSplitEnd = ("", [])
sqlSubstitute :: String -> [TH.Exp] -> TH.Exp
sqlSubstitute sql exprl = TH.AppE (TH.VarE 'BS.concat) $ TH.ListE $ ssl $ sqlSplitParams 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 '$$'"
ssl (SQLLiteral s l) = TH.VarE 'fromString `TH.AppE` stringE s : ssp l
ssl SQLSplitEnd = []
ssp (SQLPlaceholder n l) = expr n : ssl l
ssp SQLSplitEnd = []
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
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; literal occurrences of '${' may need to be escaped with '$${'"
e <- TH.newName "_tenv"
(vars, vals) <- mapAndUnzipM (\t -> do
v <- TH.newName $ 'p':BSC.unpack (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':BSC.unpack (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` stringE 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 USE_BINARY
bins
#else
[]
#endif
)
prep)
`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
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
}