keter/Keter/Plugin/Postgres.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
module Keter.Plugin.Postgres
    ( -- * Settings
      Settings
    , setupDBInfo
      -- * Functions
    , load
    ) where

import           Control.Applicative       ((<$>), (<*>), pure)
import           Control.Concurrent        (forkIO)
import           Control.Concurrent.Chan
import           Control.Concurrent.MVar
import           Control.Exception         (throwIO, try)
import           Control.Monad             (forever, mzero, replicateM, void)
import           Control.Monad.Trans.Class (lift)
import qualified Control.Monad.Trans.State as S
import qualified Data.Char                 as C
import           Data.Default
import qualified Data.HashMap.Strict       as HMap
import qualified Data.Map                  as Map
import           Data.Maybe                (fromMaybe)
import           Data.Monoid               ((<>))
import qualified Data.Text                 as T
import qualified Data.Text.Lazy            as TL
import           Data.Text.Lazy.Builder    (fromText, toLazyText)
import qualified Data.Vector               as V
import           Data.Yaml
import           Keter.Types
import           Prelude                   hiding (FilePath)
import           System.Directory          (createDirectoryIfMissing,
                                            doesFileExist, renameFile)
import           System.FilePath           (takeDirectory, (<.>))
import           System.Process            (readProcess)
import qualified System.Random             as R

data Settings = Settings
    { setupDBInfo :: DBInfo -> IO ()
      -- ^ How to create the given user/database. Default: uses the @psql@
      -- command line tool and @sudo -u postgres@.
    }

instance Default Settings where
    def = Settings
        { setupDBInfo = \DBInfo{..} -> do
            let sql = toLazyText $
                    "CREATE USER "         <> fromText dbiUser <>
                    " PASSWORD '"          <> fromText dbiPass <>
                    "';\nCREATE DATABASE " <> fromText dbiName <>
                    " OWNER "              <> fromText dbiUser <>
                    ";"
                (cmd, args) 
                    | (  dbServer dbiServer == "localhost" 
                      || dbServer dbiServer == "127.0.0.1") = 
                        ("sudo", ["-u", "postgres", "psql"])
                    | otherwise = 
                        ("psql",
                        [ "-h", (T.unpack $ dbServer dbiServer)
                        , "-p", (show $ dbPort dbiServer)
                        , "-U", "postgres"])
            _ <- readProcess cmd args $ TL.unpack sql
            return ()
        }

-- | Information on an individual PostgreSQL database.
data DBInfo = DBInfo
    { dbiName   :: Text
    , dbiUser   :: Text
    , dbiPass   :: Text
    , dbiServer :: DBServerInfo
    }
    deriving Show

data DBServerInfo = DBServerInfo
    { dbServer :: Text
    , dbPort   :: Int
    }
    deriving Show

randomDBI :: DBServerInfo -> R.StdGen -> (DBInfo, R.StdGen)
randomDBI dbsi =
    S.runState (DBInfo <$> rt <*> rt <*> rt <*> (pure dbsi)) 
  where
    rt = T.pack <$> replicateM 10 (S.state $ R.randomR ('a', 'z'))

instance ToJSON DBInfo where
    toJSON (DBInfo n u p (DBServerInfo server port)) = object
        [ "name"   .= n
        , "user"   .= u
        , "pass"   .= p
        , "server" .= server
        , "port"   .= port
        ]

instance FromJSON DBInfo where
    parseJSON (Object o) = DBInfo
        <$> o .: "name"
        <*> o .: "user"
        <*> o .: "pass"
        <*> (DBServerInfo
            <$> o .:? "server" .!= "localhost"
            <*> o .:? "port"   .!= 5432)
    parseJSON _ = mzero

instance FromJSON DBServerInfo where
    parseJSON (Object o) = DBServerInfo
        <$> o .: "server"
        <*> o .: "port"
    parseJSON _ = mzero
    
instance Default DBServerInfo where
    def = DBServerInfo "localhost" 5432

data Command = GetConfig Appname DBServerInfo (Either SomeException DBInfo -> IO ())

-- | Load a set of existing connections from a config file. If the file does
-- not exist, it will be created. Any newly created databases will
-- automatically be saved to this file.
load :: Settings -> FilePath -> IO Plugin
load Settings{..} fp = do
    createDirectoryIfMissing True $ takeDirectory fp
    e <- doesFileExist fp
    edb <- if e
        then decodeFileEither fp
        else return $ Right Map.empty
    case edb of
        Left ex -> throwIO ex
        Right db -> go db
  where
    go db0 = do
        chan <- newChan
        g0 <- R.newStdGen
        -- FIXME stop using the worker thread approach?
        void $ forkIO $ flip S.evalStateT (db0, g0) $ forever $ loop chan
        return Plugin
            { pluginGetEnv = \appname o ->
                case HMap.lookup "postgres" o of
                    Just (Array v) -> do
                        let dbServer = fromMaybe def . parseMaybe parseJSON $ V.head v
                        doenv chan appname dbServer
                    Just (Bool True) -> do
                        doenv chan appname def
                    _ -> return []
            }
      where doenv chan appname dbs = do
            x <- newEmptyMVar
            writeChan chan $ GetConfig appname dbs $ putMVar x
            edbi <- takeMVar x
            edbiToEnv edbi
                    
    tmpfp = fp <.> "tmp"

    loop chan = do
        GetConfig appname dbServer f <- lift $ readChan chan
        (db, g) <- S.get
        dbi <-
            case Map.lookup appname db of
                Just dbi -> return $ Right dbi
                Nothing -> do
                    let (dbi', g') = randomDBI dbServer g
                    let dbi = dbi'
                            { dbiName = sanitize appname <> dbiName dbi'
                            , dbiUser = sanitize appname <> dbiUser dbi'
                            }
                    ex <- lift $ try $ setupDBInfo dbi
                    case ex of
                        Left e -> return $ Left e
                        Right () -> do
                            let db' = Map.insert appname dbi db
                            ey <- lift $ try $ do
                                encodeFile tmpfp db'
                                renameFile tmpfp fp
                            case ey of
                                Left e -> return $ Left e
                                Right () -> do
                                    S.put (db', g')
                                    return $ Right dbi
        lift $ f dbi

    sanitize = T.map sanitize'
    sanitize' c
        | 'A' <= c && c <= 'Z' = C.toLower c
        | 'a' <= c && c <= 'z' = c
        | '0' <= c && c <= '9' = c
        | otherwise = '_'

edbiToEnv :: Either SomeException DBInfo
          -> IO [(Text, Text)]
edbiToEnv (Left e) = throwIO e
edbiToEnv (Right dbi) = return
    [ ("PGHOST", dbServer $ dbiServer dbi)
    , ("PGPORT", T.pack . show . dbPort $ dbiServer dbi)
    , ("PGUSER", dbiUser dbi)
    , ("PGPASS", dbiPass dbi)
    , ("PGDATABASE", dbiName dbi)
    ]