{-# OPTIONS_HADDOCK prune #-}
module Database.Postgres.Temp.Internal.Core where
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (race_)
import Control.Exception
import Control.Monad (forever, (>=>))
import qualified Data.ByteString.Char8 as BSC
import Data.Foldable (for_)
import Data.Monoid
import Data.String
import Data.Typeable
import qualified Database.PostgreSQL.Simple as PG
import qualified Database.PostgreSQL.Simple.Options as Client
import System.Exit (ExitCode(..))
import System.IO
import System.Posix.Signals (sigINT, signalProcess)
import System.Process
import System.Process.Internals
import Text.PrettyPrint.ANSI.Leijen hiding ((<$>))
data Event
= StartPlan String
| StartPostgres
| WaitForDB
| TryToConnect
deriving (Show, Eq, Ord)
data StartError
= StartPostgresFailed ExitCode
| InitDbFailed ExitCode
| CreateDbFailed ExitCode
| CompletePlanFailed String [String]
deriving (Show, Eq, Ord, Typeable)
instance Exception StartError
type Logger = Event -> IO ()
waitForDB :: Logger -> Client.Options -> IO ()
waitForDB logger options = do
logger TryToConnect
let theConnectionString = Client.toConnectionString options
startAction = PG.connectPostgreSQL theConnectionString
try (bracket startAction PG.close mempty) >>= \case
Left (_ :: IOError) -> threadDelay 10000 >> waitForDB logger options
Right () -> return ()
data CompleteProcessConfig = CompleteProcessConfig
{ completeProcessConfigEnvVars :: [(String, String)]
, completeProcessConfigCmdLine :: [String]
, completeProcessConfigStdIn :: Handle
, completeProcessConfigStdOut :: Handle
, completeProcessConfigStdErr :: Handle
}
prettyKeyPair ::(Pretty a, Pretty b) => a -> b -> Doc
prettyKeyPair k v = pretty k <> text ": " <> pretty v
instance Pretty CompleteProcessConfig where
pretty CompleteProcessConfig {..}
= text "completeProcessConfigEnvVars:"
<> softline
<> indent 2 (vsep (map (uncurry prettyKeyPair) completeProcessConfigEnvVars))
<> hardline
<> text "completeProcessConfigCmdLine:"
<> softline
<> text (unwords completeProcessConfigCmdLine)
startProcess
:: String
-> CompleteProcessConfig
-> IO ProcessHandle
startProcess name CompleteProcessConfig {..} = (\(_, _, _, x) -> x) <$>
createProcess_ name (proc name completeProcessConfigCmdLine)
{ std_err = UseHandle completeProcessConfigStdErr
, std_out = UseHandle completeProcessConfigStdOut
, std_in = UseHandle completeProcessConfigStdIn
, env = Just completeProcessConfigEnvVars
}
stopProcess :: ProcessHandle -> IO ExitCode
stopProcess = waitForProcess
executeProcess
:: String
-> CompleteProcessConfig
-> IO ExitCode
executeProcess name = startProcess name >=> waitForProcess
data CompletePostgresPlan = CompletePostgresPlan
{ completePostgresPlanProcessConfig :: CompleteProcessConfig
, completePostgresPlanClientOptions :: Client.Options
}
instance Pretty CompletePostgresPlan where
pretty CompletePostgresPlan {..}
= text "completePostgresPlanProcessConfig:"
<> softline
<> indent 2 (pretty completePostgresPlanProcessConfig)
<> hardline
<> text "completePostgresPlanClientOptions:"
<+> prettyOptions completePostgresPlanClientOptions
prettyOptions :: Client.Options -> Doc
prettyOptions = text . BSC.unpack . Client.toConnectionString
data PostgresProcess = PostgresProcess
{ postgresProcessClientOptions :: Client.Options
, postgresProcessHandle :: ProcessHandle
}
instance Pretty PostgresProcess where
pretty PostgresProcess {..}
= text "postgresProcessClientOptions:"
<+> prettyOptions postgresProcessClientOptions
terminateConnections :: Client.Options-> IO ()
terminateConnections options = do
let theConnectionString = Client.toConnectionString options
terminationQuery = fromString $ unlines
[ "SELECT pg_terminate_backend(pid)"
, "FROM pg_stat_activity"
, "WHERE datname=?;"
]
e <- try $ bracket (PG.connectPostgreSQL theConnectionString) PG.close $
\conn -> PG.execute conn terminationQuery
[getLast $ Client.dbname options]
case e of
Left (_ :: IOError) -> pure ()
Right _ -> pure ()
stopPostgresProcess :: PostgresProcess -> IO ExitCode
stopPostgresProcess PostgresProcess{..} = do
withProcessHandle postgresProcessHandle $ \case
OpenHandle p -> do
terminateConnections postgresProcessClientOptions
signalProcess sigINT p
OpenExtHandle {} -> pure ()
ClosedHandle _ -> return ()
waitForProcess postgresProcessHandle
startPostgresProcess :: Logger -> CompletePostgresPlan -> IO PostgresProcess
startPostgresProcess logger CompletePostgresPlan {..} = do
logger StartPostgres
let startAction = PostgresProcess completePostgresPlanClientOptions
<$> startProcess "postgres" completePostgresPlanProcessConfig
bracketOnError startAction stopPostgresProcess $
\result@PostgresProcess {..} -> do
let checkForCrash = do
mExitCode <- getProcessExitCode postgresProcessHandle
for_ mExitCode (throwIO . StartPostgresFailed)
logger WaitForDB
let options = completePostgresPlanClientOptions
{ Client.dbname = pure "template1"
}
waitForDB logger options
`race_` forever (checkForCrash >> threadDelay 100000)
return result
data CompletePlan = CompletePlan
{ completePlanLogger :: Logger
, completePlanInitDb :: Maybe CompleteProcessConfig
, completePlanCreateDb :: Maybe CompleteProcessConfig
, completePlanPostgres :: CompletePostgresPlan
, completePlanConfig :: String
, completePlanDataDirectory :: FilePath
}
instance Pretty CompletePlan where
pretty CompletePlan {..}
= text "completePlanInitDb:"
<> softline
<> indent 2 (pretty completePlanInitDb)
<> hardline
<> text "completePlanCreateDb:"
<> softline
<> indent 2 (pretty completePlanCreateDb)
<> hardline
<> text "completePlanPostgres:"
<> softline
<> indent 2 (pretty completePlanPostgres)
<> hardline
<> text "completePlanConfig:"
<> softline
<> indent 2 (pretty completePlanConfig)
<> hardline
<> text "completePlanDataDirectory:"
<+> pretty completePlanDataDirectory
throwIfNotSuccess :: Exception e => (ExitCode -> e) -> ExitCode -> IO ()
throwIfNotSuccess f = \case
ExitSuccess -> pure ()
e -> throwIO $ f e
startPlan :: CompletePlan -> IO PostgresProcess
startPlan plan@CompletePlan {..} = do
completePlanLogger $ StartPlan $ show $ pretty plan
for_ completePlanInitDb $ executeProcess "initdb" >=>
throwIfNotSuccess InitDbFailed
writeFile (completePlanDataDirectory <> "/postgresql.conf") completePlanConfig
let startAction = startPostgresProcess completePlanLogger completePlanPostgres
bracketOnError startAction stopPostgresProcess $ \result -> do
for_ completePlanCreateDb $ executeProcess "createdb" >=>
throwIfNotSuccess CreateDbFailed
pure result
stopPlan :: PostgresProcess -> IO ExitCode
stopPlan = stopPostgresProcess