Execution.hs

module Hasql.Postgres.Session.Execution where

import Hasql.Postgres.Prelude
import qualified Data.HashTable.IO as Hashtables
import qualified Database.PostgreSQL.LibPQ as PQ
import qualified Hasql.Postgres.Statement as Statement
import qualified Hasql.Postgres.TemplateConverter as TemplateConverter
import qualified Hasql.Postgres.Session.ResultProcessing as ResultProcessing


-- * Environment
-------------------------

type Env =
  (PQ.Connection, IORef Word16, Hashtables.BasicHashTable LocalKey RemoteKey)

newEnv :: PQ.Connection -> IO Env
newEnv c =
  (,,) <$> pure c <*> newIORef 0 <*> Hashtables.new


-- |
-- Local statement key.
data LocalKey =
  LocalKey !ByteString ![Word32]
  deriving (Show, Eq)

instance Hashable LocalKey where
  hashWithSalt salt (LocalKey template types) =
    hashWithSalt salt template

localKey :: ByteString -> [PQ.Oid] -> LocalKey
localKey t ol =
  LocalKey t (map oidMapper ol)
  where
    oidMapper (PQ.Oid x) = fromIntegral x


-- |
-- Remote statement key.
type RemoteKey =
  ByteString


data Error =
  UnexpectedResult Text |
  ErroneousResult Text |
  UnparsableTemplate ByteString Text |
  TransactionConflict


-- * Monad
-------------------------

newtype M r =
  M (ReaderT Env (EitherT Error IO) r)
  deriving (Functor, Applicative, Monad, MonadIO)

run :: Env -> M r -> IO (Either Error r)
run e (M m) =
  runEitherT $ runReaderT m e

throwError :: Error -> M a
throwError = M . lift . left

prepare :: ByteString -> [PQ.Oid] -> M RemoteKey
prepare s tl =
  do
    (c, counter, table) <- M $ ask
    let lk = localKey s tl
    rk <- liftIO $ Hashtables.lookup table lk
    ($ rk) $ ($ return) $ maybe $ do
      w <- liftIO $ readIORef counter
      let rk = fromString $ show w
      unitResult =<< do liftIO $ PQ.prepare c rk s (partial (not . null) tl)
      liftIO $ Hashtables.insert table lk rk
      liftIO $ writeIORef counter (succ w)
      return rk

statement :: Statement.Statement -> M (Maybe PQ.Result)
statement s =
  do
    (c, _, _) <- M $ ask
    let (template, params, preparable) = s
    convertedTemplate <-
      either (throwError . UnparsableTemplate template) return $ 
      TemplateConverter.convert template
    case preparable of
      True -> do
        let (tl, vl) = unzip params
        key <- prepare convertedTemplate tl
        liftIO $ PQ.execPrepared c key vl PQ.Binary
      False -> do
        let params' = map (\(t, v) -> (\(vb, vf) -> (t, vb, vf)) <$> v) params
        liftIO $ PQ.execParams c convertedTemplate params' PQ.Binary

liftResultProcessing :: ResultProcessing.M a -> M a
liftResultProcessing m =
  M $ ReaderT $ \(c, _, _) -> 
    EitherT $ fmap (either (Left . mapError) Right) $ ResultProcessing.run c m
  where
    mapError =
      \case
        ResultProcessing.UnexpectedResult t   -> UnexpectedResult t
        ResultProcessing.ErroneousResult t    -> ErroneousResult t
        ResultProcessing.TransactionConflict  -> TransactionConflict

{-# INLINE unitResult #-}
unitResult :: Maybe PQ.Result -> M ()
unitResult =
  liftResultProcessing . (ResultProcessing.unit <=< ResultProcessing.just)

{-# INLINE vectorResult #-}
vectorResult :: Maybe PQ.Result -> M (Vector (Vector (Maybe ByteString)))
vectorResult =
  liftResultProcessing . (ResultProcessing.vector <=< ResultProcessing.just)

{-# INLINE countResult #-}
countResult :: Maybe PQ.Result -> M Word64
countResult =
  liftResultProcessing . (ResultProcessing.count <=< ResultProcessing.just)