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,5 +1,5 @@ {-# LANGUAGE RecordWildCards, LambdaCase, ScopedTypeVariables, DeriveDataTypeable #-} -{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveGeneric, OverloadedStrings #-} module Database.Postgres.Temp.Internal where import System.IO.Temp import System.Process @@ -8,24 +8,25 @@ import System.IO import System.Exit import System.Directory -import Network.Socket +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 SQL +import qualified Database.PostgreSQL.Simple as PG import qualified Data.ByteString.Char8 as BSC +import Control.Monad (void) +import Network.Socket.Free (openFreePort) -openFreePort :: IO Int -openFreePort = bracket (socket AF_INET Stream defaultProtocol) close $ \s -> do - localhost <- inet_addr "127.0.0.1" - bind s (SockAddrInet aNY_PORT localhost) - listen s 1 - fmap fromIntegral $ socketPort s +getFreePort :: IO Int +getFreePort = do + (port, socket) <- openFreePort + N.close socket + pure port waitForDB :: String -> IO () waitForDB connStr = do - eresult <- try $ bracket (SQL.connectPostgreSQL (BSC.pack connStr)) SQL.close $ \_ -> return () + eresult <- try $ bracket (PG.connectPostgreSQL (BSC.pack connStr)) PG.close $ \_ -> return () case eresult of Left (_ :: IOError) -> threadDelay 10000 >> waitForDB connStr Right _ -> return () @@ -80,7 +81,7 @@ data StartError = InitDBFailed ExitCode - | CreateDBFailed ExitCode + | CreateDBFailed [String] ExitCode deriving (Show, Eq, Typeable) instance Exception StartError @@ -93,6 +94,7 @@ pidString :: ProcessHandle -> IO String pidString phandle = withProcessHandle phandle (\case OpenHandle p -> return $ show p + OpenExtHandle _ _ _ -> return "" -- TODO log windows is not supported ClosedHandle _ -> return "" ) @@ -148,14 +150,14 @@ logger InitDB initDBExitCode <- runProcessWith stdOut stdErr "initdb" - "initdb" ["-E", "UNICODE", "-A", "trust", "-D", dataDir] + "initdb" ["-E", "UNICODE", "-A", "trust", "--nosync", "-D", dataDir] throwIfError InitDBFailed initDBExitCode logger WriteConfig writeFile (dataDir ++ "/postgresql.conf") $ config $ if socketType == Unix then Just mainDir else Nothing logger FreePort - port <- openFreePort + port <- getFreePort -- slight race here, the port might not be free anymore! let host = case socketType of Localhost -> "127.0.0.1" @@ -182,9 +184,9 @@ Unix -> ["-h", mainDir] Localhost -> ["-h", "127.0.0.1"] - throwIfError CreateDBFailed =<< - runProcessWith stdOut stdErr "createDB" - "createdb" (createDBHostArgs ++ ["-p", show port, "test"]) + let createDBArgs = createDBHostArgs ++ ["-p", show port, "test"] + throwIfError (CreateDBFailed createDBArgs) =<< + runProcessWith stdOut stdErr "createDB" "createdb" createDBArgs logger Finished return result @@ -201,11 +203,30 @@ startWithHandlesAndDir Unix options mainDir stdOutFile stdErrFile +-- | 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) + PG.close + $ \conn -> do + void $ PG.execute_ conn "select pg_terminate_backend(pid) from pg_stat_activity where datname='test';" + 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. + -- | Stop postgres and clean up the temporary database folder. stop :: DB -> IO ExitCode -stop DB {..} = do +stop db@DB {..} = do + withProcessHandle pid (\case - OpenHandle p -> signalProcess sigINT p + OpenHandle p -> do + -- try to terminate the connects first. If we can't terminate still + -- keep shutting down + terminateConnections db + + signalProcess sigINT p + OpenExtHandle _ _ _ -> pure () -- TODO log windows is not supported ClosedHandle _ -> return () ) 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 @@ -12,8 +12,9 @@ import Database.PostgreSQL.Simple import qualified Data.ByteString.Char8 as BSC import System.Exit -import Control.Applicative ((<$>)) +import Data.Either + main :: IO () main = hspec spec @@ -35,8 +36,10 @@ it ("deletes the temp dir and postgres on exception in " ++ show event) $ \mainFilePath -> do -- This is not the best method ... but it works beforePostgresCount <- countPostgresProcesses + stdOut <- mkDevNull + stdErr <- mkDevNull shouldThrow - (startWithLogger (\currentEvent -> when (currentEvent == event) $ throwIO Except) Unix [] mainFilePath stdout stderr) + (startWithLogger (\currentEvent -> when (currentEvent == event) $ throwIO Except) Unix [] mainFilePath stdOut stdErr) (\Except -> True) doesDirectoryExist mainFilePath `shouldReturn` False countPostgresProcesses `shouldReturn` beforePostgresCount @@ -66,3 +69,16 @@ conn <- connectPostgreSQL $ BSC.pack $ connectionString db [Only actualDuration] <- query_ conn "SHOW log_min_duration_statement" actualDuration `shouldBe` expectedDuration + + it "terminateConnections" $ \mainFilePath -> do + stdOut <- mkDevNull + stdErr <- mkDevNull + bracket (fromRight (error "failed to start db") <$> startWithLogger (\_ -> return ()) Unix [] mainFilePath stdOut stdErr) 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)] + + terminateConnections db + + bracket (connectPostgreSQL $ BSC.pack $ connectionString db) close $ \conn3 -> + query_ conn3 "SELECT COUNT(*) FROM pg_stat_activity" `shouldReturn` [Only (1 :: Int)] 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.1.1.1 +version: 0.1.2.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. @@ -17,6 +17,10 @@ The start methods use a config based on the one used by pg_tmp (http://ephemeralpg.org/), but can be overriden by different values to the first argument of the start functions. . + MacOS and Linux are support. Windows is not. + . + Requires PostgreSQL 9.3+ + . WARNING!! Ubuntu's PostgreSQL installation does not put @initdb@ on the @PATH@. We need to add it manually. The necessary binaries are in the @\/usr\/lib\/postgresql\/VERSION\/bin\/@ directory, and should be added to the @PATH@ . @@ -43,12 +47,13 @@ , unix , directory , network - , postgresql-simple , bytestring + , postgresql-simple + , port-utils ghc-options: -Wall default-language: Haskell2010 -test-suite tmp-postgres-test +test-suite test type: exitcode-stdio-1.0 hs-source-dirs: test main-is: Spec.hs