{-# LANGUAGE OverloadedStrings #-}
module Gargoyle.PostgreSQL where

import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Search as BS
import Data.Foldable (for_)
import Data.Maybe
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import System.Directory
import System.Exit
import System.IO
import System.Posix.Escape
import System.Posix.Signals
import System.Process

import Gargoyle

-- | A 'Gargoyle' that assumes `initdb` and `postgres` are in the path and
-- will perform a 'fast shutdown' on termination (see below).
defaultPostgres :: Gargoyle FilePath ByteString
defaultPostgres :: Gargoyle FilePath ByteString
defaultPostgres = FilePath
-> (FilePath -> FilePath -> IO ()) -> Gargoyle FilePath ByteString
mkPostgresGargoyle FilePath
"pg_ctl" FilePath -> FilePath -> IO ()
shutdownPostgresFast

-- | Create a gargoyle by telling it where the relevant PostgreSQL executables are and
-- what it should do in order to shut down the server. This module provides two options:
-- 'shutdownPostgresSmart' and 'shutdownPostgresFast'.
mkPostgresGargoyle :: FilePath -- ^ Path to `pg_ctl`
                   -> (FilePath -> FilePath -> IO ()) -- ^ Shutdown function
                   -> Gargoyle FilePath ByteString
                   -- ^ The 'Gargoyle' returned provides to client code the connection
                   -- string that can be used to connect to the PostgreSQL server
mkPostgresGargoyle :: FilePath
-> (FilePath -> FilePath -> IO ()) -> Gargoyle FilePath ByteString
mkPostgresGargoyle FilePath
pgCtlPath FilePath -> FilePath -> IO ()
shutdownFun = Gargoyle :: forall pid a.
FilePath
-> (FilePath -> IO ())
-> (FilePath -> IO pid)
-> (pid -> IO ())
-> (FilePath -> IO a)
-> Gargoyle pid a
Gargoyle
  { _gargoyle_exec :: FilePath
_gargoyle_exec = FilePath
"gargoyle-postgres-monitor"
  , _gargoyle_init :: FilePath -> IO ()
_gargoyle_init = FilePath -> FilePath -> IO ()
initLocalPostgres FilePath
pgCtlPath
  , _gargoyle_start :: FilePath -> IO FilePath
_gargoyle_start = FilePath -> FilePath -> IO FilePath
startLocalPostgres FilePath
pgCtlPath
  , _gargoyle_stop :: FilePath -> IO ()
_gargoyle_stop = FilePath -> FilePath -> IO ()
shutdownFun FilePath
pgCtlPath
  , _gargoyle_getInfo :: FilePath -> IO ByteString
_gargoyle_getInfo = FilePath -> IO ByteString
getLocalPostgresConnectionString
  }

-- | Create a new PostgreSQL database in a local folder. This is a low level function used to
-- define the PostgreSQL 'Gargoyle'.
initLocalPostgres :: FilePath -- ^ Path to PostgreSQL `pg_ctl` executable
                  -> FilePath -- ^ Path in which to initialize PostgreSQL Server
                  -> IO ()
initLocalPostgres :: FilePath -> FilePath -> IO ()
initLocalPostgres FilePath
binPath FilePath
dbDir = do
  Handle
devNull <- FilePath -> IOMode -> IO Handle
openFile FilePath
"/dev/null" IOMode
WriteMode
  (Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
initdb) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess (FilePath -> [FilePath] -> CreateProcess
proc FilePath
binPath
    [ FilePath
"init"
    , FilePath
"-D", FilePath
dbDir
    , FilePath
"-o", [FilePath] -> FilePath
escapeMany
      [ FilePath
"-U", FilePath
"postgres"
      , FilePath
"--no-locale"
      , FilePath
"-E", FilePath
"UTF8"
      ]
    ]) { std_in :: StdStream
std_in = StdStream
NoStream, std_out :: StdStream
std_out = Handle -> StdStream
UseHandle Handle
devNull, std_err :: StdStream
std_err = StdStream
Inherit }
  ExitCode
r <- ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
initdb
  case ExitCode
r of
    ExitCode
ExitSuccess -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    ExitCode
_ -> do
      FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"initLocalPostgres failed: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ ExitCode -> FilePath
forall a. Show a => a -> FilePath
show ExitCode
r
      ExitCode -> IO ()
forall a. ExitCode -> IO a
exitWith ExitCode
r

-- | Produces the connection string for a local postgresql database. This is a low level function
-- used to define the PostgreSQL 'Gargoyle'
getLocalPostgresConnectionString :: FilePath -> IO ByteString
getLocalPostgresConnectionString :: FilePath -> IO ByteString
getLocalPostgresConnectionString FilePath
dbDir = do
  FilePath
absoluteDbDir <- FilePath -> IO FilePath
makeAbsolute FilePath
dbDir
  ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
forall a. Monoid a => [a] -> a
mconcat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
    [ ByteString
"postgresql://postgres@"
    , (ByteString -> ByteString
LBS.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString -> ByteString -> ByteString
forall rep.
Substitution rep =>
ByteString -> rep -> ByteString -> ByteString
BS.replace ByteString
"/" (ByteString
"%2F" :: LBS.ByteString) (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
T.encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ FilePath -> Text
T.pack FilePath
absoluteDbDir)
    , ByteString
"/postgres"
    ]

-- | Start a postgres server that is assumed to be in the given folder. This is a low level function
-- used to define the PostgreSQL 'Gargoyle'
startLocalPostgres :: FilePath -- ^ Path to PostgreSQL `pg_ctl` executable
                   -> FilePath -- ^ Path where the server to start is located
                   -> IO FilePath -- ^ handle of the PostgreSQL server
startLocalPostgres :: FilePath -> FilePath -> IO FilePath
startLocalPostgres FilePath
binPath FilePath
dbDir = do
  FilePath
absoluteDbDir <- FilePath -> IO FilePath
makeAbsolute FilePath
dbDir
  Handle
devNull <- FilePath -> IOMode -> IO Handle
openFile FilePath
"/dev/null" IOMode
WriteMode
  (Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
postgres) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess (FilePath -> [FilePath] -> CreateProcess
proc FilePath
binPath
    [ FilePath
"start"
    , FilePath
"-D", FilePath
absoluteDbDir
    , FilePath
"-w"
    , FilePath
"-o", [FilePath] -> FilePath
escapeMany
      [ FilePath
"-h", FilePath
""
      , FilePath
"-k", FilePath
absoluteDbDir
      ]
    ]) { std_in :: StdStream
std_in = StdStream
NoStream, std_out :: StdStream
std_out = Handle -> StdStream
UseHandle Handle
devNull, std_err :: StdStream
std_err = StdStream
Inherit }
  ExitCode
r <- ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
postgres
  case ExitCode
r of
    ExitCode
ExitSuccess -> FilePath -> IO FilePath
forall (m :: * -> *) a. Monad m => a -> m a
return FilePath
absoluteDbDir
    ExitCode
_ -> do
      FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"startLocalPostgres failed: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> ExitCode -> FilePath
forall a. Show a => a -> FilePath
show ExitCode
r
      ExitCode -> IO FilePath
forall a. ExitCode -> IO a
exitWith ExitCode
r

-- | Perform a "Smart Shutdown" of Postgres;
-- see http://www.postgresql.org/docs/current/static/server-shutdown.html
shutdownPostgresSmart :: FilePath -- ^ Path to PostgreSQL `pg_ctl` executable
                      -> FilePath -- ^ Path where the server to start is located
                      -> IO ()
shutdownPostgresSmart :: FilePath -> FilePath -> IO ()
shutdownPostgresSmart = FilePath -> FilePath -> FilePath -> IO ()
shutdownPostgresWithMode FilePath
"smart"

-- | Perform a "Fast Shutdown" of Postgres;
-- see http://www.postgresql.org/docs/current/static/server-shutdown.html
shutdownPostgresFast :: FilePath -- ^ Path to PostgreSQL `pg_ctl` executable
                      -> FilePath -- ^ Path where the server to start is located
                      -> IO ()
shutdownPostgresFast :: FilePath -> FilePath -> IO ()
shutdownPostgresFast = FilePath -> FilePath -> FilePath -> IO ()
shutdownPostgresWithMode FilePath
"fast"

-- | Perform a "Immediate Shutdown" of Postgres;
-- see http://www.postgresql.org/docs/current/static/server-shutdown.html
shutdownPostgresImmediate :: FilePath -- ^ Path to PostgreSQL `pg_ctl` executable
                      -> FilePath -- ^ Path where the server to start is located
                      -> IO ()
shutdownPostgresImmediate :: FilePath -> FilePath -> IO ()
shutdownPostgresImmediate = FilePath -> FilePath -> FilePath -> IO ()
shutdownPostgresWithMode FilePath
"immediate"

shutdownPostgresWithMode :: String -- ^ The shutdown mode to execute; see https://www.postgresql.org/docs/9.5/app-pg-ctl.html
                         -> FilePath -- ^ Path to PostgreSQL `pg_ctl` executable
                         -> FilePath -- ^ Path where the server to start is located
                         -> IO ()
shutdownPostgresWithMode :: FilePath -> FilePath -> FilePath -> IO ()
shutdownPostgresWithMode FilePath
mode FilePath
binPath FilePath
absoluteDbDir = do
  (Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
postgres) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess (FilePath -> [FilePath] -> CreateProcess
proc FilePath
binPath
    [ FilePath
"stop"
    , FilePath
"-D", FilePath
absoluteDbDir
    , FilePath
"-w"
    , FilePath
"-m", FilePath
mode
    ]) { std_in :: StdStream
std_in = StdStream
NoStream, std_out :: StdStream
std_out = StdStream
NoStream, std_err :: StdStream
std_err = StdStream
Inherit }
  ExitCode
r <- ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
postgres
  case ExitCode
r of
    ExitCode
ExitSuccess -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    ExitCode
_ -> do
      FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"stopLocalPostgres failed: " FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> ExitCode -> FilePath
forall a. Show a => a -> FilePath
show ExitCode
r
      ExitCode -> IO ()
forall a. ExitCode -> IO a
exitWith ExitCode
r

-- | Run `psql` against a Gargoyle managed db.
psqlLocal :: Gargoyle pid ByteString -- ^ 'Gargoyle' against which to run
          -> FilePath -- ^ The path to `psql`
          -> FilePath -- ^ The path where the managed daemon is expected
          -> Maybe String
          -- ^ Optionally provide stdin input instead of an inheriting current stdin.
          -- It will have a newline and quit command appended to it.
          -> IO ()
psqlLocal :: Gargoyle pid ByteString
-> FilePath -> FilePath -> Maybe FilePath -> IO ()
psqlLocal Gargoyle pid ByteString
g FilePath
psqlPath FilePath
dbPath Maybe FilePath
minput = Gargoyle pid ByteString
-> FilePath -> (ByteString -> IO ()) -> IO ()
forall pid a b. Gargoyle pid a -> FilePath -> (a -> IO b) -> IO b
withGargoyle Gargoyle pid ByteString
g FilePath
dbPath ((ByteString -> IO ()) -> IO ()) -> (ByteString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ByteString
dbUri -> do
  IO Handler -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Handler -> IO ()) -> IO Handler -> IO ()
forall a b. (a -> b) -> a -> b
$ Signal -> Handler -> Maybe SignalSet -> IO Handler
installHandler Signal
keyboardSignal Handler
Ignore Maybe SignalSet
forall a. Maybe a
Nothing
  let psqlProc :: CreateProcess
psqlProc = (FilePath -> [FilePath] -> CreateProcess
proc FilePath
psqlPath [ Text -> FilePath
T.unpack (Text -> FilePath) -> Text -> FilePath
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
T.decodeUtf8 ByteString
dbUri ])
        { std_in :: StdStream
std_in = case Maybe FilePath
minput of
            Maybe FilePath
Nothing -> StdStream
Inherit
            Just FilePath
_ -> StdStream
CreatePipe
        , std_out :: StdStream
std_out = StdStream
Inherit
        , std_err :: StdStream
std_err = StdStream
Inherit
        }
  (Maybe Handle
mStdin, Maybe Handle
_, Maybe Handle
_, ProcessHandle
psql) <- CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess CreateProcess
psqlProc
  case Maybe FilePath
minput of
    Maybe FilePath
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just FilePath
input -> Handle -> FilePath -> IO ()
hPutStrLn (Maybe Handle -> Handle
forall a. HasCallStack => Maybe a -> a
fromJust Maybe Handle
mStdin) (FilePath
input FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
"\n\\q")
  ExitCode
ExitSuccess <- ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
psql
  () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Run an arbitrary process against a Gargoyle-managed DB, providing connection
--   information by substituting a given argument pattern with the connection string.
runPgLocalWithSubstitution
  :: Gargoyle pid ByteString -- ^ 'Gargoyle' against which to run
  -> FilePath -- ^ The path where the managed daemon is expected
  -> FilePath -- ^ Path to process to run
  -> (String -> [String]) -- ^ Function producing arguments to the process given the connection string
  -> Maybe String -- ^ Optionally provide stdin input instead of an inheriting current stdin.
  -> IO ExitCode
runPgLocalWithSubstitution :: Gargoyle pid ByteString
-> FilePath
-> FilePath
-> (FilePath -> [FilePath])
-> Maybe FilePath
-> IO ExitCode
runPgLocalWithSubstitution Gargoyle pid ByteString
g FilePath
dbPath FilePath
procPath FilePath -> [FilePath]
mkProcArgs Maybe FilePath
mInput = Gargoyle pid ByteString
-> FilePath -> (ByteString -> IO ExitCode) -> IO ExitCode
forall pid a b. Gargoyle pid a -> FilePath -> (a -> IO b) -> IO b
withGargoyle Gargoyle pid ByteString
g FilePath
dbPath ((ByteString -> IO ExitCode) -> IO ExitCode)
-> (ByteString -> IO ExitCode) -> IO ExitCode
forall a b. (a -> b) -> a -> b
$ \ByteString
dbUri -> do
  IO Handler -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Handler -> IO ()) -> IO Handler -> IO ()
forall a b. (a -> b) -> a -> b
$ Signal -> Handler -> Maybe SignalSet -> IO Handler
installHandler Signal
keyboardSignal Handler
Ignore Maybe SignalSet
forall a. Maybe a
Nothing
  let
    procSpec :: CreateProcess
procSpec = (FilePath -> [FilePath] -> CreateProcess
proc FilePath
procPath ([FilePath] -> CreateProcess) -> [FilePath] -> CreateProcess
forall a b. (a -> b) -> a -> b
$ FilePath -> [FilePath]
mkProcArgs (FilePath -> [FilePath]) -> FilePath -> [FilePath]
forall a b. (a -> b) -> a -> b
$ Text -> FilePath
T.unpack (Text -> FilePath) -> Text -> FilePath
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
T.decodeUtf8 ByteString
dbUri)
      { std_in :: StdStream
std_in = case Maybe FilePath
mInput of
          Maybe FilePath
Nothing -> StdStream
Inherit
          Just FilePath
_ -> StdStream
CreatePipe
      , std_out :: StdStream
std_out = StdStream
Inherit
      , std_err :: StdStream
std_err = StdStream
Inherit
      }
  CreateProcess
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO ExitCode)
-> IO ExitCode
forall a.
CreateProcess
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO a)
-> IO a
withCreateProcess CreateProcess
procSpec ((Maybe Handle
  -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO ExitCode)
 -> IO ExitCode)
-> (Maybe Handle
    -> Maybe Handle -> Maybe Handle -> ProcessHandle -> IO ExitCode)
-> IO ExitCode
forall a b. (a -> b) -> a -> b
$ \Maybe Handle
mStdin Maybe Handle
_ Maybe Handle
_ ProcessHandle
procHandle -> do
    Maybe FilePath -> (FilePath -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ Maybe FilePath
mInput ((FilePath -> IO ()) -> IO ()) -> (FilePath -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
      Handle -> FilePath -> IO ()
hPutStrLn (Handle -> Maybe Handle -> Handle
forall a. a -> Maybe a -> a
fromMaybe (FilePath -> Handle
forall a. HasCallStack => FilePath -> a
error FilePath
"runPgLocalWithSubstitution: input stream was expected") Maybe Handle
mStdin)
    ProcessHandle -> IO ExitCode
waitForProcess ProcessHandle
procHandle