module Database.Postgres.Temp.Internal.Core where
import qualified Database.PostgreSQL.Simple.Options as PostgresClient
import qualified Database.PostgreSQL.Simple as PG
import System.Process.Internals
import System.Exit (ExitCode(..))
import Data.String
import System.Posix.Signals (sigINT, signalProcess)
import Control.Exception
import System.Process (getProcessExitCode, waitForProcess)
import Data.Foldable (for_)
import Control.Concurrent.Async (race_)
import Control.Monad (forever, (>=>))
import Control.Concurrent (threadDelay)
import Data.Typeable
import System.IO
import System.Process
import Data.Monoid
data Event
= StartPostgres
| WaitForDB
deriving (Show, Eq, Ord)
data StartError
= StartPostgresFailed ExitCode
| InitDbFailed ExitCode
| CreateDbFailed ExitCode
| CompletePlanFailed [String]
deriving (Show, Eq, Ord, Typeable)
instance Exception StartError
type Logger = Event -> IO ()
waitForDB :: PostgresClient.Options -> IO ()
waitForDB options = do
let theConnectionString = PostgresClient.toConnectionString options
startAction = PG.connectPostgreSQL theConnectionString
try (bracket startAction PG.close mempty) >>= \case
Left (_ :: IOError) -> threadDelay 10000 >> waitForDB options
Right () -> return ()
data ProcessConfig = ProcessConfig
{ processConfigEnvVars :: [(String, String)]
, processConfigCmdLine :: [String]
, processConfigStdIn :: Handle
, processConfigStdOut :: Handle
, processConfigStdErr :: Handle
}
startProcess
:: String
-> ProcessConfig
-> IO ProcessHandle
startProcess name ProcessConfig {..} = (\(_, _, _, x) -> x) <$>
createProcess_ name (proc name processConfigCmdLine)
{ std_err = UseHandle processConfigStdErr
, std_out = UseHandle processConfigStdOut
, std_in = UseHandle processConfigStdIn
, env = Just processConfigEnvVars
}
executeProcess
:: String
-> ProcessConfig
-> IO ExitCode
executeProcess name = startProcess name >=> waitForProcess
data PostgresPlan = PostgresPlan
{ postgresPlanProcessConfig :: ProcessConfig
, postgresPlanClientConfig :: PostgresClient.Options
}
data PostgresProcess = PostgresProcess
{ postgresProcessClientConfig :: PostgresClient.Options
, postgresProcessHandle :: ProcessHandle
}
terminateConnections :: PostgresClient.Options-> IO ()
terminateConnections options = do
let theConnectionString = PostgresClient.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 $ PostgresClient.dbname options]
case e of
Left (_ :: IOError) -> pure ()
Right _ -> pure ()
stopPostgresProcess :: PostgresProcess -> IO ExitCode
stopPostgresProcess PostgresProcess{..} = do
withProcessHandle postgresProcessHandle $ \case
OpenHandle p -> do
terminateConnections postgresProcessClientConfig
signalProcess sigINT p
OpenExtHandle {} -> pure ()
ClosedHandle _ -> return ()
exitCode <- waitForProcess postgresProcessHandle
pure exitCode
startPostgresProcess :: Logger -> PostgresPlan -> IO PostgresProcess
startPostgresProcess logger PostgresPlan {..} = do
logger StartPostgres
let startAction = PostgresProcess postgresPlanClientConfig
<$> startProcess "postgres" postgresPlanProcessConfig
bracketOnError startAction stopPostgresProcess $
\result@PostgresProcess {..} -> do
let checkForCrash = do
mExitCode <- getProcessExitCode postgresProcessHandle
for_ mExitCode (throwIO . StartPostgresFailed)
logger WaitForDB
let options = postgresPlanClientConfig
{ PostgresClient.dbname = pure "template1"
}
waitForDB options
`race_` forever (checkForCrash >> threadDelay 100000)
return result
data Plan = Plan
{ planLogger :: Logger
, planInitDb :: Maybe ProcessConfig
, planCreateDb :: Maybe ProcessConfig
, planPostgres :: PostgresPlan
, planConfig :: String
, planDataDirectory :: FilePath
}
throwIfNotSuccess :: Exception e => (ExitCode -> e) -> ExitCode -> IO ()
throwIfNotSuccess f = \case
ExitSuccess -> pure ()
e -> throwIO $ f e
initPlan :: Plan -> IO PostgresProcess
initPlan Plan {..} = do
for_ planInitDb $ executeProcess "initdb" >=>
throwIfNotSuccess InitDbFailed
writeFile (planDataDirectory <> "/postgresql.conf") planConfig
let startAction = startPostgresProcess planLogger planPostgres
bracketOnError startAction stopPostgresProcess $ \result -> do
for_ planCreateDb $ executeProcess "createdb" >=>
throwIfNotSuccess CreateDbFailed
pure result
shutdownPlan :: PostgresProcess -> IO ExitCode
shutdownPlan = stopPostgresProcess