pgmq-effectful-0.7.0.0: test/EphemeralDb.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Test database infrastructure using ephemeral-pg.
--
-- This is a copy of pgmq-hasql\'s EphemeralDb helper, adjusted for
-- pgmq-effectful\'s test suite. Keeping a local copy avoids cross-package
-- test-helper sharing (which would require a shared test-helper library
-- or fragile cross-package hs-source-dirs).
module EphemeralDb
  ( withPgmqDb,
    withPgmqPool,
    ephemeralConfig,
    ensureOwnerOnlyRoot,
    installPgmqNative,
    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.Connection.Settings qualified as ConnectionSettings
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 the
-- 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-effectful/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 -> installPgmqNative connSettings
        Just "1.13.0" -> installPgmqNative 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

-- | Apply the native pgmq migration ledger to the database at the given
-- connection settings, failing loudly on any error.
installPgmqNative :: ConnectionSettings.Settings -> IO ()
installPgmqNative 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

-- | Alias for 'withPgmqDb' kept for parallelism with pgmq-hasql.
withPgmqPool :: (Pool.Pool -> IO a) -> IO (Either StartError a)
withPgmqPool = withPgmqDb