{-# LANGUAGE OverloadedStrings #-}

-- | Schema setup, teardown, and connection helpers for the arbiter test suites.
module Arbiter.Test.Setup
  ( SetupConfig (..)
  , defaultSetupConfig
  , setupDDL
  , setupDDLWithNotify
  , cleanupData
  , execute_
  , execStatement
  , execQuery
  , setupOnce
  , addQueueTable
  , disableNoticeReporting
  , createSharedPool
  , truncateToMicros
  , seedConcurrencyPoolSQL
  , drainWith
  ) where

import Arbiter.Core.Codec (RowCodec)
import Arbiter.Core.Concurrency.Schema qualified as CC
import Arbiter.Core.Concurrency.Spec (ConcurrencyPolicy (..))
import Arbiter.Core.Job.Schema qualified as Schema
import Arbiter.Core.MonadArbiter (MonadArbiter, Params)
import Arbiter.Core.MonadArbiter qualified as MA
import Arbiter.Core.RateLimit.Schema qualified as RL
import Arbiter.Core.SchemaTables (allSchemaTables)
import Arbiter.Core.SqlLiterals (textLiteral)
import Arbiter.Migrations
  ( allTableAdmission
  , jobQueueMigrationsForTable
  , schemaLevelMigrations
  )
import Control.Concurrent (threadDelay)
import Control.Exception (throwIO, try)
import Control.Monad (void, when)
import Data.ByteString (ByteString)
import Data.Foldable (traverse_)
import Data.Int (Int32, Int64)
import Data.Pool (Pool, defaultPoolConfig, newPool, setNumStripes)
import Data.String (fromString)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (UTCTime (..), picosecondsToDiffTime)
import Database.PostgreSQL.LibPQ qualified as LibPQ
import Database.PostgreSQL.Simple (Connection, SqlError (..), close, connectPostgreSQL, execute)
import Database.PostgreSQL.Simple.Internal qualified as PGS
import Database.PostgreSQL.Simple.Migration (MigrationCommand (..))
import Database.PostgreSQL.Simple.Types (Query (..))

-- | Configuration for test setup
data SetupConfig = SetupConfig
  { SetupConfig -> Bool
setupEnableNotifications :: Bool
  -- ^ Whether to create LISTEN/NOTIFY triggers
  , SetupConfig -> Bool
setupEnableRankingIndexes :: Bool
  -- ^ Whether to create ranking indexes for optimized claim queries
  }
  deriving stock (SetupConfig -> SetupConfig -> Bool
(SetupConfig -> SetupConfig -> Bool)
-> (SetupConfig -> SetupConfig -> Bool) -> Eq SetupConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SetupConfig -> SetupConfig -> Bool
== :: SetupConfig -> SetupConfig -> Bool
$c/= :: SetupConfig -> SetupConfig -> Bool
/= :: SetupConfig -> SetupConfig -> Bool
Eq, Int -> SetupConfig -> ShowS
[SetupConfig] -> ShowS
SetupConfig -> String
(Int -> SetupConfig -> ShowS)
-> (SetupConfig -> String)
-> ([SetupConfig] -> ShowS)
-> Show SetupConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SetupConfig -> ShowS
showsPrec :: Int -> SetupConfig -> ShowS
$cshow :: SetupConfig -> String
show :: SetupConfig -> String
$cshowList :: [SetupConfig] -> ShowS
showList :: [SetupConfig] -> ShowS
Show)

-- | Default test setup configuration
defaultSetupConfig :: SetupConfig
defaultSetupConfig :: SetupConfig
defaultSetupConfig =
  SetupConfig
    { setupEnableNotifications :: Bool
setupEnableNotifications = Bool
True
    , setupEnableRankingIndexes :: Bool
setupEnableRankingIndexes = Bool
True
    }

-- | Rebuild the schema without notification triggers.
setupDDL :: Text -> Text -> Connection -> IO ()
setupDDL :: Text -> Text -> Connection -> IO ()
setupDDL = SetupConfig -> Text -> Text -> Connection -> IO ()
setupDDLWithConfig SetupConfig
defaultSetupConfig {setupEnableNotifications = False}

-- | Rebuild the schema with notification triggers.
setupDDLWithNotify :: Text -> Text -> Connection -> IO ()
setupDDLWithNotify :: Text -> Text -> Connection -> IO ()
setupDDLWithNotify = SetupConfig -> Text -> Text -> Connection -> IO ()
setupDDLWithConfig SetupConfig
defaultSetupConfig

-- | Rebuild the schema with the shipped migration scripts.
setupDDLWithConfig :: SetupConfig -> Text -> Text -> Connection -> IO ()
setupDDLWithConfig :: SetupConfig -> Text -> Text -> Connection -> IO ()
setupDDLWithConfig SetupConfig
config Text
schemaName Text
tableName Connection
conn = do
  IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Text -> IO ()
execute_ Connection
conn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"DROP SCHEMA IF EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" CASCADE"
  IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Text -> IO ()
execute_ Connection
conn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> Text
Schema.createSchemaSQL Text
schemaName
  (MigrationCommand -> IO ()) -> [MigrationCommand] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ MigrationCommand -> IO ()
runScript ([MigrationCommand] -> IO ()) -> [MigrationCommand] -> IO ()
forall a b. (a -> b) -> a -> b
$
    Text -> [MigrationCommand]
schemaLevelMigrations Text
schemaName
      [MigrationCommand] -> [MigrationCommand] -> [MigrationCommand]
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> TableAdmission -> [MigrationCommand]
jobQueueMigrationsForTable Text
schemaName Text
tableName TableAdmission
allTableAdmission
  Bool -> Text -> Text -> Connection -> IO ()
createNotifyObjects (SetupConfig -> Bool
setupEnableNotifications SetupConfig
config) Text
schemaName Text
tableName Connection
conn
  where
    skipped :: [String]
skipped
      | SetupConfig -> Bool
setupEnableRankingIndexes SetupConfig
config = []
      | Bool
otherwise = ShowS -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ((Text -> String
T.unpack Text
tableName String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"-") String -> ShowS
forall a. Semigroup a => a -> a -> a
<>) [String
"create-group-key-index", String
"migrate-ungrouped-ready-split-indexes"]
    runScript :: MigrationCommand -> IO ()
runScript (MigrationScript String
name ByteString
sql)
      | String
name String -> [String] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
skipped = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> () -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
execute Connection
conn (ByteString -> Query
Query ByteString
sql) ()
    runScript MigrationCommand
_ = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Truncate the queue's tables between tests.
cleanupData :: Text -> Text -> Connection -> IO ()
cleanupData :: Text -> Text -> Connection -> IO ()
cleanupData Text
schemaName Text
tableName Connection
conn = do
  Connection -> Text -> IO ()
execute_ Connection
conn Text
"SET client_min_messages = WARNING"
  -- A draining worker pool can still hold locks on the job and groups tables.
  -- Bound the wait and retry on a deadlock or lock timeout.
  Connection -> Text -> IO ()
execute_ Connection
conn Text
"SET lock_timeout = '5s'"
  -- Rate-limit policies are seeded once per suite.
  let truncated :: [Text]
truncated = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
RL.arbiterRateLimitPoliciesTableName) ([Text] -> [Text]
allSchemaTables [Text
tableName])
      truncateSql :: Text
truncateSql =
        Text
"TRUNCATE "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Text -> Text
Schema.qualifiedTable Text
schemaName) [Text]
truncated)
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" CASCADE"
      go :: Int -> IO ()
go Int
remaining = do
        outcome <- IO () -> IO (Either SqlError ())
forall e a. Exception e => IO a -> IO (Either e a)
try (Connection -> Text -> IO ()
execute_ Connection
conn Text
truncateSql) :: IO (Either SqlError ())
        case outcome of
          Right () -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          -- 40P01 deadlock_detected, 55P03 lock_not_available
          Left SqlError
sqlErr
            | SqlError -> ByteString
sqlState SqlError
sqlErr ByteString -> [ByteString] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [ByteString
"40P01", ByteString
"55P03"] Bool -> Bool -> Bool
&& Int
remaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> (Int
0 :: Int) ->
                Int -> IO ()
threadDelay Int
100_000 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO ()
go (Int
remaining Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            | Bool
otherwise -> SqlError -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SqlError
sqlErr
  Int -> IO ()
go Int
10
  Connection -> Text -> IO ()
execute_ Connection
conn Text
"RESET lock_timeout"
  Connection -> Text -> IO ()
execute_ Connection
conn Text
"SET client_min_messages = NOTICE"

-- | Run a statement with no parameters (test setup only).
execute_ :: Connection -> Text -> IO ()
execute_ :: Connection -> Text -> IO ()
execute_ Connection
conn Text
sql = IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> () -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
execute Connection
conn (String -> Query
forall a. IsString a => String -> a
fromString (Text -> String
T.unpack Text
sql) :: Query) ()

-- | Run a pre-rendered statement with positional parameters (test setup only).
execStatement :: (MonadArbiter m) => Text -> Params -> m Int64
execStatement :: forall (m :: * -> *). MonadArbiter m => Text -> Params -> m Int64
execStatement Text
sql Params
params = Query () -> m Int64
forall a. Query a -> m Int64
forall (m :: * -> *) a. MonadArbiter m => Query a -> m Int64
MA.executeStatement (Text -> Params -> RowCodec () -> Query ()
forall a. Text -> Params -> RowCodec a -> Query a
MA.Query Text
sql Params
params (() -> RowCodec ()
forall a. a -> Ap NullCol a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))

-- | Run a pre-rendered query with positional parameters and a decoder (test setup only).
execQuery :: (MonadArbiter m) => Text -> Params -> RowCodec a -> m [a]
execQuery :: forall (m :: * -> *) a.
MonadArbiter m =>
Text -> Params -> RowCodec a -> m [a]
execQuery Text
sql Params
params RowCodec a
codec = Query a -> m [a]
forall a. Query a -> m [a]
forall (m :: * -> *) a. MonadArbiter m => Query a -> m [a]
MA.executeQuery (Text -> Params -> RowCodec a -> Query a
forall a. Text -> Params -> RowCodec a -> Query a
MA.Query Text
sql Params
params RowCodec a
codec)

-- | Connect and rebuild the schema once, for a suite's outer bracket.
setupOnce :: ByteString -> Text -> Text -> Bool -> IO ()
setupOnce :: ByteString -> Text -> Text -> Bool -> IO ()
setupOnce ByteString
connStr Text
schemaName Text
tableName Bool
withNotify = do
  conn <- ByteString -> IO Connection
connectPostgreSQL ByteString
connStr
  disableNoticeReporting conn
  let config = SetupConfig
defaultSetupConfig {setupEnableNotifications = withNotify}
  setupDDLWithConfig config schemaName tableName conn
  close conn

-- | Add one more job-queue table to an existing schema.
addQueueTable :: ByteString -> Text -> Text -> Bool -> IO ()
addQueueTable :: ByteString -> Text -> Text -> Bool -> IO ()
addQueueTable ByteString
connStr Text
schemaName Text
tableName Bool
withNotify = do
  conn <- ByteString -> IO Connection
connectPostgreSQL ByteString
connStr
  disableNoticeReporting conn
  traverse_ (runScript conn) (jobQueueMigrationsForTable schemaName tableName allTableAdmission)
  createNotifyObjects withNotify schemaName tableName conn
  close conn
  where
    runScript :: Connection -> MigrationCommand -> IO ()
runScript Connection
conn (MigrationScript String
_ ByteString
sql) = IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> () -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
execute Connection
conn (ByteString -> Query
Query ByteString
sql) ()
    runScript Connection
_ MigrationCommand
_ = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Install the notification function and trigger the reconciler would.
createNotifyObjects :: Bool -> Text -> Text -> Connection -> IO ()
createNotifyObjects :: Bool -> Text -> Text -> Connection -> IO ()
createNotifyObjects Bool
withNotify Text
schemaName Text
tableName Connection
conn =
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
withNotify (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    Connection -> Text -> IO ()
execute_ Connection
conn (Text -> Text -> Text
Schema.createNotifyFunctionSQL Text
schemaName Text
tableName)
    Connection -> Text -> IO ()
execute_ Connection
conn (Text -> Text -> Text
Schema.createNotifyTriggerSQL Text
schemaName Text
tableName)

-- | Silence libpq notice output for the connection.
disableNoticeReporting :: Connection -> IO ()
disableNoticeReporting :: Connection -> IO ()
disableNoticeReporting Connection
conn =
  Connection -> (Connection -> IO ()) -> IO ()
forall a. Connection -> (Connection -> IO a) -> IO a
PGS.withConnection Connection
conn Connection -> IO ()
LibPQ.disableNoticeReporting

-- | A single-stripe pool of 5 connections for tests.
createSharedPool :: ByteString -> IO (Pool Connection)
createSharedPool :: ByteString -> IO (Pool Connection)
createSharedPool ByteString
connStr =
  PoolConfig Connection -> IO (Pool Connection)
forall a. PoolConfig a -> IO (Pool a)
newPool (PoolConfig Connection -> IO (Pool Connection))
-> PoolConfig Connection -> IO (Pool Connection)
forall a b. (a -> b) -> a -> b
$ Maybe Int -> PoolConfig Connection -> PoolConfig Connection
forall a. Maybe Int -> PoolConfig a -> PoolConfig a
setNumStripes (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
1) (PoolConfig Connection -> PoolConfig Connection)
-> PoolConfig Connection -> PoolConfig Connection
forall a b. (a -> b) -> a -> b
$ IO Connection
-> (Connection -> IO ()) -> Double -> Int -> PoolConfig Connection
forall a. IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
defaultPoolConfig (ByteString -> IO Connection
connectPostgreSQL ByteString
connStr) Connection -> IO ()
close Double
60 Int
5

-- | Seed a concurrency pool's default limit and clear any override, as SQL statements.
seedConcurrencyPoolSQL :: Text -> Text -> Int32 -> [Text]
seedConcurrencyPoolSQL :: Text -> Text -> Int32 -> [Text]
seedConcurrencyPoolSQL Text
schema Text
prefix Int32
lim =
  [ Text -> ConcurrencyPolicy -> Text
CC.upsertConcurrencyPolicyRowSQL Text
schema (Text -> Int32 -> ConcurrencyPolicy
ConcurrencyPolicy Text
prefix Int32
lim)
  , Text
"UPDATE "
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
CC.arbiterConcurrencyPoliciesTable Text
schema
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" SET override_limit = NULL WHERE prefix_id = "
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
textLiteral Text
prefix
  ]

-- | Repeat a batch action until it returns empty, accumulating the results.
drainWith :: IO [a] -> IO [a]
drainWith :: forall a. IO [a] -> IO [a]
drainWith IO [a]
fetch = [[a]] -> IO [a]
go []
  where
    go :: [[a]] -> IO [a]
go [[a]]
batches = do
      batch <- IO [a]
fetch
      if null batch then pure (concat (reverse batches)) else go (batch : batches)

-- | Truncate to microsecond precision to match PostgreSQL @timestamptz@.
truncateToMicros :: UTCTime -> UTCTime
truncateToMicros :: UTCTime -> UTCTime
truncateToMicros (UTCTime Day
day DiffTime
dayTime) =
  let micros :: Integer
micros = DiffTime -> Integer
forall b. Integral b => DiffTime -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (DiffTime
dayTime DiffTime -> DiffTime -> DiffTime
forall a. Num a => a -> a -> a
* DiffTime
1e6) :: Integer
   in Day -> DiffTime -> UTCTime
UTCTime Day
day (Integer -> DiffTime
picosecondsToDiffTime (Integer
micros Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
1000000))