Postgres.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedRecordDot #-}

-- | Suite-level PostgreSQL test fixture for Keiro test suites.
--
-- This follows the @ephemeral-pg@ "suite-level template databases" best practice:
-- start one cached PostgreSQL server for the whole suite, migrate a template
-- database once, and clone a fresh, isolated database per example with
-- PostgreSQL's @CREATE DATABASE ... TEMPLATE ...@. Each example still receives an
-- empty, isolated database, but the expensive work — server startup and running
-- the Kiroku and Keiro migrations — happens once per suite rather than once per
-- example.
--
-- Usage from @hspec@:
--
-- @
-- main :: IO ()
-- main =
--   'withMigratedSuiteWith' [pgmqMigrations] \\fixture ->
--     hspec $
--       describe "..." $ around ('withFreshStore' fixture) $ do
--         it "..." $ \\store -> ...
-- @
module Keiro.Test.Postgres
  ( Fixture,
    withMigratedSuite,
    withMigratedSuiteWith,
    withMigratedBenchmarkSuite,
    withFreshDatabase,
    withFreshStore,
    withFreshStoreWith,
    StoreRunner (..),
    withFreshResourceStore,
    withFreshResourceStoreWith,
    withFreshResourceStorePrepared,
    withFreshStores2,
    withFreshResourceStoresSameDatabase,
    withFreshResourceStoresSameDatabaseWithConnection,
  )
where

import Control.Concurrent.STM (TVar, atomically, newTVarIO, stateTVar)
import Control.Exception (bracket, onException)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Monoid (Last (..))
import Data.Text (Text)
import Data.Text qualified as Text
import Database.PostgreSQL.Migrate (MigrationComponent, defaultRunOptions, migrationPlan, runMigrationPlan)
import Effectful (Eff, IOE, Limit (..), Persistence (..), UnliftStrategy (..), runEff, withEffToIO)
import Effectful.Error.Static (Error, runErrorNoCallStack)
import EphemeralPg qualified as Pg
import GHC.Clock (getMonotonicTimeNSec)
import Hasql.Connection.Settings qualified as Conn
import Hasql.Pool qualified as Pool
import Hasql.Pool.Config qualified as Pool.Config
import Hasql.Session qualified as Session
import Keiro.Migrations (keiroMigrations)
import Kiroku.Store qualified as Store
import Kiroku.Store.Effect (Store, runStoreResource)
import Kiroku.Store.Effect.Resource (KirokuStoreResource, getKirokuStore, withKirokuStore)
import Kiroku.Store.Error (StoreError)
import Kiroku.Store.Migrations qualified as Kiroku
import System.Directory (createDirectoryIfMissing)
import System.Environment (lookupEnv)
import System.Posix.User (getEffectiveUserID)

-- | A running, migrated suite fixture: one cached PostgreSQL server owning a
-- single migrated template database, plus a counter for unique clone names.
data Fixture = Fixture
  { administrativeConnection :: Text,
    databaseConnection :: Text -> Text,
    templateName :: Text,
    nextId :: TVar Int
  }

templateDbName :: Text
templateDbName = "keiro_template"

-- | The ephemeral PostgreSQL configuration for Keiro suites: 'Pg.defaultConfig'
-- with a stable, per-user temporary root.
--
-- ephemeral-pg sweeps clusters abandoned by killed runs on the next startup, but
-- only inside the configured temporary root. Left unset, the root is @$TMPDIR@,
-- which @nix develop@ and some CI runners make per-session, so each run sweeps an
-- empty directory and orphaned postmasters accumulate. A fixed root shared by
-- every Keiro suite lets any run reclaim what an earlier run left behind. The
-- effective uid keys the root so a build sandbox running as another user does not
-- collide with a developer's @0700@ directory. The keiro-migrations test suites
-- use the same root.
ephemeralPgConfig :: IO Pg.Config
ephemeralPgConfig = do
  uid <- getEffectiveUserID
  let root = "/tmp/ephpg-keiro-" <> show uid
  createDirectoryIfMissing True root
  pure Pg.defaultConfig {Pg.temporaryRoot = Last (Just root)}

-- | Start one cached PostgreSQL server, create a template database, apply the
-- Kiroku event-store schema and Keiro framework schema to it once, then run
-- @action@ with the resulting 'Fixture'. The server is stopped on exit.
--
-- 'EphemeralPg.startCached' restores a clean @initdb@ cluster (no schema), and
-- 'Kiroku.Store.withStore' deliberately creates no tables, so the migrations
-- are applied here against the template before any example clones it.
withMigratedSuite :: (Fixture -> IO a) -> IO a
withMigratedSuite = withMigratedSuiteWith []

-- | Like 'withMigratedSuite', but appends extra @pg-migrate@ components to the
-- framework plan so the template database also carries their schema — for example
-- @pgmq-migration@'s @pgmqMigrations@. The extra components are applied after
-- Kiroku's and Keiro's, in the order given, and before any example database is
-- cloned.
--
-- They join the framework plan rather than running as a separate plan because a
-- @pg-migrate@ ledger is shared by every component in it: a plan that omits a
-- component already recorded in the ledger fails strict verification with
-- @UnknownStoredMigration@.
withMigratedSuiteWith :: [MigrationComponent] -> (Fixture -> IO a) -> IO a
withMigratedSuiteWith extraComponents action = do
  config <- ephemeralPgConfig
  withConfiguredSuite config extraComponents action

-- | Durable benchmark databases. An exclusively leased remote server may be
-- supplied through KEIRO_GUARDED_CATALOG_PG_ADMIN, a libpq keyword connection
-- string. Only uniquely named databases created here are dropped. Ordinary
-- correctness fixtures retain their fast local test configuration.
withMigratedBenchmarkSuite :: (Fixture -> IO a) -> IO a
withMigratedBenchmarkSuite action =
  lookupEnv "KEIRO_GUARDED_CATALOG_PG_ADMIN" >>= \case
    Nothing -> do
      config <- ephemeralPgConfig
      let durable = config {Pg.postgresSettings = filter (\(key, _) -> key `notElem` ["fsync", "synchronous_commit", "full_page_writes", "wal_level"]) config.postgresSettings <> [("fsync", "on"), ("synchronous_commit", "on"), ("full_page_writes", "on"), ("wal_level", "replica")]}
      withConfiguredSuite durable [] action
    Just prefix -> do
      let connection = Text.pack prefix
      if Text.null connection || "postgres://" `Text.isPrefixOf` connection || "postgresql://" `Text.isPrefixOf` connection
        then fail "KEIRO_GUARDED_CATALOG_PG_ADMIN requires a libpq keyword connection string"
        else do
          unique <- getMonotonicTimeNSec
          let template = "keiro_bench_" <> Text.pack (show unique)
              connect name = connection <> " dbname=" <> name
          bracket
            ( do
                runSqlOn connection ("CREATE DATABASE " <> quoteIdentifier template)
                migrateTemplate [] (connect template)
                  `onException` runSqlOn connection ("DROP DATABASE " <> quoteIdentifier template <> " WITH (FORCE)")
                counter <- newTVarIO 0
                pure (Fixture connection connect template counter)
            )
            (\_ -> runSqlOn connection ("DROP DATABASE " <> quoteIdentifier template <> " WITH (FORCE)"))
            action

withConfiguredSuite :: Pg.Config -> [MigrationComponent] -> (Fixture -> IO a) -> IO a
withConfiguredSuite config extraComponents action = do
  started <- Pg.startCached config Pg.defaultCacheConfig
  case started of
    Left err -> fail (Text.unpack (Pg.renderStartError err))
    Right server ->
      bracket
        (setup server `onException` Pg.stop server)
        (const (Pg.stop server))
        action
  where
    setup server = do
      counter <- newTVarIO 0
      runSql server ("CREATE DATABASE " <> quoteIdentifier templateDbName)
      let templateConnStr = connectionStringFor server templateDbName
      -- Apply migrations through short-lived pg-migrate connections, all of which are
      -- released before any clone, so the template has no active sessions when
      -- PostgreSQL copies it.
      migrateTemplate extraComponents templateConnStr
      pure (Fixture (Pg.connectionString server) (connectionStringFor server) templateDbName counter)

-- | Clone a fresh, empty, migrated database from the template, open a
-- 'Store.KirokuStore' against it, run @action@, then drop the clone. The store
-- (including its notifier and publisher connections) is torn down before the
-- database is dropped.
withFreshStore :: Fixture -> (Store.KirokuStore -> IO ()) -> IO ()
withFreshStore fixture = withFreshStoreWith fixture id

-- | Like 'withFreshStore' but applies @modify@ to the default connection
-- settings before opening the store — for example to add an application projection
-- schema to @extraSearchPath@ so a read-model table in that schema resolves on the
-- store pool (see 'Keiro.Connection.withProjectionSchema'). The store @schema@
-- field itself is left at kiroku's default.
withFreshStoreWith ::
  Fixture ->
  (Store.ConnectionSettings -> Store.ConnectionSettings) ->
  (Store.KirokuStore -> IO ()) ->
  IO ()
withFreshStoreWith fixture modify action =
  withFreshDatabase fixture \connStr ->
    Store.withStore (modify (Store.defaultConnectionSettings connStr)) action

-- | A reusable interpreter for computations that need both the dynamic
-- 'Store' effect and the resource effect carrying the same live store handle.
newtype StoreRunner
  = StoreRunner
      (forall a. Eff '[Store, Error StoreError, KirokuStoreResource, IOE] a -> IO (Either StoreError a))

-- | Like 'withFreshStore', but also supplies the resource-aware interpreter
-- required by Keiro's transactional command runners.
withFreshResourceStore :: Fixture -> ((Store.KirokuStore, StoreRunner) -> IO a) -> IO a
withFreshResourceStore fixture = withFreshResourceStoreWith fixture id

-- | Like 'withFreshResourceStore', with a connection-settings modifier for
-- projection schemas or store hooks.
withFreshResourceStoreWith ::
  Fixture ->
  (Store.ConnectionSettings -> Store.ConnectionSettings) ->
  ((Store.KirokuStore, StoreRunner) -> IO a) ->
  IO a
withFreshResourceStoreWith fixture modify action =
  withFreshDatabase fixture \connStr ->
    runResourceStore (modify (Store.defaultConnectionSettings connStr)) action

-- | Clone a database, let a privileged store prepare database roles or ACLs,
-- then open the resource-aware store with modified application settings.
-- The preparation store is closed before the application store is opened.
withFreshResourceStorePrepared ::
  Fixture ->
  (Store.KirokuStore -> IO ()) ->
  (Store.ConnectionSettings -> Store.ConnectionSettings) ->
  ((Store.KirokuStore, StoreRunner) -> IO ()) ->
  IO ()
withFreshResourceStorePrepared fixture prepare modify action =
  withFreshDatabase fixture \connStr -> do
    Store.withStore (Store.defaultConnectionSettings connStr) prepare
    runResourceStore (modify (Store.defaultConnectionSettings connStr)) action

runResourceStore ::
  Store.ConnectionSettings ->
  ((Store.KirokuStore, StoreRunner) -> IO a) ->
  IO a
runResourceStore settings action =
  runEff $
    withKirokuStore settings $ do
      store <- getKirokuStore
      withEffToIO (ConcUnlift Persistent Unlimited) \unlift ->
        action
          ( store,
            StoreRunner (unlift . runErrorNoCallStack . runStoreResource)
          )

-- | Independent resource-aware pools sharing one migrated database. Unlike
-- 'withFreshStores2', this fixture exercises contention in the same event store.
withFreshResourceStoresSameDatabase ::
  Fixture ->
  (((Store.KirokuStore, StoreRunner), (Store.KirokuStore, StoreRunner)) -> IO ()) ->
  IO ()
withFreshResourceStoresSameDatabase fixture action =
  withFreshResourceStoresSameDatabaseWithConnection fixture (\_ -> action)

-- | Also expose the clone's connection string for an independent SQL observer
-- or controller. All supplied stores are closed before the clone is dropped.
withFreshResourceStoresSameDatabaseWithConnection ::
  Fixture ->
  (Text -> ((Store.KirokuStore, StoreRunner), (Store.KirokuStore, StoreRunner)) -> IO ()) ->
  IO ()
withFreshResourceStoresSameDatabaseWithConnection fixture action =
  withFreshDatabase fixture \connStr ->
    runResourceStore (Store.defaultConnectionSettings connStr) \first ->
      runResourceStore (Store.defaultConnectionSettings connStr) \second ->
        action connStr (first, second)

-- | Like 'withFreshStore' but provides two independent migrated databases (and
-- two stores) cloned from the same template — used by cross-context
-- integration tests that need two isolated PostgreSQL databases.
withFreshStores2 :: Fixture -> ((Store.KirokuStore, Store.KirokuStore) -> IO ()) -> IO ()
withFreshStores2 fixture action =
  withFreshDatabase fixture \connStrA ->
    withFreshDatabase fixture \connStrB ->
      Store.withStore (Store.defaultConnectionSettings connStrA) \storeA ->
        Store.withStore (Store.defaultConnectionSettings connStrB) \storeB ->
          action (storeA, storeB)

-- | Clone a fresh database from the template, pass its connection string to
-- @action@, and drop it afterwards. Database names are unique per clone.
withFreshDatabase :: Fixture -> (Text -> IO a) -> IO a
withFreshDatabase fixture action =
  bracket create dropDb \dbName ->
    action (fixture.databaseConnection dbName)
  where
    create = do
      n <- atomically $ stateTVar fixture.nextId \i -> (i + 1, i + 1)
      let dbName = fixture.templateName <> "_" <> Text.pack (show n)
      runSqlOn fixture.administrativeConnection $
        "CREATE DATABASE "
          <> quoteIdentifier dbName
          <> " TEMPLATE "
          <> quoteIdentifier fixture.templateName
      pure dbName

    dropDb dbName =
      runSqlOn fixture.administrativeConnection $
        "DROP DATABASE IF EXISTS " <> quoteIdentifier dbName <> " WITH (FORCE)"

migrateTemplate :: [MigrationComponent] -> Text -> IO ()
migrateTemplate extraComponents connStr = do
  kiroku <- either (fail . show) pure Kiroku.kirokuMigrations
  keiro <- either (fail . show) pure keiroMigrations
  plan <- either (fail . show) pure (migrationPlan (kiroku :| keiro : extraComponents))
  result <-
    runMigrationPlan
      defaultRunOptions
      (Conn.connectionString connStr)
      plan
  either (fail . show) (const (pure ())) result

-- | Build a libpq connection string for a named database on the fixture's
-- server, addressing it over the server's Unix socket.
connectionStringFor :: Pg.Database -> Text -> Text
connectionStringFor db dbName =
  Text.unwords
    [ "host=" <> Text.pack db.socketDirectory,
      "port=" <> Text.pack (show db.port),
      "dbname=" <> dbName,
      "user=" <> db.user
    ]

-- | Run a single SQL command against the server's default database (used for
-- @CREATE DATABASE@ / @DROP DATABASE@, which cannot run inside a transaction).
runSql :: Pg.Database -> Text -> IO ()
runSql db = runSqlOn (Pg.connectionString db)

runSqlOn :: Text -> Text -> IO ()
runSqlOn connStr sql =
  bracket acquire Pool.release \pool ->
    Pool.use pool (Session.script sql) >>= either (fail . show) pure
  where
    acquire =
      Pool.acquire $
        Pool.Config.settings
          [ Pool.Config.staticConnectionSettings (Conn.connectionString connStr),
            Pool.Config.size 1
          ]

quoteIdentifier :: Text -> Text
quoteIdentifier ident =
  "\"" <> Text.replace "\"" "\"\"" ident <> "\""