Reusable PostgreSQL TestContainer

PostgreSQL.hs

module TestUtil.PostgreSQL
  ( PostgreSQL (..)
  , withPostgreSQL
  , withSharedPostgreSQL
  , runQueryBeforeAll
  , runQueryBefore
  , stopAllContainers
  ) where

import Control.Monad.Logger (MonadLoggerIO, runNoLoggingT, runStderrLoggingT)
import Control.Monad.Trans.Resource
  ( getInternalState
  , liftResourceT
  )
import Control.Monad.Trans.Resource.Internal
  ( ReleaseMap
  , stateAlloc
  , stateCleanup
  )
import Data.Acquire (ReleaseType (ReleaseNormal))
import Database.Persist.Postgresql (ConnectionString, withPostgresqlPool)
import Database.Persist.Sql (ConnectionPool, Migration, SqlBackend, addMigration, runMigration, runMigrationQuiet, runSqlPool)
import RIO
import RIO.Map qualified as M
import RIO.Text qualified as T
import RIO.Text.Lazy qualified as TL
import System.IO.Unsafe (unsafePerformIO)
import Test.Hspec
import TestContainers.Docker qualified as TC
import TestContainers.Monad (runTestContainer)
import TestUtil (requireVerbose)

--------------------------------------------------------------------------------

data PostgreSQL = PostgreSQL
  { connectionString :: ConnectionString
  , pool :: ConnectionPool
  , container :: TC.Container
  }

-- | 配下の各テストに対してPostgreSQLのコンテナを起動する。
withPostgreSQL
  :: Migration
  -> SpecWith PostgreSQL
  -> Spec
withPostgreSQL = withPostgreSQLGeneric around after id

-- | 配下のテストで共有されるPostgreSQLのコンテナを起動する。
withSharedPostgreSQL
  :: Migration
  -> SpecWith PostgreSQL
  -> Spec
withSharedPostgreSQL = withPostgreSQLGeneric aroundAll afterAll sequential

--------------------------------------------------------------------------------
-- Hook

-- |
-- @spec@の前にクエリを実行する。
-- テーブルのリセットに使用することを想定。
runQueryBeforeAll
  :: ReaderT SqlBackend IO ()
  -> SpecWith PostgreSQL
  -> SpecWith PostgreSQL
runQueryBeforeAll query spec =
  flip beforeAllWith spec $ \psql -> do
    runSqlPool query psql.pool
    pure psql

-- |
-- @spec@内の各テストの前にクエリを実行する。
-- テーブルのリセットに使用することを想定。
runQueryBefore
  :: ReaderT SqlBackend IO ()
  -> SpecWith PostgreSQL
  -> SpecWith PostgreSQL
runQueryBefore query spec =
  flip beforeWith spec $ \psql -> do
    runSqlPool query psql.pool
    pure psql

--------------------------------------------------------------------------------
-- コンテナ停止（mainから呼ぶ）

stopAllContainers :: IO ()
stopAllContainers = do
  m <- readIORef containerPool
  forM_ (M.elems m) $ \c -> do
    stopContainer c.releaseMap

--------------------------------------------------------------------------------
-- Implementation

withPostgreSQLGeneric
  :: ((ActionWith PostgreSQL -> IO ()) -> SpecWith PostgreSQL -> Spec)
  -> (ActionWith PostgreSQL -> SpecWith PostgreSQL -> SpecWith PostgreSQL)
  -> (SpecWith PostgreSQL -> SpecWith PostgreSQL)
  -> Migration
  -> SpecWith PostgreSQL
  -> Spec
withPostgreSQLGeneric around' after' sequencialOrNot setup action = do
  verbose <- runIO requireVerbose
  around'
    do
      \k ->
        withPostgreSQLNewOrReuse verbose \psql -> do
          runSqlPool (migrate verbose) psql.pool
          k psql
    do
      after' markAsUnused $
        sequencialOrNot action
  where
    withPostgreSQLNewOrReuse verbose k = do
      getReusableContainer >>= \case
        Nothing -> do
          ((container, connectionString), releaseMap) <- startContainer (spawnPostgreSQL verbose)
          addToContainerPool container connectionString releaseMap
          withPostgreSQL' (container, connectionString) verbose $ \psql -> do
            runSqlPool (migrate verbose) psql.pool
            k psql
        Just c -> do
          withPostgreSQL' (c.container, c.connectionString) verbose $ \psql -> do
            runSqlPool (dropAll >> migrate verbose) psql.pool
            k psql

    withPostgreSQL' (container, connectionString) verbose k = do
      let m :: (MonadUnliftIO m, MonadLoggerIO m) => m ()
          m =
            withPostgresqlPool connectionString 24 $ -- 24並列実行くらい想定すれば十分でしょう
              \connectionPool -> do
                let psql =
                      PostgreSQL
                        { pool = connectionPool
                        , container
                        , connectionString
                        }
                liftIO $ k psql
      if verbose then runStderrLoggingT m else runNoLoggingT m

    migrate :: (MonadIO m) => Bool -> ReaderT SqlBackend m ()
    migrate verbose
      | verbose = runMigration setup
      | otherwise = void $ runMigrationQuiet setup

    dropAll :: (MonadIO m) => ReaderT SqlBackend m ()
    dropAll = void $ runMigrationQuiet do
      addMigration
        False
        "DO $$ DECLARE\
        \  r RECORD;\
        \BEGIN\
        \  FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = current_schema()) LOOP\
        \    EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';\
        \  END LOOP;\
        \END $$;"

spawnPostgreSQL :: Bool -> TC.TestContainer (TC.Container, ByteString)
spawnPostgreSQL verbose = do
  container <-
    TC.run $
      TC.containerRequest (TC.fromDockerfile dockerfile)
        & TC.setExpose [5432]
        & TC.setEnv [("POSTGRES_HOST_AUTH_METHOD", "trust")]
        & TC.setWaitingFor (TC.waitForLogLine TC.Stderr (logStr `TL.isInfixOf`))
        & if verbose then TC.withFollowLogs TC.consoleLogConsumer else id
  let port = TC.containerPort container 5432
      connectionString =
        fromString $ "postgresql://postgres@localhost:" <> show port <> "/postgres"
  pure (container, connectionString)
  where
    logStr = "database system is ready to accept connections"
    -- DB初期化済みイメージを作る
    dockerfile =
      T.unlines
        [ "FROM postgres:15.8"
        , "ENV POSTGRES_HOST_AUTH_METHOD=trust"
        , "RUN [\"/bin/bash\", \"-o\", \"pipefail\", \"-c\", \\"
        , "  \"docker-entrypoint.sh postgres 2>&1 1>/dev/null | \\"
        , "  {  grep -q 'database system is ready to accept connections' \\"
        , "  && su postgres -c 'pg_ctl -D /var/lib/postgresql/data stop'; \\"
        , "  }\"]"
        , "USER postgres"
        , "ENTRYPOINT [\"docker-entrypoint.sh\", \"postgres\"]"
        ]

data Container = Container
  { container :: TC.Container
  , connectionString :: ConnectionString
  , releaseMap :: IORef ReleaseMap
  , used :: Bool
  }

{-# NOINLINE containerPool #-}
containerPool :: IORef (Map ConnectionString Container)
containerPool = unsafePerformIO $ newIORef mempty


addToContainerPool :: TC.Container -> ConnectionString -> IORef ReleaseMap -> IO ()
addToContainerPool container connectionString releaseMap = atomicModifyIORef' containerPool $
  \m ->
    ( M.insert
        connectionString
        Container
          { container
          , connectionString
          , releaseMap = releaseMap
          , used = True
          }
        m
    , ()
    )

markAsUnused :: PostgreSQL -> IO ()
markAsUnused psql = atomicModifyIORef' containerPool $
  \m -> (M.adjust (\c -> c {used = False}) psql.connectionString m, ())

getReusableContainer :: IO (Maybe Container)
getReusableContainer = do
  atomicModifyIORef' containerPool $ \m ->
    case M.assocs (M.filterWithKey (\_ c -> not c.used) m) of
      [] -> (m, Nothing)
      (k, c) : _ -> (M.insert k c' m, Just c')
        where
          c' = c {used = True}

startContainer :: TC.TestContainer a -> IO (a, IORef ReleaseMap)
startContainer m = do
  config <- TC.determineConfig
  runTestContainer config do
    result <- m
    releaseMap <- liftResourceT getInternalState
    liftIO $ stateAlloc releaseMap
    pure (result, releaseMap)

stopContainer :: IORef ReleaseMap -> IO ()
stopContainer = stateCleanup ReleaseNormal