pgmq-config-0.7.0.0: test/EphemeralDb.hs

{-# LANGUAGE OverloadedStrings #-}

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

    -- * Re-exports
    StartError,
  )
where

import Control.Monad (filterM, unless, when)
import Data.Bits ((.&.))
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 (doesFileExist)
import System.Environment (lookupEnv)
import System.IO.Error (catchIOError, isAlreadyExistsError)
import System.Posix.Directory qualified as Posix
import System.Posix.Files (deviceID, fileID, fileMode, fileOwner, getSymbolicLinkStatus, isDirectory, setFileMode)
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)

-- | Establish a real effective-uid-owned root without deleting its contents.
ensureOwnerOnlyRoot :: FilePath -> IO ()
ensureOwnerOnlyRoot root = do
  catchIOError (Posix.createDirectory root 0o700) $ \err ->
    unless (isAlreadyExistsError err) (ioError err)
  uid <- getEffectiveUserID
  before <- getSymbolicLinkStatus root
  unless (isDirectory before && fileOwner before == uid) $
    ioError (userError ("Unsafe ephemeral PostgreSQL root: " <> root))
  setFileMode root 0o700
  after <- getSymbolicLinkStatus root
  unless
    ( isDirectory after
        && fileOwner after == uid
        && deviceID before == deviceID after
        && fileID before == fileID after
        && fileMode after .&. 0o7777 == 0o700
    )
    (ioError (userError ("Ephemeral PostgreSQL root changed or is not owner-only: " <> root)))

-- | Cached-startup configuration pinned to 'ephemeralRoot'.
ephemeralConfig :: IO Config
ephemeralConfig = do
  root <- ephemeralRoot
  ensureOwnerOnlyRoot 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