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"