Postgres.hs

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

import           Keter.Common
import           Control.Applicative       ((<$>), (<*>), pure)
import           Keter.Aeson.KeyHelper      as AK (lookup)
import           Control.Concurrent        (forkIO)
import           Control.Concurrent.Chan
import           Control.Concurrent.MVar
import           Control.Exception         (fromException, 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 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           Prelude                   hiding (FilePath)
import           System.Directory          (createDirectoryIfMissing,
                                            doesFileExist, renameFile)
import           System.FilePath           (takeDirectory, (<.>))
import           System.IO.Error           (annotateIOError,
                                            ioeGetFileName,
                                            isDoesNotExistError)
import           System.Process            (readProcess)
import qualified System.Random             as R
import           Data.Text                  (Text)
import           System.FilePath            (FilePath)
import           Control.Exception          (SomeException)

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

defaultDBServerInfo :: DBServerInfo
defaultDBServerInfo = 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 AK.lookup "postgres" o of
                    Just (Array v) -> do
                        let dbServer = fromMaybe defaultDBServerInfo . parseMaybe parseJSON $ V.head v
                        doenv chan appname dbServer
                    Just (Bool True) -> do
                        doenv chan appname defaultDBServerInfo
                    _ -> 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) = case fromException e of
                       Just e' -> if isDoesNotExistError e'
                         && ioeGetFileName e' == Just "sudo"
                         then throwIO $
                         annotateIOError e' "\nWe are unable to find sudo in your local path, this could be because you don't have sudo installed. Sudo is necessary for keter to connect to postgres running on the local server.\nsudo" Nothing Nothing
                         else throwIO e
                       Nothing -> 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)
    ]