multipool-postgresql-simple-0.1.0.0: src/Data/MultiPool/PostgreSQLSimple.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE RecordWildCards #-}

module Data.MultiPool.PostgreSQLSimple
    ( MultiPoolBackend(..)
    ) where

import Control.Exception
import Control.Monad.IO.Unlift
import Control.Monad.Reader
import Data.Pool
import Data.MultiPool
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import Database.PostgreSQL.Simple

withResource' :: MonadUnliftIO m => Pool a -> (a -> m b) -> m b
withResource' pool action =
  withRunInIO $ \io ->
    liftIO $ withResource pool $ \a ->
      io $ action a

instance MonadUnliftIO m => MultiPoolBackend m Connection where
  type MasterConnection Connection = Connection
  type ReplicaConnection Connection = Connection
  type ReplicaIdentifier Connection = InstanceName Connection

  runWriteAny b m = runWrite b () m
  runWrite b () m = withResource' (multiPoolMaster b) $ runReaderT m

  runReadMaster b () m = runReadAnyMaster b m
  runReadAnyMaster b m = withResource' (multiPoolMaster b) $ runReaderT m
  runReadAny b m = do
    mident <- liftIO $ multiPoolAnyReplicaSelector b b
    case mident of
      Nothing -> runReadAnyMaster b m
      Just ident -> runRead b ident m
  runRead b ident m = case HM.lookup ident (multiPoolReplica b) of
    Nothing -> throw (InstanceDoesNotExist ident)
    Just repl -> withResource' repl $ runReaderT m