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"