1 {-# LANGUAGE FlexibleInstances, FlexibleContexts, 2 MultiParamTypeClasses, ConstraintKinds, TypeFamilies, 3 OverloadedStrings #-}module Service.DB 4 ( DBPool 5 , DBConn 6 , initDB 7 , finiDB 8 , withDB 9 , MonadDB 10 , DBM 11 , runDBM 12 , liftDBM 13 , dbTryJust 14 , dbRunQuery 15 , dbExecute 16 , dbExecuteSimple 17 , dbExecute1 18 , dbExecute1' 19 , dbExecute_ 20 , dbQuery 21 , dbQuery1 22 , dbQuery1' 23 , dbTransaction 24 , dbTransaction' 25 , runDBConnection 26 , useTDB 27 , runTDB 28 , mapQuery2 29 , mapPrepQuery 30 , mapRunPrepQuery 31 , mapRunPrepQuery1 32 -- FIXME: added for tests 33 , loadPGDatabase 34 , pgConnect 35 ) where 36 37 import Control.Exception (tryJust, bracket) 38 import Control.Monad (unless) 39 import Control.Monad.IO.Class (MonadIO) 40 import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_) 41 import Control.Monad.Trans.Reader (ReaderT(..)) 42 import qualified Data.ByteString.Lazy as BSL 43 import qualified Data.ByteString.Char8 as BS 44 import Data.IORef (IORef, newIORef, atomicModifyIORef') 45 import Data.Maybe (fromMaybe, isJust, fromJust) 46 import Data.Pool (Pool, withResource, createPool, destroyAllResources) 47 import qualified Database.PostgreSQL.Simple as PGSimple 48 import Database.PostgreSQL.Typed.Protocol 49 import Database.PostgreSQL.Typed.Query 50 import Database.PostgreSQL.Typed.TH (withTPGConnection, useTPGDatabase) 51 import Database.PostgreSQL.Typed.Types (PGValue) 52 import qualified Language.Haskell.TH as TH 53 import Network (PortID(..)) 54 import System.IO.Unsafe (unsafePerformIO) 55 56 import Has 57 import qualified Store.Config as C 58 59 confPGDatabase :: C.Config -> PGDatabase 60 confPGDatabase conf = defaultPGDatabase 61 { pgDBHost = fromMaybe "localhost" host 62 , pgDBPort = if isJust host 63 then PortNumber (maybe 5432 fromInteger $ conf C.! "port") 64 else UnixSocket (fromJust $ conf C.! "sock") 65 , pgDBName = fromMaybe user $ conf C.! "db" 66 , pgDBUser = user 67 , pgDBPass = fromMaybe "" $ conf C.! "pass" 68 , pgDBDebug = fromMaybe False $ conf C.! "debug" 69 } 70 where 71 host = conf C.! "host" 72 user = conf C.! "user" 73 74 75 data DBPool = DBPool (Pool PGConnection) (Pool PGSimple.Connection) 76 type DBConn = PGConnection 77 78 initDB :: C.Config -> IO DBPool 79 initDB conf = 80 DBPool 81 <$> createPool' (pgConnect db) pgDisconnect 82 <*> createPool' (PGSimple.connect simpleConnInfo) (PGSimple.close) 83 where 84 createPool' :: IO a -> (a -> IO ()) -> IO (Pool a) 85 createPool' get release = 86 createPool get release stripes (fromInteger idle) conn 87 db = confPGDatabase conf 88 simpleConnInfo = PGSimple.defaultConnectInfo 89 { PGSimple.connectHost = pgDBHost db 90 , PGSimple.connectPort = case pgDBPort db of 91 PortNumber x -> fromIntegral x -- x is opaque 92 UnixSocket _ -> 93 PGSimple.connectPort PGSimple.defaultConnectInfo 94 Service _ -> 95 PGSimple.connectPort PGSimple.defaultConnectInfo 96 , PGSimple.connectUser = BS.unpack (pgDBUser db) 97 , PGSimple.connectPassword = BS.unpack (pgDBPass db) 98 , PGSimple.connectDatabase = BS.unpack (pgDBName db) 99 } 100 stripes = fromMaybe 1 $ conf C.! "stripes" 101 idle = fromMaybe 300 $ conf C.! "idle" 102 conn = fromMaybe 16 $ conf C.! "maxconn" 103 104 finiDB :: DBPool -> IO () 105 finiDB (DBPool p p') = do 106 -- Different types -> no 'travers'ing 107 destroyAllResources p 108 destroyAllResources p' 109 110 withDB :: DBPool -> (DBConn -> IO a) -> IO a 111 withDB (DBPool p _) = withResource p 112 113 type MonadDB c m = (MonadIO m, MonadHas DBConn c m) 114 115 {-# INLINE liftDB #-} 116 liftDB :: MonadDB c m => (PGConnection -> IO a) -> m a 117 liftDB = focusIO 118 119 type DBM a = ReaderT PGConnection IO a 120 121 runDBM :: DBPool -> DBM a -> IO a 122 runDBM p = withDB p . runReaderT 123 124 liftDBM :: MonadDB c m => DBM a -> m a 125 liftDBM = liftDB . runReaderT 126 127 -- |Combination of 'liftDBM' and lifted 'tryJust' 128 dbTryJust :: MonadDB c m => (PGError -> Maybe e) -> DBM a -> m (Either e a) 129 dbTryJust err q = liftDB $ tryJust err . runReaderT q 130 131 dbRunQuery :: (MonadDB c m, PGQuery q a) => q -> m (Int, [a]) 132 dbRunQuery q = liftDB $ \c -> pgRunQuery c q 133 134 dbExecute :: (MonadDB c m, PGQuery q ()) => q -> m Int 135 dbExecute q = liftDB $ \c -> pgExecute c q 136 137 dbExecuteSimple :: MonadDB c m => PGSimpleQuery () -> m Int 138 dbExecuteSimple = dbExecute 139 140 dbExecute1 :: (MonadDB c m, PGQuery q (), Show q) => q -> m Bool 141 dbExecute1 q = do 142 r <- dbExecute q 143 case r of 144 0 -> return False 145 1 -> return True 146 _ -> fail $ "pgExecute1 " ++ show q ++ ": " ++ show r ++ " rows" 147 148 dbExecute1' :: (MonadDB c m, PGQuery q (), Show q) => q -> m () 149 dbExecute1' q = do 150 r <- dbExecute1 q 151 unless r $ fail $ "pgExecute1' " ++ show q ++ ": failed" 152 153 dbExecute_ :: (MonadDB c m) => BSL.ByteString -> m () 154 dbExecute_ q = liftDB $ \c -> pgSimpleQueries_ c q 155 156 dbQuery :: (MonadDB c m, PGQuery q a) => q -> m [a] 157 dbQuery q = liftDB $ \c -> pgQuery c q 158 159 dbQuery1 :: (MonadDB c m, PGQuery q a, Show q) => q -> m (Maybe a) 160 dbQuery1 q = do 161 r <- dbQuery q 162 case r of 163 [] -> return Nothing 164 [x] -> return $ Just x 165 _ -> fail $ "pgQuery1 " ++ show q ++ ": too many results" 166 167 dbQuery1' :: (MonadDB c m, PGQuery q a, Show q) => q -> m a 168 dbQuery1' q = maybe (fail $ "pgQuery1' " ++ show q ++ ": no results") return =<< dbQuery1 q 169 170 dbTransaction :: MonadDB c m => DBM a -> m a 171 dbTransaction f = liftDB $ \c -> pgTransaction c (runReaderT f c) 172 173 dbTransaction' :: (MonadBaseControl IO m, MonadDB c m) => m a -> m a 174 dbTransaction' f = do 175 c <- peek 176 liftBaseOp_ (pgTransaction c) f 177 178 -- For connections outside runtime: 179 180 loadPGDatabase :: IO PGDatabase 181 loadPGDatabase = confPGDatabase . C.get "db" <$> C.load "databrary.conf" 182 183 runDBConnection :: DBM a -> IO a 184 runDBConnection f = bracket 185 (pgConnect =<< loadPGDatabase) 186 pgDisconnect 187 (runReaderT f) 188 189 loadTDB :: TH.DecsQ 190 loadTDB = do 191 database <- TH.runIO loadPGDatabase 192 useTPGDatabase database 193 194 {-# NOINLINE usedTDB #-} 195 usedTDB :: IORef Bool 196 usedTDB = unsafePerformIO $ newIORef False 197 useTDB :: TH.DecsQ 198 useTDB = do 199 d <- TH.runIO $ atomicModifyIORef' usedTDB ((,) True) 200 if d 201 then return [] 202 else loadTDB 203 204 runTDB :: DBM a -> TH.Q a 205 runTDB f = do 206 _ <- useTDB 207 TH.runIO $ withTPGConnection $ runReaderT f 208 209 -- Temporary helpers while removing postgresql-typed, remove after complete 210 mapQuery2 :: BS.ByteString -> ([PGValue] -> a) -> PGSimpleQuery a -- mapQuery is same as mapQuery2, both will be deleted 211 mapQuery2 qry mkResult = 212 fmap mkResult (rawPGSimpleQuery qry) 213 214 mapPrepQuery :: (BS.ByteString, [PGValue]) -> ([PGValue] -> a) -> PGPreparedQuery a 215 mapPrepQuery (qry, params) mkResult = 216 fmap mkResult (rawPGPreparedQuery qry params) 217 218 mapRunPrepQuery :: (MonadDB c m) => (BS.ByteString, [PGValue], [Bool]) -> ([PGValue] -> a) -> m [a] 219 mapRunPrepQuery (qry, params, bc) mkResult = do 220 rows <- liftDB $ \c -> snd <$> pgPreparedQuery c qry [] params bc 221 pure (fmap mkResult rows) 222 223 mapRunPrepQuery1 :: (MonadDB c m) => (BS.ByteString, [PGValue], [Bool]) -> ([PGValue] -> a) -> m (Maybe a) 224 mapRunPrepQuery1 args@(q, _, _) mkResult = do 225 rows <- mapRunPrepQuery args mkResult 226 case rows of 227 [] -> return Nothing 228 [x] -> return $ Just x 229 _ -> fail $ "pgQuery1 " ++ show q ++ ": too many results"