module Database.PostgreSQL.ORM.DBSelect (
DBSelect(..), FromClause(..)
, dbSelectParams, dbSelect
, Cursor(..), curSelect, curNext
, dbFold, dbFoldM, dbFoldM_
, dbCollect
, renderDBSelect, buildDBSelect
, emptyDBSelect, expressionDBSelect
, modelDBSelect
, dbJoin, dbJoinModels
, dbProject, dbProject'
, dbNest, dbChain
, addWhere_, addWhere, setOrderBy, setLimit, setOffset, addExpression
) where
import Control.Monad.IO.Class
import Blaze.ByteString.Builder
import Blaze.ByteString.Builder.Char.Utf8 (fromChar)
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as S8
import Data.Functor
import Data.Monoid
import Data.String
import Data.IORef
import Database.PostgreSQL.Simple
import Database.PostgreSQL.Simple.Internal
import Database.PostgreSQL.Simple.Types
import GHC.Generics
import Database.PostgreSQL.Escape
import Database.PostgreSQL.ORM.Model
data FromClause = FromModel {
fcVerbatim :: !Query
, fcCanonical :: !S.ByteString
}
| FromJoin {
fcLeft :: !FromClause
, fcJoinOp :: !Query
, fcRight :: !FromClause
, fcOnClause :: !Query
, fcCanonical :: !S.ByteString
}
deriving Show
nullFrom :: FromClause -> Bool
nullFrom (FromModel q _) | qNull q = True
nullFrom _ = False
data DBSelect a = DBSelect {
selWith :: !Query
, selSelectKeyword :: !Query
, selFields :: Query
, selFrom :: !FromClause
, selWhereKeyword :: !Query
, selWhere :: !Query
, selGroupBy :: !Query
, selHaving :: !Query
, selOrderBy :: !Query
, selLimit :: !Query
, selOffset :: !Query
} deriving (Generic)
instance Show (DBSelect a) where
show = S8.unpack . fromQuery . renderDBSelect
space :: Builder
space = fromChar ' '
qNull :: Query -> Bool
qNull = S.null . fromQuery
qBuilder :: Query -> Builder
qBuilder = fromByteString . fromQuery
toQuery :: Builder -> Query
toQuery = Query . toByteString
buildFromClause :: FromClause -> Builder
buildFromClause (FromModel q _) | qNull q = mempty
buildFromClause cl0 = fromByteString " FROM " <> go cl0
where go (FromModel q _) = qBuilder q
go (FromJoin left joinkw right onClause _) = mconcat [
fromChar '(', go left, space, qBuilder joinkw, space, go right
, if qNull onClause then mempty else space <> qBuilder onClause
, fromChar ')' ]
class GDBS f where
gdbsDefault :: f p
gdbsQuery :: f p -> Builder
instance GDBS (K1 i Query) where
gdbsDefault = K1 (Query S.empty)
gdbsQuery (K1 q) | qNull q = mempty
| otherwise = space <> qBuilder q
instance GDBS (K1 i FromClause) where
gdbsDefault = K1 (FromModel "" "")
gdbsQuery (K1 fc) = buildFromClause fc
instance (GDBS a, GDBS b) => GDBS (a :*: b) where
gdbsDefault = gdbsDefault :*: gdbsDefault
gdbsQuery (a :*: b) = gdbsQuery a <> gdbsQuery b
instance (GDBS f) => GDBS (M1 i c f) where
gdbsDefault = M1 gdbsDefault
gdbsQuery = gdbsQuery . unM1
emptyDBSelect :: DBSelect a
emptyDBSelect = (to gdbsDefault) { selSelectKeyword = fromString "SELECT" }
expressionDBSelect :: (Model r) => Query -> DBSelect r
expressionDBSelect q = emptyDBSelect { selFields = q }
buildDBSelect :: DBSelect a -> Builder
buildDBSelect dbs = gdbsQuery $ from dbs
renderDBSelect :: DBSelect a -> Query
renderDBSelect = Query . S.tail . toByteString . buildDBSelect
catQueries :: Query -> Query -> Query -> Query
catQueries left delim right
| qNull left = right
| qNull right = left
| otherwise = Query $ S.concat $ map fromQuery [left, delim, right]
addWhere_ :: Query -> DBSelect a -> DBSelect a
addWhere_ q dbs
| qNull q = dbs
| otherwise = dbs { selWhereKeyword = "WHERE"
, selWhere = catQueries (selWhere dbs) " AND " q }
addWhere :: (ToRow p) => Query -> p -> DBSelect a -> DBSelect a
addWhere q p dbs
| qNull q = dbs
| otherwise = dbs {
selWhereKeyword = "WHERE"
, selWhere = if qNull $ selWhere dbs
then toQuery clause
else toQuery $ qBuilder (selWhere dbs) <>
fromByteString " AND " <> clause
}
where clause = mconcat [fromChar '(', buildSql q p, fromChar ')']
setOrderBy :: Query -> DBSelect a -> DBSelect a
setOrderBy (Query ob) dbs = dbs { selOrderBy = Query $ "ORDER BY " <> ob }
setLimit :: Int -> DBSelect a -> DBSelect a
setLimit i dbs = dbs { selLimit = fmtSql "LIMIT ?" (Only i) }
setOffset :: Int -> DBSelect a -> DBSelect a
setOffset i dbs = dbs { selOffset = fmtSql "OFFSET ?" (Only i) }
addExpression :: (Model r) => Query -> DBSelect a -> DBSelect (a :. r)
addExpression q dbs = dbs {
selFields = if qNull $ selFields dbs then q
else Query $ S.concat $ map fromQuery [selFields dbs, ", ", q]
}
modelDBSelect :: forall a. (Model a) => DBSelect a
modelDBSelect = r
where mi = modelIdentifiers :: ModelIdentifiers a
r = emptyDBSelect {
selFields = Query $ S.intercalate ", " $ modelQColumns mi
, selFrom = FromModel (Query $ modelQTable mi) (modelQTable mi)
}
dbSelectParams :: (Model a, ToRow p) => DBSelect a -> Connection -> p -> IO [a]
dbSelectParams dbs = \c p -> map lookupRow <$> query c q p
where
q = renderDBSelect dbs
dbSelect :: (Model a) => Connection -> DBSelect a -> IO [a]
dbSelect c dbs = map lookupRow <$> query_ c q
where
q = renderDBSelect dbs
data Cursor a = Cursor { curConn :: !Connection
, curName :: !Query
, curChunkSize :: !Query
, curCache :: IORef [a] }
curSelect :: Model a => Connection -> DBSelect a -> IO (Cursor a)
curSelect c dbs = do
name <- newTempName c
execute_ c $
mconcat [ "DECLARE ", name, " NO SCROLL CURSOR FOR ", q ]
cacheRef <- newIORef []
return $ Cursor c name "256" cacheRef
where q = renderDBSelect dbs
curNext :: Model a => Cursor a -> IO (Maybe a)
curNext Cursor{..} = do
cache <- readIORef curCache
case cache of
x:xs -> do
writeIORef curCache xs
return $ Just x
[] -> do
res <- map lookupRow <$> query_ curConn (mconcat
[ "FETCH FORWARD ", curChunkSize, " FROM ", curName])
case res of
[] -> return Nothing
x:xs -> do
writeIORef curCache xs
return $ Just x
dbFold :: Model model
=> Connection -> (b -> model -> b) -> b -> DBSelect model -> IO b
dbFold c act initial dbs = do
cur <- curSelect c dbs
go cur initial
where go cur accm = do
mres <- curNext cur
case mres of
Nothing -> return accm
Just res -> go cur (act accm res)
dbFoldM :: (MonadIO m, Model model)
=> Connection -> (b -> model -> m b) -> b -> DBSelect model -> m b
dbFoldM c act initial dbs = do
cur <- liftIO $ curSelect c dbs
go cur initial
where go cur accm = do
mres <- liftIO $ curNext cur
case mres of
Nothing -> return accm
Just res -> act accm res >>= go cur
dbFoldM_ :: (MonadIO m, Model model)
=> Connection -> (model -> m ()) -> DBSelect model -> m ()
dbFoldM_ c act dbs = dbFoldM c (const act) () dbs
dbCollect :: (Model a, Model b)
=> Connection -> DBSelect (a :. b) -> IO [(a, [b])]
dbCollect c ab = dbFold c group [] ab
where
group :: (Model a, Model b) => [(a, [b])] -> (a :. b) -> [(a, [b])]
group [] (a :. b) = [(a, [b])]
group ls@(l:_) (a :. b) | primaryKey a /= primaryKey (fst l) = (a, [b]):ls
group (l:ls) (_ :. b) = (fst l, b:(snd l)):ls
dbJoin :: forall a b.
(Model a, Model b) =>
DBSelect a
-> Query
-> DBSelect b
-> Query
-> DBSelect (a :. b)
dbJoin left joinOp right onClause = addWhere_ (selWhere left) right {
selFields = Query $ S.concat [fromQuery $ selFields left, ", ",
fromQuery $ selFields right]
, selFrom = newfrom
}
where idab = modelIdentifiers :: ModelIdentifiers (a :. b)
newfrom | nullFrom $ selFrom right = selFrom left
| nullFrom $ selFrom left = selFrom right
| otherwise = FromJoin (selFrom left) joinOp (selFrom right)
onClause (modelQTable idab)
dbJoinModels :: (Model a, Model b) =>
Query
-> Query
-> DBSelect (a :. b)
dbJoinModels kw on = dbJoin modelDBSelect kw modelDBSelect on
dbProject :: forall a something_containing_a.
(Model a) => DBSelect something_containing_a -> DBSelect a
dbProject dbs = r
where sela = modelDBSelect :: DBSelect a
r = dbs { selFields = selFields sela }
dbProject' :: forall a something_containing_a.
(Model a) => DBSelect something_containing_a -> DBSelect a
dbProject' dbs = r
where sela = modelDBSelect :: DBSelect a
ida = modelIdentifiers :: ModelIdentifiers a
Just mq = modelQualifier ida
q = toQuery $ fromChar '(' <>
buildDBSelect dbs { selFields = selFields sela } <>
fromByteString ") AS " <> fromByteString mq
r = sela { selFrom = FromModel q $ modelQTable ida }
mergeFromClauses :: S.ByteString -> FromClause -> FromClause -> FromClause
mergeFromClauses canon left right =
case go left of
(fc, 1) -> fc
(_, 0) -> error $ "mergeFromClauses could not find " ++ show canon
(_, _) -> error $ "mergeFromClauses found duplicate " ++ show canon
where go fc | fcCanonical fc == canon = (right, 1 :: Int)
go (FromJoin l op r on ffc) =
case (go l, go r) of
((lfc, ln), (rfc, rn)) -> (FromJoin lfc op rfc on ffc, ln + rn)
go fc = (fc, 0)
dbNest :: forall a b c. (Model a, Model b) =>
DBSelect (a :. b) -> DBSelect (b :. c) -> DBSelect (a :. b :. c)
dbNest left right = addWhere_ (selWhere left) right {
selFields = fields
, selFrom = mergeFromClauses nameb (selFrom left) (selFrom right)
}
where nameb = modelQTable (modelIdentifiers :: ModelIdentifiers b)
acols = modelQColumns (modelIdentifiers :: ModelIdentifiers a)
colcomma c r = fromByteString c <> fromByteString ", " <> r
fields = toQuery $ foldr colcomma (qBuilder $ selFields right)
acols
dbChain :: (Model a, Model b, Model c) =>
DBSelect (a :. b) -> DBSelect (b :. c) -> DBSelect (a :. c)
dbChain left right = dbProject $ dbNest left right