{-# OPTIONS_HADDOCK prune #-}
module Database.Postgres.Temp.Internal.Core where
import qualified Database.PostgreSQL.Simple.Options as Client
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 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
import Text.PrettyPrint.ANSI.Leijen hiding ((<$>))
import qualified Data.ByteString.Char8 as BSC
data Event
= StartPostgres
| WaitForDB
| StartPlan String
| 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 ProcessConfig = ProcessConfig
{ processConfigEnvVars :: [(String, String)]
, processConfigCmdLine :: [String]
, processConfigStdIn :: Handle
, processConfigStdOut :: Handle
, processConfigStdErr :: Handle
}
prettyKeyPair ::(Pretty a, Pretty b) => a -> b -> Doc
prettyKeyPair k v = pretty k <> text ": " <> pretty v
instance Pretty ProcessConfig where
pretty ProcessConfig {..}
= text "processConfigEnvVars:"
<> softline
<> indent 2 (vsep (map (uncurry prettyKeyPair) processConfigEnvVars))
<> hardline
<> text "processConfigCmdLine:"
<> softline
<> text (unwords processConfigCmdLine)
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
, postgresPlanClientOptions :: Client.Options
}
instance Pretty PostgresPlan where
pretty PostgresPlan {..}
= text "postgresPlanProcessConfig:"
<> softline
<> indent 2 (pretty postgresPlanProcessConfig)
<> hardline
<> text "postgresPlanClientOptions:"
<+> prettyOptions postgresPlanClientOptions
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 -> PostgresPlan -> IO PostgresProcess
startPostgresProcess logger PostgresPlan {..} = do
logger StartPostgres
let startAction = PostgresProcess postgresPlanClientOptions
<$> startProcess "postgres" postgresPlanProcessConfig
bracketOnError startAction stopPostgresProcess $
\result@PostgresProcess {..} -> do
let checkForCrash = do
mExitCode <- getProcessExitCode postgresProcessHandle
for_ mExitCode (throwIO . StartPostgresFailed)
logger WaitForDB
let options = postgresPlanClientOptions
{ Client.dbname = pure "template1"
}
waitForDB logger 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
}
instance Pretty Plan where
pretty Plan {..}
= text "planInitDb:"
<> softline
<> indent 2 (pretty planInitDb)
<> hardline
<> text "planCreateDb:"
<> softline
<> indent 2 (pretty planCreateDb)
<> hardline
<> text "planPostgres:"
<> softline
<> indent 2 (pretty planPostgres)
<> hardline
<> text "planConfig:"
<> softline
<> indent 2 (pretty planConfig)
<> hardline
<> text "planDataDirectory:"
<+> pretty planDataDirectory
throwIfNotSuccess :: Exception e => (ExitCode -> e) -> ExitCode -> IO ()
throwIfNotSuccess f = \case
ExitSuccess -> pure ()
e -> throwIO $ f e
startPlan :: Plan -> IO PostgresProcess
startPlan plan@Plan {..} = do
planLogger $ StartPlan $ show $ pretty plan
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
stopPlan :: PostgresProcess -> IO ExitCode
stopPlan = stopPostgresProcess