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