{-# 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
    { Settings -> DBInfo -> IO ()
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
defaultSettings = Settings
        { setupDBInfo :: DBInfo -> IO ()
setupDBInfo = \DBInfo{Text
DBServerInfo
dbiServer :: DBInfo -> DBServerInfo
dbiPass :: DBInfo -> Text
dbiUser :: DBInfo -> Text
dbiName :: DBInfo -> Text
dbiServer :: DBServerInfo
dbiPass :: Text
dbiUser :: Text
dbiName :: Text
..} -> do
            let sql :: Text
sql = Builder -> Text
toLazyText forall a b. (a -> b) -> a -> b
$
                    Builder
"CREATE USER "         forall a. Semigroup a => a -> a -> a
<> Text -> Builder
fromText Text
dbiUser forall a. Semigroup a => a -> a -> a
<>
                    Builder
" PASSWORD '"          forall a. Semigroup a => a -> a -> a
<> Text -> Builder
fromText Text
dbiPass forall a. Semigroup a => a -> a -> a
<>
                    Builder
"';\nCREATE DATABASE " forall a. Semigroup a => a -> a -> a
<> Text -> Builder
fromText Text
dbiName forall a. Semigroup a => a -> a -> a
<>
                    Builder
" OWNER "              forall a. Semigroup a => a -> a -> a
<> Text -> Builder
fromText Text
dbiUser forall a. Semigroup a => a -> a -> a
<>
                    Builder
";"
                (FilePath
cmd, [FilePath]
args) 
                    | (  DBServerInfo -> Text
dbServer DBServerInfo
dbiServer forall a. Eq a => a -> a -> Bool
== Text
"localhost" 
                      Bool -> Bool -> Bool
|| DBServerInfo -> Text
dbServer DBServerInfo
dbiServer forall a. Eq a => a -> a -> Bool
== Text
"127.0.0.1") = 
                        (FilePath
"sudo", [FilePath
"-u", FilePath
"postgres", FilePath
"psql"])
                    | Bool
otherwise = 
                        (FilePath
"psql",
                        [ FilePath
"-h", (Text -> FilePath
T.unpack forall a b. (a -> b) -> a -> b
$ DBServerInfo -> Text
dbServer DBServerInfo
dbiServer)
                        , FilePath
"-p", (forall a. Show a => a -> FilePath
show forall a b. (a -> b) -> a -> b
$ DBServerInfo -> Int
dbPort DBServerInfo
dbiServer)
                        , FilePath
"-U", FilePath
"postgres"])
            FilePath
_ <- FilePath -> [FilePath] -> FilePath -> IO FilePath
readProcess FilePath
cmd [FilePath]
args forall a b. (a -> b) -> a -> b
$ Text -> FilePath
TL.unpack Text
sql
            forall (m :: * -> *) a. Monad m => a -> m a
return ()
        }

-- | Information on an individual PostgreSQL database.
data DBInfo = DBInfo
    { DBInfo -> Text
dbiName   :: Text
    , DBInfo -> Text
dbiUser   :: Text
    , DBInfo -> Text
dbiPass   :: Text
    , DBInfo -> DBServerInfo
dbiServer :: DBServerInfo
    }
    deriving Int -> DBInfo -> ShowS
[DBInfo] -> ShowS
DBInfo -> FilePath
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [DBInfo] -> ShowS
$cshowList :: [DBInfo] -> ShowS
show :: DBInfo -> FilePath
$cshow :: DBInfo -> FilePath
showsPrec :: Int -> DBInfo -> ShowS
$cshowsPrec :: Int -> DBInfo -> ShowS
Show

data DBServerInfo = DBServerInfo
    { DBServerInfo -> Text
dbServer :: Text
    , DBServerInfo -> Int
dbPort   :: Int
    }
    deriving Int -> DBServerInfo -> ShowS
[DBServerInfo] -> ShowS
DBServerInfo -> FilePath
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [DBServerInfo] -> ShowS
$cshowList :: [DBServerInfo] -> ShowS
show :: DBServerInfo -> FilePath
$cshow :: DBServerInfo -> FilePath
showsPrec :: Int -> DBServerInfo -> ShowS
$cshowsPrec :: Int -> DBServerInfo -> ShowS
Show

randomDBI :: DBServerInfo -> R.StdGen -> (DBInfo, R.StdGen)
randomDBI :: DBServerInfo -> StdGen -> (DBInfo, StdGen)
randomDBI DBServerInfo
dbsi =
    forall s a. State s a -> s -> (a, s)
S.runState (Text -> Text -> Text -> DBServerInfo -> DBInfo
DBInfo forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT StdGen Identity Text
rt forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> StateT StdGen Identity Text
rt forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> StateT StdGen Identity Text
rt forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall (f :: * -> *) a. Applicative f => a -> f a
pure DBServerInfo
dbsi)) 
  where
    rt :: StateT StdGen Identity Text
rt = FilePath -> Text
T.pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
10 (forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
S.state forall a b. (a -> b) -> a -> b
$ forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
R.randomR (Char
'a', Char
'z'))

instance ToJSON DBInfo where
    toJSON :: DBInfo -> Value
toJSON (DBInfo Text
n Text
u Text
p (DBServerInfo Text
server Int
port)) = [Pair] -> Value
object
        [ Key
"name"   forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text
n
        , Key
"user"   forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text
u
        , Key
"pass"   forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text
p
        , Key
"server" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text
server
        , Key
"port"   forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Int
port
        ]

instance FromJSON DBInfo where
    parseJSON :: Value -> Parser DBInfo
parseJSON (Object Object
o) = Text -> Text -> Text -> DBServerInfo -> DBInfo
DBInfo
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
o forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"name"
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
o forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"user"
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
o forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"pass"
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Text -> Int -> DBServerInfo
DBServerInfo
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
o forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"server" forall a. Parser (Maybe a) -> a -> Parser a
.!= Text
"localhost"
            forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
o forall a. FromJSON a => Object -> Key -> Parser (Maybe a)
.:? Key
"port"   forall a. Parser (Maybe a) -> a -> Parser a
.!= Int
5432)
    parseJSON Value
_ = forall (m :: * -> *) a. MonadPlus m => m a
mzero

instance FromJSON DBServerInfo where
    parseJSON :: Value -> Parser DBServerInfo
parseJSON (Object Object
o) = Text -> Int -> DBServerInfo
DBServerInfo
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
o forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"server"
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
o forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"port"
    parseJSON Value
_ = forall (m :: * -> *) a. MonadPlus m => m a
mzero

defaultDBServerInfo :: DBServerInfo
defaultDBServerInfo :: DBServerInfo
defaultDBServerInfo = Text -> Int -> DBServerInfo
DBServerInfo Text
"localhost" Int
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 -> FilePath -> IO Plugin
load Settings{DBInfo -> IO ()
setupDBInfo :: DBInfo -> IO ()
setupDBInfo :: Settings -> DBInfo -> IO ()
..} FilePath
fp = do
    Bool -> FilePath -> IO ()
createDirectoryIfMissing Bool
True forall a b. (a -> b) -> a -> b
$ ShowS
takeDirectory FilePath
fp
    Bool
e <- FilePath -> IO Bool
doesFileExist FilePath
fp
    Either ParseException (Map Text DBInfo)
edb <- if Bool
e
        then forall a. FromJSON a => FilePath -> IO (Either ParseException a)
decodeFileEither FilePath
fp
        else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right forall k a. Map k a
Map.empty
    case Either ParseException (Map Text DBInfo)
edb of
        Left ParseException
ex -> forall e a. Exception e => e -> IO a
throwIO ParseException
ex
        Right Map Text DBInfo
db -> Map Text DBInfo -> IO Plugin
go Map Text DBInfo
db
  where
    go :: Map Text DBInfo -> IO Plugin
go Map Text DBInfo
db0 = do
        Chan Command
chan <- forall a. IO (Chan a)
newChan
        StdGen
g0 <- forall (m :: * -> *). MonadIO m => m StdGen
R.newStdGen
        -- FIXME stop using the worker thread approach?
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
S.evalStateT (Map Text DBInfo
db0, StdGen
g0) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ Chan Command -> StateT (Map Text DBInfo, StdGen) IO ()
loop Chan Command
chan
        forall (m :: * -> *) a. Monad m => a -> m a
return Plugin
            { pluginGetEnv :: Text -> Object -> IO [(Text, Text)]
pluginGetEnv = \Text
appname Object
o ->
                case forall v. Key -> KeyMap v -> Maybe v
AK.lookup Key
"postgres" Object
o of
                    Just (Array Array
v) -> do
                        let dbServer :: DBServerInfo
dbServer = forall a. a -> Maybe a -> a
fromMaybe DBServerInfo
defaultDBServerInfo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Parser b) -> a -> Maybe b
parseMaybe forall a. FromJSON a => Value -> Parser a
parseJSON forall a b. (a -> b) -> a -> b
$ forall a. Vector a -> a
V.head Array
v
                        Chan Command -> Text -> DBServerInfo -> IO [(Text, Text)]
doenv Chan Command
chan Text
appname DBServerInfo
dbServer
                    Just (Bool Bool
True) -> do
                        Chan Command -> Text -> DBServerInfo -> IO [(Text, Text)]
doenv Chan Command
chan Text
appname DBServerInfo
defaultDBServerInfo
                    Maybe Value
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return []
            }
      where doenv :: Chan Command -> Text -> DBServerInfo -> IO [(Text, Text)]
doenv Chan Command
chan Text
appname DBServerInfo
dbs = do
            MVar (Either SomeException DBInfo)
x <- forall a. IO (MVar a)
newEmptyMVar
            forall a. Chan a -> a -> IO ()
writeChan Chan Command
chan forall a b. (a -> b) -> a -> b
$ Text
-> DBServerInfo
-> (Either SomeException DBInfo -> IO ())
-> Command
GetConfig Text
appname DBServerInfo
dbs forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> a -> IO ()
putMVar MVar (Either SomeException DBInfo)
x
            Either SomeException DBInfo
edbi <- forall a. MVar a -> IO a
takeMVar MVar (Either SomeException DBInfo)
x
            Either SomeException DBInfo -> IO [(Text, Text)]
edbiToEnv Either SomeException DBInfo
edbi
                    
    tmpfp :: FilePath
tmpfp = FilePath
fp FilePath -> ShowS
<.> FilePath
"tmp"

    loop :: Chan Command -> StateT (Map Text DBInfo, StdGen) IO ()
loop Chan Command
chan = do
        GetConfig Text
appname DBServerInfo
dbServer Either SomeException DBInfo -> IO ()
f <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a. Chan a -> IO a
readChan Chan Command
chan
        (Map Text DBInfo
db, StdGen
g) <- forall (m :: * -> *) s. Monad m => StateT s m s
S.get
        Either SomeException DBInfo
dbi <-
            case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
appname Map Text DBInfo
db of
                Just DBInfo
dbi -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right DBInfo
dbi
                Maybe DBInfo
Nothing -> do
                    let (DBInfo
dbi', StdGen
g') = DBServerInfo -> StdGen -> (DBInfo, StdGen)
randomDBI DBServerInfo
dbServer StdGen
g
                    let dbi :: DBInfo
dbi = DBInfo
dbi'
                            { dbiName :: Text
dbiName = Text -> Text
sanitize Text
appname forall a. Semigroup a => a -> a -> a
<> DBInfo -> Text
dbiName DBInfo
dbi'
                            , dbiUser :: Text
dbiUser = Text -> Text
sanitize Text
appname forall a. Semigroup a => a -> a -> a
<> DBInfo -> Text
dbiUser DBInfo
dbi'
                            }
                    Either SomeException ()
ex <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ DBInfo -> IO ()
setupDBInfo DBInfo
dbi
                    case Either SomeException ()
ex of
                        Left SomeException
e -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left SomeException
e
                        Right () -> do
                            let db' :: Map Text DBInfo
db' = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
appname DBInfo
dbi Map Text DBInfo
db
                            Either SomeException ()
ey <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ do
                                forall a. ToJSON a => FilePath -> a -> IO ()
encodeFile FilePath
tmpfp Map Text DBInfo
db'
                                FilePath -> FilePath -> IO ()
renameFile FilePath
tmpfp FilePath
fp
                            case Either SomeException ()
ey of
                                Left SomeException
e -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left SomeException
e
                                Right () -> do
                                    forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (Map Text DBInfo
db', StdGen
g')
                                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right DBInfo
dbi
        forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ Either SomeException DBInfo -> IO ()
f Either SomeException DBInfo
dbi

    sanitize :: Text -> Text
sanitize = (Char -> Char) -> Text -> Text
T.map Char -> Char
sanitize'
    sanitize' :: Char -> Char
sanitize' Char
c
        | Char
'A' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'Z' = Char -> Char
C.toLower Char
c
        | Char
'a' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'z' = Char
c
        | Char
'0' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'9' = Char
c
        | Bool
otherwise = Char
'_'

edbiToEnv :: Either SomeException DBInfo
          -> IO [(Text, Text)]
edbiToEnv :: Either SomeException DBInfo -> IO [(Text, Text)]
edbiToEnv (Left SomeException
e) = case forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
                       Just IOError
e' -> if IOError -> Bool
isDoesNotExistError IOError
e'
                         Bool -> Bool -> Bool
&& IOError -> Maybe FilePath
ioeGetFileName IOError
e' forall a. Eq a => a -> a -> Bool
== forall a. a -> Maybe a
Just FilePath
"sudo"
                         then forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$
                         IOError -> FilePath -> Maybe Handle -> Maybe FilePath -> IOError
annotateIOError IOError
e' FilePath
"\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" forall a. Maybe a
Nothing forall a. Maybe a
Nothing
                         else forall e a. Exception e => e -> IO a
throwIO SomeException
e
                       Maybe IOError
Nothing -> forall e a. Exception e => e -> IO a
throwIO SomeException
e
edbiToEnv (Right DBInfo
dbi) = forall (m :: * -> *) a. Monad m => a -> m a
return
    [ (Text
"PGHOST", DBServerInfo -> Text
dbServer forall a b. (a -> b) -> a -> b
$ DBInfo -> DBServerInfo
dbiServer DBInfo
dbi)
    , (Text
"PGPORT", FilePath -> Text
T.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> FilePath
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. DBServerInfo -> Int
dbPort forall a b. (a -> b) -> a -> b
$ DBInfo -> DBServerInfo
dbiServer DBInfo
dbi)
    , (Text
"PGUSER", DBInfo -> Text
dbiUser DBInfo
dbi)
    , (Text
"PGPASS", DBInfo -> Text
dbiPass DBInfo
dbi)
    , (Text
"PGDATABASE", DBInfo -> Text
dbiName DBInfo
dbi)
    ]