bein-0.3: Bein/Commands.hs

{-# LANGUAGE FlexibleContexts #-}
module Bein.Commands where
       
import Prelude hiding (catch)
import System.IO
import Data.Convertible
import Control.Monad.Trans ()
import Control.Exception (catch)
import Control.Concurrent ( forkIO, ThreadId )
import Control.Monad.Reader
    ( MonadReader(ask), MonadIO(..), ReaderT(runReaderT) )
import Bein.Types ( BeinState(..), BeinM )
import Database.HDBC.PostgreSQL (Connection)
import Database.HDBC
    ( SqlValue, fromSql, IConnection(commit, run), quickQuery', handleSqlError, withTransaction )


encodeResponse :: String -> String -> String  
encodeResponse header body = header ++ "\n" ++ body'' ++ ".\n"
  where body' = unlines $ map dotEscape $ lines body
        body'' = if body' == "" || last body' == '\n' then body' else body' ++ "\n"
        dotEscape ('.':cs) = ".." ++ cs
        dotEscape cs = cs

forkR :: BeinState s => BeinM s () -> BeinM s ThreadId
forkR f = do st <- ask
             liftIO $ forkIO $ runReaderT f st

catchR :: BeinM s a -> (IOError -> BeinM s a) -> BeinM s a
a `catchR` b = ask >>= \st -> liftIO $ handleSqlError (runReaderT a st) `catch` (\e -> runReaderT (b e) st)

database :: BeinState s => ReaderT s IO Connection
database = ask >>= return.db

query :: BeinState s => String -> [SqlValue] -> BeinM s [[SqlValue]]
query q vs = database >>= \conn -> liftIO $ withTransaction conn (\c -> quickQuery' c q vs)

update :: BeinState s => String -> [SqlValue] -> BeinM s ()
update q vs = database >>= \conn -> liftIO $ withTransaction conn (\c -> run c q vs >> commit c)

updateWithResponse :: BeinState s => String -> [SqlValue] -> BeinM s [[SqlValue]]
updateWithResponse q vs = database >>= \conn -> do res <- liftIO $ quickQuery' conn q vs
                                                   liftIO $ commit conn
                                                   return res

oneValueUpdate :: Convertible SqlValue a => BeinState s => String -> [SqlValue] -> BeinM s (Maybe a)
oneValueUpdate q vs = do
  res <- updateWithResponse q vs
  case res of
    [] -> return Nothing
    [[x]] -> return (Just (fromSql x))
    r -> error $ "Received invalidly shaped response from database: " ++ show r
    
getCommandBlock :: Handle -> IO String
getCommandBlock h = f h ""
  where f h' acc = do l <- hGetLine h
                      case l of
                        "." -> return acc
                        ('.':'.':rest) -> f h' (acc ++ "." ++ rest ++ "\n")
                        _ -> f h' (acc ++ l ++ "\n")

maybeRowQuery :: BeinState s => String -> [SqlValue] -> BeinM s (Maybe [SqlValue])
maybeRowQuery qstr vals = query qstr vals >>= \r -> case r of
  [] -> return Nothing
  [x] -> return (Just x)
  _ -> error "Received multiple rows to a maybeRowQuery."