module Main where
import Control.Concurrent.Async (race)
import qualified Data.ByteString.Char8 as B8
import qualified Data.Text as Text
import qualified Hasql.Connection as Connection
import qualified Hasql.Decoders as Decoders
import qualified Hasql.Encoders as Encoders
import Hasql.Pool
import qualified Hasql.Pool.Config as Config
import qualified Hasql.Session as Session
import qualified Hasql.Statement as Statement
import qualified System.Environment
import qualified System.Random.Stateful as Random
import Test.Hspec
import Prelude
main :: IO ()
main = do
connectionSettings <- getConnectionSettings
let withPool poolSize acqTimeout maxLifetime maxIdletime connectionSettings =
bracket
( acquire
( Config.settings
[ Config.size poolSize,
Config.acquisitionTimeout acqTimeout,
Config.agingTimeout maxLifetime,
Config.idlenessTimeout maxIdletime,
Config.staticConnectionSettings connectionSettings
]
)
)
release
withDefaultPool =
withPool 3 10 1_800 1_800 connectionSettings
hspec . describe "" $ do
it "Releases a spot in the pool when there is a query error" $ withDefaultPool $ \pool -> do
use pool badQuerySession `shouldNotReturn` (Right ())
use pool selectOneSession `shouldReturn` (Right 1)
it "Simulation of connection error works" $ withDefaultPool $ \pool -> do
res <- use pool $ closeConnSession >> selectOneSession
shouldSatisfy res $ \case
Left (SessionUsageError (Session.QueryError _ _ (Session.ClientError _))) -> True
_ -> False
it "Connection errors cause eviction of connection" $ withDefaultPool $ \pool -> do
res <- use pool $ closeConnSession >> selectOneSession
res <- use pool $ closeConnSession >> selectOneSession
res <- use pool $ closeConnSession >> selectOneSession
res <- use pool $ selectOneSession
shouldSatisfy res $ isRight
it "Connection gets returned to the pool after normal use" $ withDefaultPool $ \pool -> do
res <- use pool $ selectOneSession
res <- use pool $ selectOneSession
res <- use pool $ selectOneSession
res <- use pool $ selectOneSession
res <- use pool $ selectOneSession
shouldSatisfy res $ isRight
it "Connection gets returned to the pool after non-connection error" $ withDefaultPool $ \pool -> do
res <- use pool $ badQuerySession
res <- use pool $ badQuerySession
res <- use pool $ badQuerySession
res <- use pool $ badQuerySession
res <- use pool $ selectOneSession
shouldSatisfy res $ isRight
it "The pool remains usable after release" $ withDefaultPool $ \pool -> do
res <- use pool $ selectOneSession
release pool
res <- use pool $ selectOneSession
shouldSatisfy res $ isRight
it "Getting and setting session variables works" $ withDefaultPool $ \pool -> do
res <- use pool $ getSettingSession "testing.foo"
res `shouldBe` Right Nothing
res <- use pool $ do
setSettingSession "testing.foo" "hello world"
getSettingSession "testing.foo"
res `shouldBe` Right (Just "hello world")
it "Session variables stay set when a connection gets reused" $ withPool 1 10 1_800 1_800 connectionSettings $ \pool -> do
res <- use pool $ setSettingSession "testing.foo" "hello world"
res `shouldBe` Right ()
res2 <- use pool $ getSettingSession "testing.foo"
res2 `shouldBe` Right (Just "hello world")
it "Releasing the pool resets session variables" $ withPool 1 10 1_800 1_800 connectionSettings $ \pool -> do
res <- use pool $ setSettingSession "testing.foo" "hello world"
res `shouldBe` Right ()
release pool
res <- use pool $ getSettingSession "testing.foo"
res `shouldBe` Right Nothing
it "Times out connection acquisition"
$
-- 1ms timeout
withPool 1 0.001 1_800 1_800 connectionSettings
$ \pool -> do
sleeping <- newEmptyMVar
t0 <- getCurrentTime
res <-
race
( use pool
$ liftIO
$ do
putMVar sleeping ()
threadDelay 1_000_000 -- 1s
)
( do
takeMVar sleeping
use pool $ selectOneSession
)
t1 <- getCurrentTime
res `shouldBe` Right (Left AcquisitionTimeoutUsageError)
diffUTCTime t1 t0 `shouldSatisfy` (< 0.5) -- 0.5s
it "Passively times out old connections (maxLifetime)"
$
-- 0.5s connection lifetime
withPool 1 10 0.5 1_800 connectionSettings
$ \pool -> do
res <- use pool $ setSettingSession "testing.foo" "hello world"
res `shouldBe` Right ()
res2 <- use pool $ getSettingSession "testing.foo"
res2 `shouldBe` Right (Just "hello world")
threadDelay 1_000_000 -- 1s
res3 <- use pool $ getSettingSession "testing.foo"
res3 `shouldBe` Right Nothing
it "Counts active connections" $ do
(taggedConnectionSettings, appName) <- tagConnection connectionSettings
withPool 3 10 1_800 1_800 taggedConnectionSettings $ \pool -> do
res <- use pool $ countConnectionsSession appName
res `shouldBe` Right 1
it "Actively times out old connections (maxLifetime)" $ do
withDefaultPool $ \countPool -> do
(taggedConnectionSettings, appName) <- tagConnection connectionSettings
withPool 3 10 0.5 1_800 taggedConnectionSettings $ \limitedPool -> do
res <- use limitedPool $ selectOneSession
res `shouldBe` Right 1
res2 <- use countPool $ countConnectionsSession appName
res2 `shouldBe` Right 1
threadDelay 1_000_000 -- 1s
res3 <- use countPool $ countConnectionsSession appName
res3 `shouldBe` Right 0
it "Times out old connections (maxIdletime)" $ do
-- 0.5s connection idle time
withPool 1 10 1_800 0.5 connectionSettings $ \pool -> do
res <- use pool $ setSettingSession "testing.foo" "hello world"
res `shouldBe` Right ()
res2 <- use pool $ getSettingSession "testing.foo"
res2 `shouldBe` Right (Just "hello world")
-- busy sleep, to keep connection alive
forM_ [1 .. 10] $ \_ -> do
r <- use pool $ selectOneSession
r `shouldBe` Right 1
threadDelay 100_000 -- 0.1s
res3 <- use pool $ getSettingSession "testing.foo"
res3 `shouldBe` Right (Just "hello world")
-- idle sleep, connection times out
threadDelay 1_000_000 -- 1s
res4 <- use pool $ getSettingSession "testing.foo"
res4 `shouldBe` Right Nothing
getConnectionSettings :: IO Connection.Settings
getConnectionSettings =
B8.unwords
. catMaybes
<$> sequence
[ setting "host" $ defaultEnv "POSTGRES_HOST" "localhost",
setting "port" $ defaultEnv "POSTGRES_PORT" "5432",
setting "user" $ defaultEnv "POSTGRES_USER" "postgres",
setting "password" $ defaultEnv "POSTGRES_PASSWORD" "postgres",
setting "dbname" $ defaultEnv "POSTGRES_DBNAME" "postgres"
]
where
maybeEnv env = fmap B8.pack <$> System.Environment.lookupEnv env
defaultEnv env val = Just . fromMaybe val <$> maybeEnv env
setting label getEnv = do
val <- getEnv
return $ (\v -> label <> "=" <> v) <$> val
tagConnection :: Connection.Settings -> IO (Connection.Settings, Text)
tagConnection connectionSettings = do
tag <- Random.uniformWord32 Random.globalStdGen
let appName = "hasql-pool-test-" <> show tag
return (connectionSettings <> " application_name=" <> B8.pack appName, Text.pack appName)
selectOneSession :: Session.Session Int64
selectOneSession =
Session.statement () statement
where
statement = Statement.Statement "SELECT 1" Encoders.noParams decoder True
decoder = Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int8))
badQuerySession :: Session.Session ()
badQuerySession =
Session.statement () statement
where
statement = Statement.Statement "zzz" Encoders.noParams Decoders.noResult True
closeConnSession :: Session.Session ()
closeConnSession = do
conn <- ask
liftIO $ Connection.release conn
setSettingSession :: Text -> Text -> Session.Session ()
setSettingSession name value = do
Session.statement (name, value) statement
where
statement = Statement.Statement "SELECT set_config($1, $2, false)" encoder Decoders.noResult True
encoder =
contramap fst (Encoders.param (Encoders.nonNullable Encoders.text))
<> contramap snd (Encoders.param (Encoders.nonNullable Encoders.text))
getSettingSession :: Text -> Session.Session (Maybe Text)
getSettingSession name = do
Session.statement name statement
where
statement = Statement.Statement "SELECT current_setting($1, true)" encoder decoder True
encoder = Encoders.param (Encoders.nonNullable Encoders.text)
decoder = Decoders.singleRow (Decoders.column (Decoders.nullable Decoders.text))
countConnectionsSession :: Text -> Session.Session Int64
countConnectionsSession appName = do
Session.statement appName statement
where
statement = Statement.Statement "SELECT count(*) FROM pg_stat_activity WHERE application_name = $1" encoder decoder True
encoder = Encoders.param (Encoders.nonNullable Encoders.text)
decoder = Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int8))