diff --git a/src/Database/Postgres/Temp.hs b/src/Database/Postgres/Temp.hs --- a/src/Database/Postgres/Temp.hs +++ b/src/Database/Postgres/Temp.hs @@ -25,7 +25,6 @@ -} -{-# LANGUAGE RecordWildCards, LambdaCase, ScopedTypeVariables #-} module Database.Postgres.Temp ( -- * Types DB (..) @@ -36,6 +35,10 @@ , startLocalhost , startAndLogToTmp , startWithHandles + , Options(..) + , defaultOptions + , InitDbOptions(..) + , defaultInitDbOptions -- * Stopping @postgres@ , stop -- * Starting and Stopping postgres without removing the temporary directory diff --git a/src/Database/Postgres/Temp/Internal.hs b/src/Database/Postgres/Temp/Internal.hs --- a/src/Database/Postgres/Temp/Internal.hs +++ b/src/Database/Postgres/Temp/Internal.hs @@ -1,25 +1,39 @@ {-# LANGUAGE RecordWildCards, LambdaCase, ScopedTypeVariables, DeriveDataTypeable #-} {-# LANGUAGE DeriveGeneric, OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} + module Database.Postgres.Temp.Internal where -import System.IO.Temp -import System.Process -import System.Process.Internals -import Control.Concurrent -import System.IO -import System.Exit -import System.Directory -import qualified Network.Socket as N -import Control.Exception -import Data.Typeable -import GHC.Generics -import System.Posix.Signals -import qualified Database.PostgreSQL.Simple as PG + +import Control.Concurrent (MVar, newMVar, threadDelay, withMVar) +import Control.Concurrent.Async (race_) +import Control.Exception (Exception, IOException, bracket, bracketOnError, catch, onException, throwIO, try) +import Control.Monad (forever, void) import qualified Data.ByteString.Char8 as BSC -import Control.Monad (void, forever) +import Data.Foldable (for_) +import Data.IORef (IORef, newIORef, readIORef, writeIORef) +import Data.Maybe (catMaybes) +import Data.Typeable (Typeable) +import qualified Database.PostgreSQL.Simple as PG +import qualified Database.PostgreSQL.Simple.Options as Options +import GHC.Generics (Generic) +import qualified Network.Socket as N import Network.Socket.Free (openFreePort) -import Data.Foldable -import Control.Concurrent.Async(race_) -import Data.IORef +import System.Directory (removeDirectoryRecursive) +import System.Exit (ExitCode(..)) +import System.IO (Handle, IOMode(WriteMode), openFile, stderr, stdout) +import System.IO.Temp (createTempDirectory) +import System.Posix.Signals (sigHUP, sigINT, signalProcess) +import System.Process (getPid, getProcessExitCode, proc, waitForProcess) +import System.Process.Internals + ( CreateProcess + , ProcessHandle + , ProcessHandle__(..) + , StdStream(UseHandle) + , createProcess_ + , std_err + , std_out + , withProcessHandle + ) getFreePort :: IO Int getFreePort = do @@ -27,11 +41,11 @@ N.close socket pure port -waitForDB :: String -> IO () -waitForDB connStr = do - eresult <- try $ bracket (PG.connectPostgreSQL (BSC.pack connStr)) PG.close $ \_ -> return () +waitForDB :: Options.Options -> IO () +waitForDB options = do + eresult <- try $ bracket (PG.connectPostgreSQL (Options.toConnectionString options)) PG.close $ \_ -> return () case eresult of - Left (_ :: IOError) -> threadDelay 10000 >> waitForDB connStr + Left (_ :: IOError) -> threadDelay 10000 >> waitForDB options Right _ -> return () -- A helper for dealing with locks @@ -41,7 +55,7 @@ data DB = DB { mainDir :: FilePath -- ^ Temporary directory where the unix socket, logs and data directory live. - , connectionString :: String + , options :: Options.Options -- ^ PostgreSQL connection string. , extraOptions :: [(String, String)] -- ^ Additionally options passed to the postgres command line @@ -59,19 +73,38 @@ -- ^ The process handle for the @postgres@ process. } +connectionString :: DB -> String +connectionString = BSC.unpack . Options.toConnectionString . options + data SocketClass = Localhost | Unix deriving (Show, Eq, Read, Ord, Enum, Bounded, Generic, Typeable) +-- ^ +data Options = Options { + tmpDbName :: String + -- ^ The database name to use. Defaults to 'test' + , tmpInitDbOptions :: InitDbOptions + -- ^ Options to pass to initdb + , tmpCmdLineOptions :: [(String, String)] + -- ^ Extra options which override the defaults +} + +defaultOptions :: Options +defaultOptions = Options { + tmpDbName = "test" + , tmpInitDbOptions = defaultInitDbOptions + , tmpCmdLineOptions = [] +} + -- | start postgres and use the current processes stdout and stderr -start :: [(String, String)] +start :: Options -- ^ Extra options which override the defaults -> IO (Either StartError DB) start options = startWithHandles Unix options stdout stderr -- | start postgres and use the current processes stdout and stderr -- but use TCP on localhost instead of a unix socket. -startLocalhost :: [(String, String)] - -- ^ Extra options which override the defaults +startLocalhost :: Options -> IO (Either StartError DB) startLocalhost options = startWithHandles Localhost options stdout stderr @@ -114,7 +147,7 @@ pidString :: ProcessHandle -> IO String pidString phandle = withProcessHandle phandle (\case OpenHandle p -> return $ show p - OpenExtHandle _ _ _ -> return "" -- TODO log windows is not supported + OpenExtHandle {} -> return "" -- TODO log windows is not supported ClosedHandle _ -> return "" ) @@ -125,7 +158,7 @@ -- | Start postgres and pass in handles for stdout and stderr startWithHandles :: SocketClass - -> [(String, String)] + -> Options -- ^ Extra options which override the defaults -> Handle -- ^ @stdout@ @@ -137,7 +170,7 @@ startWithHandlesAndDir socketClass options mainDir stdOut stdErr startWithHandlesAndDir :: SocketClass - -> [(String, String)] + -> Options -> FilePath -> Handle -> Handle @@ -153,7 +186,7 @@ -- A helper that attempts to blocks until a connection can be made, throws -- 'StartPostgresFailed' if the postgres process fails or throws --- 'StartPostgresDisappeared' if the 'pid' somehow becomes 'Nothinng'. +-- 'StartPostgresDisappeared' if the 'pid' somehow becomes 'Nothing'. waitOnPostgres :: DB -> IO () waitOnPostgres DB {..} = do let postgresOptions = makePostgresOptions extraOptions (mainDir ++ "/data") port @@ -162,13 +195,7 @@ Just thePid -> do mExitCode <- getProcessExitCode thePid for_ mExitCode (throwIO . StartPostgresFailed postgresOptions) - host = case socketClass of - Localhost -> "127.0.0.1" - Unix -> mainDir - makeConnectionString dbName = "postgresql:///" - ++ dbName ++ "?host=" ++ host ++ "&port=" ++ show port - - waitForDB (makeConnectionString "template1") `race_` + waitForDB options `race_` forever (checkForCrash >> threadDelay 100000) -- | Send the SIGHUP signal to the postgres process to start a config reload @@ -186,7 +213,7 @@ -- If the postgres process becomes 'Nothing' while starting -- this function throws 'StartPostgresDisappeared'. startPostgres :: DB -> IO () -startPostgres db@DB {..} = withLock pidLock $ do +startPostgres db@DB {..} = withLock pidLock $ readIORef pid >>= \case Just _ -> throwIO AnotherPostgresProcessActive Nothing -> do @@ -212,7 +239,7 @@ terminateConnections db signalProcess sigINT p - OpenExtHandle _ _ _ -> pure () -- TODO log windows is not supported + OpenExtHandle {} -> pure () -- TODO log windows is not supported ClosedHandle _ -> return () ) @@ -232,7 +259,7 @@ -> Handle -> [String] -> IO ProcessHandle -runPostgres theStdOut theStdErr postgresOptions = do +runPostgres theStdOut theStdErr postgresOptions = fmap fourth $ createProcess_ "postgres" $ procWith theStdOut theStdErr "postgres" postgresOptions @@ -252,17 +279,19 @@ startWithLogger :: (Event -> IO ()) -> SocketClass - -> [(String, String)] + -> Options -> FilePath -> Handle -> Handle -> IO (Either StartError DB) -startWithLogger logger socketClass options mainDir stdOut stdErr = try $ flip onException (rmDirIgnoreErrors mainDir) $ do +startWithLogger logger socketClass (Options {..}) mainDir stdOut stdErr = + try $ flip onException (rmDirIgnoreErrors mainDir) $ do + let dataDir = mainDir ++ "/data" logger InitDB initDBExitCode <- runProcessWith stdOut stdErr "initdb" - "initdb" ["-E", "UNICODE", "-A", "trust", "--nosync", "-D", dataDir] + "initdb" $ initDbToCommandLingArgs tmpInitDbOptions { initDbPgData = Just dataDir } throwIfError InitDBFailed initDBExitCode logger WriteConfig @@ -271,20 +300,24 @@ logger FreePort port <- getFreePort -- slight race here, the port might not be free anymore! - let host = case socketClass of + let user = initDbUser tmpInitDbOptions + host = case socketClass of Localhost -> "127.0.0.1" Unix -> mainDir - let makeConnectionString dbName = "postgresql:///" - ++ dbName ++ "?host=" ++ host ++ "&port=" ++ show port - connectionString = makeConnectionString "test" + let mkOptions dbName = (Options.defaultOptions dbName) { + Options.oHost = Just host + , Options.oPort = Just port + , Options.oUser = user + } + logger StartPostgres pidLock <- newMVar () - let postgresOptions = makePostgresOptions options dataDir port + let postgresOptions = makePostgresOptions tmpCmdLineOptions dataDir port createDBResult = do thePid <- runPostgres stdOut stdErr postgresOptions pid <- newIORef $ Just thePid - pure $ DB mainDir connectionString options stdErr stdOut pidLock port socketClass pid + pure $ DB mainDir (mkOptions tmpDbName) tmpCmdLineOptions stdErr stdOut pidLock port socketClass pid bracketOnError createDBResult stop $ \result -> do let checkForCrash = readIORef (pid result) >>= \case @@ -294,7 +327,7 @@ for_ mExitCode (throwIO . StartPostgresFailed postgresOptions) logger WaitForDB - waitForDB (makeConnectionString "template1") `race_` + waitForDB (mkOptions "template1") `race_` forever (checkForCrash >> threadDelay 100000) logger CreateDB @@ -302,7 +335,8 @@ Unix -> ["-h", mainDir] Localhost -> ["-h", "127.0.0.1"] - let createDBArgs = createDBHostArgs ++ ["-p", show port, "test"] + createDBUserArg = maybe [] (\u->["--username="++u]) user + createDBArgs = createDBHostArgs ++ ["-p", show port, tmpDbName] ++ createDBUserArg throwIfError (CreateDBFailed createDBArgs) =<< runProcessWith stdOut stdErr "createDB" "createdb" createDBArgs @@ -310,7 +344,7 @@ return result -- | Start postgres and log it's all stdout to {'mainDir'}\/output.txt and {'mainDir'}\/error.txt -startAndLogToTmp :: [(String, String)] +startAndLogToTmp :: Options -- ^ Extra options which override the defaults -> IO (Either StartError DB) startAndLogToTmp options = do @@ -324,11 +358,12 @@ -- | Force all connections to the database to close. Can be useful in some testing situations. -- Called during shutdown as well. terminateConnections :: DB -> IO () -terminateConnections DB {..} = do - e <- try $ bracket (PG.connectPostgreSQL $ BSC.pack connectionString) +terminateConnections db@DB {..} = do + e <- try $ bracket (PG.connectPostgreSQL $ BSC.pack $ connectionString db) PG.close $ \conn -> do - void $ PG.execute_ conn "select pg_terminate_backend(pid) from pg_stat_activity where datname='test';" + let q = "select pg_terminate_backend(pid) from pg_stat_activity where datname=?;" + void $ PG.execute conn q [Options.oDbname options] case e of Left (_ :: IOError) -> pure () -- expected Right _ -> pure () -- Surprising ... but I do not know yet if this is a failure of termination or not. @@ -339,3 +374,40 @@ result <- stopPostgres db removeDirectoryRecursive mainDir return result + +data InitDbOptions = InitDbOptions + { initDbUser :: Maybe String + , initDbEncoding :: Maybe String + , initDbAuth :: Maybe String + , initDbPgData :: Maybe String + , initDbNoSync :: Bool + , initDbDebug :: Bool + , initDbExtraOptions :: [String] + } deriving (Show, Eq, Read, Ord, Generic, Typeable) + +defaultInitDbOptions :: InitDbOptions +defaultInitDbOptions = InitDbOptions { + initDbUser = Nothing + , initDbEncoding = Just "UNICODE" + , initDbAuth = Just "trust" + , initDbPgData = Nothing + , initDbNoSync = True + , initDbDebug = False + , initDbExtraOptions = [] +} + +initDbToCommandLingArgs :: InitDbOptions -> [String] +initDbToCommandLingArgs InitDbOptions {..} = strArgs <> boolArgs <> initDbExtraOptions + where + strArgs = fmap (\(a,b) -> "--" <> a <> "=" <> b) . catMaybes $ strength <$> + [ ("username", initDbUser) + , ("encoding", initDbEncoding) + , ("auth", initDbAuth) + , ("pgdata", initDbPgData) + ] + boolArgs = fmap fst . filter snd $ + [ ("--nosync", initDbNoSync) + , ("--debug", initDbDebug) ] + + strength :: Functor f => (a, f b) -> f (a, b) + strength (a, fb) = (a,) <$> fb diff --git a/test/Database/Postgres/Temp/InternalSpec.hs b/test/Database/Postgres/Temp/InternalSpec.hs --- a/test/Database/Postgres/Temp/InternalSpec.hs +++ b/test/Database/Postgres/Temp/InternalSpec.hs @@ -14,6 +14,12 @@ import System.Exit import System.Timeout(timeout) import Data.Either +{- +import qualified Database.PostgreSQL.LibPQ as LibPQ +import qualified Database.PostgreSQL.Simple.Internal as PS +import Control.Monad.Reader +import Data.Traversable +-} mkDevNull :: IO Handle mkDevNull = openFile "/dev/null" WriteMode @@ -24,8 +30,13 @@ instance Exception Except countPostgresProcesses :: IO Int -countPostgresProcesses = length . lines <$> readProcess "pgrep" ["postgres"] [] +countPostgresProcesses = do + (exitCode, xs, _) <- readProcessWithExitCode "pgrep" ["postgres"] [] + unless (exitCode == ExitSuccess || exitCode == ExitFailure 1) $ throwIO exitCode + + pure $ length $ lines xs + spec :: Spec spec = describe "Database.Postgres.Temp.Internal" $ do before (createTempDirectory "/tmp" "tmp-postgres") $ after rmDirIgnoreErrors $ describe "startWithLogger/stop" $ do @@ -36,7 +47,7 @@ theStdOut <- mkDevNull theStdErr <- mkDevNull shouldThrow - (startWithLogger (\currentEvent -> when (currentEvent == event) $ throwIO Except) Unix [] mainFilePath theStdOut theStdErr) + (startWithLogger (\currentEvent -> when (currentEvent == event) $ throwIO Except) Unix defaultOptions mainFilePath theStdOut theStdErr) (\Except -> True) doesDirectoryExist mainFilePath `shouldReturn` False countPostgresProcesses `shouldReturn` beforePostgresCount @@ -45,7 +56,7 @@ beforePostgresCount <- countPostgresProcesses theStdOut <- mkDevNull theStdErr <- mkDevNull - result <- startWithLogger (\_ -> return ()) Unix [] mainFilePath theStdOut theStdErr + result <- startWithLogger (\_ -> return ()) Unix defaultOptions mainFilePath theStdOut theStdErr db <- case result of Right x -> return x Left err -> error $ show err @@ -61,7 +72,7 @@ theStdOut <- mkDevNull theStdErr <- mkDevNull bracket (startWithLogger (const $ pure ()) Unix - [("log_min_duration_statement", expectedDuration)] + defaultOptions { tmpCmdLineOptions = [("log_min_duration_statement", expectedDuration)] } mainFilePath theStdOut theStdErr ) (either (\_ -> return ()) (void . stop)) $ \result -> do @@ -76,9 +87,9 @@ theStdOut <- mkDevNull theStdErr <- mkDevNull r <- timeout 5000000 $ startWithLogger (const $ pure ()) Unix - [ ("log_directory", "/this/does/not/exist") + defaultOptions { tmpCmdLineOptions = [ ("log_directory", "/this/does/not/exist") , ("logging_collector", "true") - ] mainFilePath theStdOut theStdErr + ] } mainFilePath theStdOut theStdErr case r of Nothing -> -- bad test, shouldSatisfy is difficult because it wants Show on DB. @@ -95,24 +106,25 @@ it "terminateConnections" $ \mainFilePath -> do theStdOut <- mkDevNull theStdErr <- mkDevNull - bracket (fromRight (error "failed to start db") <$> startWithLogger (\_ -> return ()) Unix [] mainFilePath theStdOut theStdErr) stop $ \db -> do + bracket (fromRight (error "failed to start db") <$> startWithLogger (\_ -> return ()) Unix defaultOptions mainFilePath theStdOut theStdErr) stop $ \db -> do bracket (connectPostgreSQL $ BSC.pack $ connectionString db) close $ \_ -> bracket (connectPostgreSQL $ BSC.pack $ connectionString db) close $ \conn2 -> do - query_ conn2 "SELECT COUNT(*) FROM pg_stat_activity" `shouldReturn` [Only (2 :: Int)] + query_ conn2 "SELECT COUNT(*) FROM pg_stat_activity where backend_type='client backend'" `shouldReturn` [Only (2 :: Int)] terminateConnections db bracket (connectPostgreSQL $ BSC.pack $ connectionString db) close $ \conn3 -> - query_ conn3 "SELECT COUNT(*) FROM pg_stat_activity" `shouldReturn` [Only (1 :: Int)] + query_ conn3 "SELECT COUNT(*) FROM pg_stat_activity where backend_type='client backend'" `shouldReturn` [Only (1 :: Int)] -- The point of stopPostgres and startPostgres is to test recovery so -- let's make sure that works it "stopPostgres/startPostgres works" $ \mainFilePath -> do - let extraOpts = + let extraOpts = defaultOptions { tmpCmdLineOptions = [ ("wal_level", "replica") , ("archive_mode", "on") , ("max_wal_senders", "2") ] + } theStdOut <- mkDevNull theStdErr <- mkDevNull @@ -138,8 +150,8 @@ _ <- execute_ conn "CREATE TABLE foo(id int PRIMARY KEY);" _ <- execute_ conn "BEGIN ISOLATION LEVEL READ COMMITTED READ WRITE; INSERT INTO foo (id) VALUES (1); COMMIT" - _ :: [Only String] <- query_ conn "SELECT pg_xlogfile_name(pg_switch_xlog())" - _ :: [Only String] <- query_ conn "SELECT pg_xlogfile_name(pg_create_restore_point('pitr'))" + _ :: [Only String] <- query_ conn "SELECT pg_walfile_name(pg_switch_wal())" + _ :: [Only String] <- query_ conn "SELECT pg_walfile_name(pg_create_restore_point('pitr'))" _ <- execute_ conn "BEGIN ISOLATION LEVEL READ COMMITTED READ WRITE; INSERT INTO foo (id) VALUES (2); COMMIT" query_ conn "SELECT id FROM foo ORDER BY id ASC" diff --git a/tmp-postgres.cabal b/tmp-postgres.cabal --- a/tmp-postgres.cabal +++ b/tmp-postgres.cabal @@ -1,5 +1,5 @@ name: tmp-postgres -version: 0.2.0.0 +version: 0.3.0.0 synopsis: Start and stop a temporary postgres for testing description: This module provides functions creating a temporary postgres instance on a random port for testing. @@ -49,6 +49,7 @@ , network , bytestring , postgresql-simple + , postgres-options , port-utils , async ghc-options: -Wall @@ -60,6 +61,7 @@ main-is: Spec.hs other-modules: Database.Postgres.Temp.InternalSpec build-depends: base + , postgresql-libpq , tmp-postgres , hspec , hspec-discover @@ -68,6 +70,7 @@ , process , postgresql-simple , bytestring + , mtl ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N default-language: Haskell2010