pgmq-config-0.6.1.0: test/EphemeralDb.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Test database infrastructure using ephemeral-pg
module EphemeralDb
  ( -- * Database setup
    withPgmqDb,

    -- * Re-exports
    StartError,
  )
where

import Control.Monad (filterM, when)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Monoid (Last (..))
import Data.Text.IO qualified as TextIO
import Database.PostgreSQL.Migrate
  ( defaultRunOptions,
    migrationPlan,
    runMigrationPlan,
  )
import EphemeralPg
  ( Config (temporaryRoot),
    StartError,
    connectionSettings,
    defaultCacheConfig,
    defaultConfig,
    withCachedConfig,
  )
import Hasql.Pool qualified as Pool
import Hasql.Pool.Config qualified as PoolConfig
import Hasql.Session qualified as Session
import Pgmq.Migration qualified as Migration
import System.Directory (createDirectoryIfMissing, doesFileExist)
import System.Environment (lookupEnv)
import System.Posix.User (getEffectiveUserID)

-- | Root directory for ephemeral PostgreSQL clusters.
--
-- ephemeral-pg reaps abandoned clusters at startup, but only within its own
-- temporary root. With 'temporaryRoot' unset that root is @$TMPDIR@, which
-- @nix develop@ makes unique per shell, so a run never reclaims what an earlier
-- session abandoned. Pinning one root across sessions keeps them reachable.
--
-- The path is keyed by effective uid. It is a fixed location created @0700@, so
-- a run under a different user -- the Nix build sandbox -- must not collide with
-- a directory it cannot write to.
ephemeralRoot :: IO FilePath
ephemeralRoot = do
  uid <- getEffectiveUserID
  pure ("/tmp/ephpg-pgmq-hs-" <> show uid)

-- | Cached-startup configuration pinned to 'ephemeralRoot'.
ephemeralConfig :: IO Config
ephemeralConfig = do
  root <- ephemeralRoot
  createDirectoryIfMissing True root
  pure defaultConfig {temporaryRoot = Last (Just root)}

-- | Run an action with a temporary PostgreSQL database that has pgmq schema installed
withPgmqDb :: (Pool.Pool -> IO a) -> IO (Either StartError a)
withPgmqDb action =
  ephemeralConfig >>= \config ->
    withCachedConfig config defaultCacheConfig $ \db -> do
      let connSettings = connectionSettings db
          poolConfig =
            PoolConfig.settings
              [ PoolConfig.size 3,
                PoolConfig.staticConnectionSettings connSettings
              ]
      pool <- Pool.acquire poolConfig
      version <- lookupEnv "PGMQ_TEST_SCHEMA_VERSION"
      case version of
        Just "1.12.0" -> do
          paths <- filterM doesFileExist ["test/fixtures/pgmq-1.12.0.sql", "pgmq-config/test/fixtures/pgmq-1.12.0.sql"]
          path <- case paths of
            candidate : _ -> pure candidate
            [] -> error "Missing packaged PGMQ 1.12.0 test fixture"
          sql <- TextIO.readFile path
          Pool.use pool (Session.script sql) >>= either (error . show) pure
        Nothing -> installNative connSettings
        Just "1.13.0" -> installNative connSettings
        Just invalid -> error ("Invalid PGMQ_TEST_SCHEMA_VERSION: " <> invalid)
      required <- (== Just "1") <$> lookupEnv "PGMQ_REQUIRE_PARTMAN"
      -- Required runs fail for absent or unusable pg_partman, not just missing metadata.
      partman <- Pool.use pool (Session.script "CREATE SCHEMA IF NOT EXISTS partman; CREATE EXTENSION IF NOT EXISTS pg_partman SCHEMA partman")
      case partman of
        Left err -> when required (error ("PGMQ_REQUIRE_PARTMAN=1: " <> show err))
        Right () -> pure ()
      action pool
  where
    installNative connSettings = do
      component <- either (error . ("Invalid PGMQ migration component: " <>) . show) pure Migration.pgmqMigrations
      plan <- either (error . ("Invalid PGMQ migration plan: " <>) . show) pure (migrationPlan (component :| []))
      installResult <- runMigrationPlan defaultRunOptions connSettings plan
      either (error . ("Migration failed: " <>) . show) (const (pure ())) installResult