{-# OPTIONS_HADDOCK prune #-}
module Database.Postgres.Temp.Internal.Core where
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (race_, withAsync)
import Control.Exception
import Control.Monad (forever, unless, void)
import Crypto.Hash.SHA1 (hash)
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Base64.URL as Base64
import Data.Char
import Data.Foldable (for_)
import Data.IORef
import Data.Maybe
import Data.Monoid
import Data.List
import Data.String
import Data.Typeable
import qualified Database.PostgreSQL.Simple as PG
import qualified Database.PostgreSQL.Simple.Options as Client
import System.Directory
import System.Exit (ExitCode(..))
import System.IO
import System.Posix.Signals (sigQUIT, signalProcess)
import System.Process
import System.Process.Internals
import System.Timeout
import Text.PrettyPrint.ANSI.Leijen hiding ((<$>))
data Event
= StartPlan String
| StartPostgres
| WaitForDB
| TryToConnect
deriving (Show, Eq, Ord)
data StartError
= StartPostgresFailed ExitCode
| InitDbFailed
{ startErrorStdOut :: String
, startErrorStdErr :: String
, startErrorExitCode :: ExitCode
}
| CreateDbFailed
{ startErrorStdOut :: String
, startErrorStdErr :: String
, startErrorExitCode :: ExitCode
}
| CompletePlanFailed String [String]
| CompleteProcessConfigFailed String [String]
| ConnectionTimedOut
| DeleteDbError PG.SqlError
| EmptyDataDirectory
| CopyCachedInitDbFailed String ExitCode
| FailedToFindDataDirectory String
deriving (Show, Eq, 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 1000 >> waitForDB logger options
Right () -> return ()
teeHandle :: Handle -> (Handle -> IO a) -> IO (a, String)
teeHandle orig f =
bracket createPipe (\(x, y) -> hClose x >> hClose y) $ \(readEnd, writeEnd) -> do
outputRef <- newIORef []
let readerLoop = forever $ do
theLine <- hGetLine readEnd
modifyIORef outputRef (<>theLine)
hPutStrLn orig theLine
res <- withAsync readerLoop $ \_ -> f writeEnd
(res,) <$> readIORef outputRef
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)
makeCommandLine :: String -> CompleteProcessConfig -> String
makeCommandLine command CompleteProcessConfig {..} =
let envs = unwords $ map (\(x, y) -> x <> "=" <> y) completeProcessConfigEnvVars
args = unwords completeProcessConfigCmdLine
in envs <> " " <> command <> args
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 conf =
bracket (startProcess name conf) terminateProcess waitForProcess
executeProcessAndTee
:: String
-> CompleteProcessConfig
-> IO (ExitCode, String, String)
executeProcessAndTee name config = fmap (\((x, y), z) -> (x, z, y)) $
teeHandle (completeProcessConfigStdOut config) $ \newOut ->
teeHandle (completeProcessConfigStdErr config) $ \newErr ->
executeProcess name $ config
{ completeProcessConfigStdErr = newErr
, completeProcessConfigStdOut = newOut
}
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
{ Client.dbname = pure "template1"
}
terminationQuery = fromString $ unlines
[ "SELECT pg_terminate_backend(pid)"
, "FROM pg_stat_activity"
, "WHERE datname=?;"
]
e <- try $ bracket (PG.connectPostgreSQL theConnectionString) PG.close $
\conn -> PG.query conn terminationQuery
[getLast $ Client.dbname options]
case e of
Left (_ :: IOError) -> pure ()
Right (_ :: [PG.Only Bool]) -> pure ()
stopPostgresProcess :: PostgresProcess -> IO ExitCode
stopPostgresProcess PostgresProcess{..} = do
withProcessHandle postgresProcessHandle $ \case
OpenHandle p ->
signalProcess sigQUIT p
OpenExtHandle {} -> pure ()
ClosedHandle _ -> return ()
waitForProcess postgresProcessHandle
startPostgresProcess :: Int -> Logger -> CompletePostgresPlan -> IO PostgresProcess
startPostgresProcess time logger CompletePostgresPlan {..} = do
logger StartPostgres
let startAction = PostgresProcess completePostgresPlanClientOptions
<$> startProcess "postgres" completePostgresPlanProcessConfig
bracketOnError startAction stopPostgresProcess $
\result@PostgresProcess {..} -> do
logger WaitForDB
let options = completePostgresPlanClientOptions
{ Client.dbname = pure "template1"
}
checkForCrash = do
mExitCode <- getProcessExitCode postgresProcessHandle
for_ mExitCode (throwIO . StartPostgresFailed)
timeoutAndThrow = timeout time (waitForDB logger options) >>= \case
Just () -> pure ()
Nothing -> throwIO ConnectionTimedOut
timeoutAndThrow `race_` forever (checkForCrash >> threadDelay 100000)
return result
getInitDbVersion :: IO String
getInitDbVersion = readProcessWithExitCode "initdb" ["--version"] "" >>= \case
(ExitSuccess, outputString, _) -> do
let
theLastPart = last $ words outputString
versionPart = takeWhile (\x -> isDigit x || x == '.' || x == '-') theLastPart
pure $ if last versionPart == '.'
then init versionPart
else versionPart
(startErrorExitCode, startErrorStdOut, startErrorStdErr) ->
throwIO InitDbFailed {..}
makeInitDbCommandLine :: CompleteProcessConfig -> String
makeInitDbCommandLine = makeCommandLine "initdb"
makeArgumentHash :: String -> String
makeArgumentHash = BSC.unpack . Base64.encode . hash . BSC.pack
splitDataDirectory :: CompleteProcessConfig -> (Maybe String, CompleteProcessConfig)
splitDataDirectory old =
let isDataDirectoryFlag xs = "-D" `isPrefixOf` xs || "--pgdata=" `isPrefixOf` xs
(dataDirectoryArgs, otherArgs) =
partition isDataDirectoryFlag $ completeProcessConfigCmdLine old
firstDataDirectoryArg = flip fmap (listToMaybe dataDirectoryArgs) $ \case
'-':'D':' ':theDir -> theDir
'-':'D':theDir -> theDir
'-':'-':'p':'g':'d':'a':'t':'a':'=':theDir -> theDir
_ -> error "splitDataDirectory not possible"
filteredEnvs = filter (not . ("PGDATA"==) . fst) $
completeProcessConfigEnvVars old
clearedConfig = old
{ completeProcessConfigCmdLine = otherArgs
, completeProcessConfigEnvVars = filteredEnvs
}
in (firstDataDirectoryArg, clearedConfig)
makeCachePath :: FilePath -> String -> IO String
makeCachePath cacheFolder cmdLine = do
version <- getInitDbVersion
let theHash = makeArgumentHash cmdLine
pure $ cacheFolder <> "/" <> version <> "/" <> theHash
addDataDirectory :: String -> CompleteProcessConfig -> CompleteProcessConfig
addDataDirectory theDataDirectory x = x
{ completeProcessConfigCmdLine =
("--pgdata=" <> theDataDirectory) : completeProcessConfigCmdLine x
}
executeInitDb :: Maybe (Bool, FilePath) -> CompleteProcessConfig -> IO ()
executeInitDb cache config = do
let runInitDb theConfig = do
(res, stdOut, stdErr) <- executeProcessAndTee "initdb" theConfig
throwIfNotSuccess (InitDbFailed stdOut stdErr) res
void $ case cache of
Nothing -> runInitDb config
Just (copyOnWrite, directoryType) -> do
let (mtheDataDirectory, clearedConfig) = splitDataDirectory config
theDataDirectory <- maybe
(throwIO $ FailedToFindDataDirectory (show $ pretty config))
pure
mtheDataDirectory
let theCommandLine = makeInitDbCommandLine clearedConfig
cachePath <- makeCachePath directoryType theCommandLine
let newDataDirectory = cachePath <> "/data"
doesDirectoryExist cachePath >>= \case
True -> pure ()
False -> do
createDirectoryIfMissing True cachePath
writeFile (cachePath <> "/commandLine.log") theCommandLine
runInitDb $ addDataDirectory newDataDirectory clearedConfig
let
#ifdef darwin_HOST_OS
cpFlags = if copyOnWrite then "cp -Rc " else "cp -R "
#else
cpFlags = if copyOnWrite then "cp -R --reflink=auto " else "cp -R "
#endif
copyCommand = cpFlags <> newDataDirectory <> "/* " <> theDataDirectory
throwIfNotSuccess (CopyCachedInitDbFailed copyCommand) =<< system copyCommand
data CompletePlan = CompletePlan
{ completePlanLogger :: Logger
, completePlanInitDb :: Maybe CompleteProcessConfig
, completePlanCreateDb :: Maybe CompleteProcessConfig
, completePlanPostgres :: CompletePostgresPlan
, completePlanConfig :: String
, completePlanDataDirectory :: FilePath
, completePlanConnectionTimeout :: Int
, completePlanCacheDirectory :: Maybe (Bool, 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
<> hardline
<> text "completePlanCacheDirectory:"
<+> pretty completePlanCacheDirectory
throwIfNotSuccess :: Exception e => (ExitCode -> e) -> ExitCode -> IO ()
throwIfNotSuccess f = \case
ExitSuccess -> pure ()
e -> throwIO $ f e
executeCreateDb :: CompleteProcessConfig -> IO ()
executeCreateDb config = do
(res, stdOut, stdErr) <- executeProcessAndTee "createdb" config
throwIfNotSuccess (CreateDbFailed stdOut stdErr) res
startPlan :: CompletePlan -> IO PostgresProcess
startPlan plan@CompletePlan {..} = do
completePlanLogger $ StartPlan $ show $ pretty plan
for_ completePlanInitDb $ executeInitDb completePlanCacheDirectory
versionFileExists <- doesFileExist $ completePlanDataDirectory <> "/PG_VERSION"
unless versionFileExists $ throwIO EmptyDataDirectory
writeFile (completePlanDataDirectory <> "/postgresql.conf") completePlanConfig
let startAction = startPostgresProcess
completePlanConnectionTimeout completePlanLogger completePlanPostgres
bracketOnError startAction stopPostgresProcess $ \result -> do
for_ completePlanCreateDb executeCreateDb
pure result
stopPlan :: PostgresProcess -> IO ExitCode
stopPlan = stopPostgresProcess