{-# LANGUAGE OverloadedStrings #-}

-- | Versioned, tracked migrations for job queue schemas, run once in order.
-- History lives in a @schema_migrations@ table inside the target schema.
-- Each run also seeds the admission policies and reconciles the rate-limit durability
-- and optional triggers in 'MigrationConfig'.
module Arbiter.Migrations
  ( -- * Configuration
    MigrationConfig (..)
  , defaultMigrationConfig

    -- * Tracked migrations
  , runMigrationsForRegistry
  , runMigrationsTrackedForTables
  , jobQueueMigrationsForTable
  , schemaLevelMigrations
  , AdmissionSeeds (..)
  , noAdmissionSeeds
  , TableAdmission (..)
  , allTableAdmission

    -- * Validation
  , validateRegistryNames
  , maxQueueNameBytes
  , conflictingPolicyPrefixes

    -- * Re-exports
  , QueueSpec (..)
  , Queue
  , MigrationCommand (..)
  , MigrationResult (..)
  , Durability (..)
  , RateLimitPolicy (..)
  , ConcurrencyPolicy (..)
  , SchemaName
  , TableName
  ) where

import Arbiter.Core.Concurrency.Schema
  ( addConcurrencyColumnsSQL
  , createConcurrencyIndexSQL
  , createConcurrencyPoliciesTableSQL
  , createConcurrencyTableSQL
  , createConcurrencyTriggerFunctionsSQL
  , createConcurrencyTriggersSQL
  , upsertConcurrencyPolicyRowSQL
  )
import Arbiter.Core.Concurrency.Spec (ConcurrencyPolicy (..), registryConcurrencyPolicies, registryConcurrencyTables)
import Arbiter.Core.CronSchedule
  ( addLastManualRunColumnSQL
  , addQueueNameColumnSQL
  , addRunRequestedColumnSQL
  , addTimezoneColumnSQL
  , createCronSchedulesTableSQL
  )
import Arbiter.Core.Enum (showText)
import Arbiter.Core.Exceptions (displayEx)
import Arbiter.Core.Gates (addGateMetadataColumnSQL, createGatesTableSQL)
import Arbiter.Core.Job.Schema
  ( SchemaName
  , TableName
  , addArchiveForColumnSQL
  , addCancelRequestedAtColumnSQL
  , addClaimSeqColumnSQL
  , addClaimedByColumnSQL
  , addKindColumnSQL
  , addTraceContextColumnSQL
  , cancelNotifyChannel
  , createArchiveCompletedAtIndexSQL
  , createArchiveExpiresAtIndexSQL
  , createArchiveGroupKeyIndexSQL
  , createArchiveJobIdIndexSQL
  , createArchiveParentIdIndexSQL
  , createDLQFailedAtIndexSQL
  , createDLQGroupKeyIndexSQL
  , createDLQParentIdIndexSQL
  , createDedupKeyIndexSQL
  , createEventStreamingFunctionSQL
  , createEventStreamingTriggersSQL
  , createJobQueueArchiveTableSQL
  , createJobQueueDLQTableSQL
  , createJobQueueTableSQL
  , createNotifyFunctionSQL
  , createNotifyTriggerSQL
  , createParentIdIndexSQL
  , createResultsTableSQL
  , createSchemaSQL
  , dlqSuffix
  , dropEventStreamingFunctionSQL
  , eventStreamingAdoptedObjectComment
  , eventStreamingDLQTriggerName
  , eventStreamingFunctionName
  , eventStreamingObjectComment
  , eventStreamingObjectCommentPrefix
  , eventStreamingTriggerName
  , jobQueueTable
  , legacyEventStreamingTriggers
  , migrateUngroupedReadySplitIndexesSQL
  , notifyAdoptedObjectComment
  , notifyFunctionName
  , notifyObjectComment
  , notifyObjectCommentPrefix
  , notifyTriggerName
  , pauseNotifyChannel
  , queueTableNames
  , setMaxAttemptsDefaultSQL
  )
import Arbiter.Core.Job.Schema.Groups
  ( createGroupsEmptiedIndexSQL
  , createGroupsTableSQL
  , createGroupsTriggerFunctionsSQL
  , createGroupsTriggersSQL
  , createJobQueueGroupInFlightIndexSQL
  , createJobQueueGroupKeyIndexSQL
  , createJobQueueGroupRetriedIndexSQL
  , createJobQueueGroupedDueIndexSQL
  , migrateGroupsReadyRankingSQL
  )
import Arbiter.Core.Job.Types (RegistryAdmissionPolicies)
import Arbiter.Core.QueueRegistry (Queue, QueueSpec (..), RegistryTables (..))
import Arbiter.Core.Queues (addDLQRetentionColumnSQL, createQueuesTableSQL, registerQueuesSQL)
import Arbiter.Core.RateLimit.Schema
  ( addRateLimitColumnsSQL
  , addRateLimitCostColumnSQL
  , alterRateLimitsDurabilitySQL
  , arbiterRateLimitsTableName
  , createRateLimitBucketTriggerFunctionsSQL
  , createRateLimitBucketTriggersSQL
  , createRateLimitPoliciesTableSQL
  , createRateLimitsTableSQL
  , createThrottledIndexSQL
  , upsertRateLimitPolicyRowSQL
  )
import Arbiter.Core.RateLimit.Spec
  ( Durability (..)
  , RateLimitPolicy (..)
  , registryRateLimitPolicies
  , registryRateLimitTables
  )
import Arbiter.Core.SchemaTables (sharedArbiterTables)
import Arbiter.Core.SqlLiterals (maxIdentifierBytes)
import Arbiter.Core.Worker (createWorkersTableSQL)
import Control.Exception (bracket, catch, throwIO, try)
import Control.Monad (unless, void, when)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Foldable (find, traverse_)
import Data.Int (Int32)
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust, listToMaybe)
import Data.Proxy (Proxy (..))
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import Data.Time (NominalDiffTime)
import Database.PostgreSQL.LibPQ qualified as LibPQ
import Database.PostgreSQL.Simple (Only (..), close, connectPostgreSQL, execute_, query)
import Database.PostgreSQL.Simple qualified as PG
import Database.PostgreSQL.Simple.Internal (withConnection)
import Database.PostgreSQL.Simple.Migration
  ( MigrationCommand (..)
  , MigrationOptions (..)
  , MigrationResult (..)
  , Verbosity (..)
  , defaultOptions
  , runMigrations
  )
import Database.PostgreSQL.Simple.Types (PGArray (..), Query (..))
import UnliftIO.Exception (tryAny)

-- | Migration options and the desired state reconciled after a schema's tracked migrations run.
data MigrationConfig = MigrationConfig
  { MigrationConfig -> Bool
enableNotifications :: Bool
  -- ^ Install the LISTEN\/NOTIFY triggers that wake workers. Each run adds or drops
  -- them to match. Default: 'True'.
  , MigrationConfig -> Bool
enableEventStreaming :: Bool
  -- ^ Install the admin UI's event-streaming triggers, which send a JSON event via
  -- @pg_notify@ for every insert, update and delete on a queue table and every insert on
  -- its DLQ. A lease extend sends no event. Each run adds or drops them to match, and
  -- drops the shared function when no trigger uses it.
  -- Default: 'False'.
  , MigrationConfig -> Durability
rateLimitDurability :: Durability
  -- ^ WAL-logging for the rate-limit bucket table in this schema. 'Unlogged'
  -- (default) resets buckets on crash\/failover. 'Durable' preserves them at a
  -- throughput cost.
  , MigrationConfig -> Maybe NominalDiffTime
migrationLockTimeout :: Maybe NominalDiffTime
  -- ^ Seconds to wait for the schema's migration lock, then return 'MigrationError'.
  -- 'Nothing' (default) waits indefinitely.
  }
  deriving stock (MigrationConfig -> MigrationConfig -> Bool
(MigrationConfig -> MigrationConfig -> Bool)
-> (MigrationConfig -> MigrationConfig -> Bool)
-> Eq MigrationConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MigrationConfig -> MigrationConfig -> Bool
== :: MigrationConfig -> MigrationConfig -> Bool
$c/= :: MigrationConfig -> MigrationConfig -> Bool
/= :: MigrationConfig -> MigrationConfig -> Bool
Eq, Int -> MigrationConfig -> ShowS
[MigrationConfig] -> ShowS
MigrationConfig -> String
(Int -> MigrationConfig -> ShowS)
-> (MigrationConfig -> String)
-> ([MigrationConfig] -> ShowS)
-> Show MigrationConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MigrationConfig -> ShowS
showsPrec :: Int -> MigrationConfig -> ShowS
$cshow :: MigrationConfig -> String
show :: MigrationConfig -> String
$cshowList :: [MigrationConfig] -> ShowS
showList :: [MigrationConfig] -> ShowS
Show)

-- | Notify triggers on, event streaming off, unlogged rate-limit buckets, no lock timeout.
defaultMigrationConfig :: MigrationConfig
defaultMigrationConfig :: MigrationConfig
defaultMigrationConfig =
  MigrationConfig
    { enableNotifications :: Bool
enableNotifications = Bool
True
    , enableEventStreaming :: Bool
enableEventStreaming = Bool
False
    , rateLimitDurability :: Durability
rateLimitDurability = Durability
Unlogged
    , migrationLockTimeout :: Maybe NominalDiffTime
migrationLockTimeout = Maybe NominalDiffTime
forall a. Maybe a
Nothing
    }

-- | Migrate every queue in a registry into one schema. The schema itself is created
-- first, outside migration tracking. Names that fail 'validateRegistryNames' return
-- 'MigrationError' before the connection opens.
--
-- @
-- type AppRegistry =
--   '[ Queue "email_jobs" EmailPayload
--    , Queue "order_jobs" OrderPayload
--    ]
--
-- main :: IO ()
-- main = do
--   result <- runMigrationsForRegistry
--               (Proxy \@AppRegistry)
--               "host=localhost dbname=mydb"
--               "arbiter"
--               defaultMigrationConfig
--   print result
-- @
runMigrationsForRegistry
  :: forall registry
   . ( RegistryAdmissionPolicies registry
     , RegistryTables registry
     )
  => Proxy registry
  -- ^ Proxy for the job payload registry
  -> ByteString
  -- ^ Database connection string
  -> SchemaName
  -- ^ Schema for the arbiter tables
  -> MigrationConfig
  -- ^ Migration options
  -> IO (MigrationResult String)
runMigrationsForRegistry :: forall (registry :: JobPayloadRegistry).
(RegistryAdmissionPolicies registry, RegistryTables registry) =>
Proxy registry
-> ByteString
-> SchemaName
-> MigrationConfig
-> IO (MigrationResult String)
runMigrationsForRegistry Proxy registry
proxy ByteString
connStr SchemaName
schemaName MigrationConfig
config = do
  let ccTables :: Map SchemaName Bool
ccTables = [(SchemaName, Bool)] -> Map SchemaName Bool
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (forall (registry :: JobPayloadRegistry).
RegistryConcurrencyPolicies registry =>
[(SchemaName, Bool)]
registryConcurrencyTables @registry)
      rlTables :: Map SchemaName Bool
rlTables = [(SchemaName, Bool)] -> Map SchemaName Bool
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (forall (registry :: JobPayloadRegistry).
RegistryRateLimitPolicies registry =>
[(SchemaName, Bool)]
registryRateLimitTables @registry)
      admissionFor :: SchemaName -> TableAdmission
admissionFor SchemaName
table =
        TableAdmission
          { tableConcurrency :: Bool
tableConcurrency = Bool -> SchemaName -> Map SchemaName Bool -> Bool
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Bool
False SchemaName
table Map SchemaName Bool
ccTables
          , tableRateLimit :: Bool
tableRateLimit = Bool -> SchemaName -> Map SchemaName Bool -> Bool
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Bool
False SchemaName
table Map SchemaName Bool
rlTables
          }
      tables :: [(SchemaName, TableAdmission)]
tables = [(SchemaName
table, SchemaName -> TableAdmission
admissionFor SchemaName
table) | SchemaName
table <- Proxy registry -> [SchemaName]
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> [SchemaName]
registryTableNames Proxy registry
proxy]
      -- Policies come from each payload's 'rateLimitFor' and 'concurrencyFor' selectors.
      seeds :: AdmissionSeeds
seeds =
        AdmissionSeeds
          { seedRateLimitPolicies :: [RateLimitPolicy]
seedRateLimitPolicies = Set RateLimitPolicy -> [RateLimitPolicy]
forall a. Set a -> [a]
Set.toList (forall (registry :: JobPayloadRegistry).
RegistryRateLimitPolicies registry =>
Set RateLimitPolicy
registryRateLimitPolicies @registry)
          , seedConcurrencyPolicies :: [ConcurrencyPolicy]
seedConcurrencyPolicies = Set ConcurrencyPolicy -> [ConcurrencyPolicy]
forall a. Set a -> [a]
Set.toList (forall (registry :: JobPayloadRegistry).
RegistryConcurrencyPolicies registry =>
Set ConcurrencyPolicy
registryConcurrencyPolicies @registry)
          }
  ByteString
-> SchemaName
-> [(SchemaName, TableAdmission)]
-> MigrationConfig
-> AdmissionSeeds
-> IO (MigrationResult String)
runMigrationsTrackedForTables ByteString
connStr SchemaName
schemaName [(SchemaName, TableAdmission)]
tables MigrationConfig
config AdmissionSeeds
seeds

-- | Admission policy rows to seed after a successful migration.
data AdmissionSeeds = AdmissionSeeds
  { AdmissionSeeds -> [RateLimitPolicy]
seedRateLimitPolicies :: [RateLimitPolicy]
  -- ^ Rate-limit policies to upsert into the policies table.
  , AdmissionSeeds -> [ConcurrencyPolicy]
seedConcurrencyPolicies :: [ConcurrencyPolicy]
  -- ^ Concurrency policies to upsert into the policies table.
  }

-- | Seeds for a deployment with no admission policies.
noAdmissionSeeds :: AdmissionSeeds
noAdmissionSeeds :: AdmissionSeeds
noAdmissionSeeds = AdmissionSeeds {seedRateLimitPolicies :: [RateLimitPolicy]
seedRateLimitPolicies = [], seedConcurrencyPolicies :: [ConcurrencyPolicy]
seedConcurrencyPolicies = []}

-- | Which admission trigger kinds a table's payload declares. Trigger migrations are
-- install-only. A kind removed from a payload keeps its triggers.
data TableAdmission = TableAdmission
  { TableAdmission -> Bool
tableConcurrency :: Bool
  -- ^ Install the concurrency triggers.
  , TableAdmission -> Bool
tableRateLimit :: Bool
  -- ^ Install the rate-limit bucket triggers.
  }
  deriving stock (TableAdmission -> TableAdmission -> Bool
(TableAdmission -> TableAdmission -> Bool)
-> (TableAdmission -> TableAdmission -> Bool) -> Eq TableAdmission
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TableAdmission -> TableAdmission -> Bool
== :: TableAdmission -> TableAdmission -> Bool
$c/= :: TableAdmission -> TableAdmission -> Bool
/= :: TableAdmission -> TableAdmission -> Bool
Eq, Int -> TableAdmission -> ShowS
[TableAdmission] -> ShowS
TableAdmission -> String
(Int -> TableAdmission -> ShowS)
-> (TableAdmission -> String)
-> ([TableAdmission] -> ShowS)
-> Show TableAdmission
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TableAdmission -> ShowS
showsPrec :: Int -> TableAdmission -> ShowS
$cshow :: TableAdmission -> String
show :: TableAdmission -> String
$cshowList :: [TableAdmission] -> ShowS
showList :: [TableAdmission] -> ShowS
Show)

-- | Install every admission trigger kind.
allTableAdmission :: TableAdmission
allTableAdmission :: TableAdmission
allTableAdmission = TableAdmission {tableConcurrency :: Bool
tableConcurrency = Bool
True, tableRateLimit :: Bool
tableRateLimit = Bool
True}

-- | The longest queue name whose generated identifiers stay distinct after PostgreSQL
-- truncates them to 'maxIdentifierBytes'.
maxQueueNameBytes :: Int
maxQueueNameBytes :: Int
maxQueueNameBytes = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Int -> Bool
identifiersDistinct [Int
1 .. Int
maxIdentifierBytes])
  where
    identifiersDistinct :: Int -> Bool
identifiersDistinct Int
len = [SchemaName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ((SchemaName -> SchemaName)
-> (SchemaName -> SchemaName) -> [SchemaName] -> [SchemaName]
forall params row.
Ord params =>
(row -> SchemaName) -> (row -> params) -> [row] -> [SchemaName]
conflictingPrefixes (Int -> SchemaName -> SchemaName
T.take Int
maxIdentifierBytes) SchemaName -> SchemaName
forall a. a -> a
id (SchemaName -> [SchemaName]
renderedIdentifiers (Int -> SchemaName -> SchemaName
T.replicate Int
len SchemaName
"q")))

-- | All identifiers in a queue's rendered DDL.
renderedIdentifiers :: TableName -> [Text]
renderedIdentifiers :: SchemaName -> [SchemaName]
renderedIdentifiers SchemaName
table =
  (SchemaName -> [SchemaName]) -> [SchemaName] -> [SchemaName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap SchemaName -> [SchemaName]
quoted ([SchemaName] -> [SchemaName]) -> [SchemaName] -> [SchemaName]
forall a b. (a -> b) -> a -> b
$
    [ByteString -> SchemaName
decodeUtf8 ByteString
body | MigrationScript String
_ ByteString
body <- SchemaName -> SchemaName -> TableAdmission -> [MigrationCommand]
jobQueueMigrationsForTable SchemaName
probeSchema SchemaName
table TableAdmission
allTableAdmission]
      [SchemaName] -> [SchemaName] -> [SchemaName]
forall a. Semigroup a => a -> a -> a
<> [ SchemaName -> SchemaName -> SchemaName
createNotifyFunctionSQL SchemaName
probeSchema SchemaName
table
         , SchemaName -> SchemaName -> SchemaName
createNotifyTriggerSQL SchemaName
probeSchema SchemaName
table
         , SchemaName -> SchemaName -> SchemaName
createEventStreamingTriggersSQL SchemaName
probeSchema SchemaName
table
         ]
  where
    probeSchema :: SchemaName
probeSchema = SchemaName
"arbiter"
    quoted :: SchemaName -> [SchemaName]
quoted = [SchemaName] -> [SchemaName]
forall {a}. [a] -> [a]
everyOther ([SchemaName] -> [SchemaName])
-> (SchemaName -> [SchemaName]) -> SchemaName -> [SchemaName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => SchemaName -> SchemaName -> [SchemaName]
SchemaName -> SchemaName -> [SchemaName]
T.splitOn SchemaName
"\""
    everyOther :: [a] -> [a]
everyOther (a
_ : a
name : [a]
rest) = a
name a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
everyOther [a]
rest
    everyOther [a]
_ = []

-- | Reject an empty or over-long schema name, and queue names that are empty, exceed
-- 'maxQueueNameBytes', generate a schema-wide arbiter table, or generate the same
-- object or channel name as another queue.
validateRegistryNames :: SchemaName -> [TableName] -> Either Text ()
validateRegistryNames :: SchemaName -> [SchemaName] -> Either SchemaName ()
validateRegistryNames SchemaName
schemaName [SchemaName]
tables
  | SchemaName -> Bool
T.null SchemaName
schemaName = SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left SchemaName
"Arbiter schema name must not be empty"
  | SchemaName -> Int
byteLength SchemaName
schemaName Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIdentifierBytes =
      SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left (SchemaName
"Arbiter schema name exceeds PostgreSQL's " SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> Int -> SchemaName
forall a. Show a => a -> SchemaName
showText Int
maxIdentifierBytes SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
"-byte identifier limit")
  | (SchemaName -> Bool) -> [SchemaName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any SchemaName -> Bool
T.null [SchemaName]
tables = SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left SchemaName
"Arbiter queue name must not be empty"
  | Just SchemaName
table <- (SchemaName -> Bool) -> [SchemaName] -> Maybe SchemaName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxQueueNameBytes) (Int -> Bool) -> (SchemaName -> Int) -> SchemaName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SchemaName -> Int
byteLength) [SchemaName]
tables =
      SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left
        ( SchemaName
"Arbiter queue name exceeds the "
            SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> Int -> SchemaName
forall a. Show a => a -> SchemaName
showText Int
maxQueueNameBytes
            SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
"-byte generated-identifier limit: "
            SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
table
        )
  | Just SchemaName
reserved <- Maybe SchemaName
reservedCollision =
      SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left (SchemaName
"Arbiter queue name generates a reserved arbiter table: " SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
reserved)
  | Just SchemaName
generated <- Maybe SchemaName
generatedCollision =
      SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left (SchemaName
"Arbiter queue names generate the same PostgreSQL object: " SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
generated)
  | Just SchemaName
channel <- Maybe SchemaName
channelCollision =
      SchemaName -> Either SchemaName ()
forall a b. a -> Either a b
Left (SchemaName
"Arbiter queue names generate the same notification channel: " SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
channel)
  | Bool
otherwise = () -> Either SchemaName ()
forall a b. b -> Either a b
Right ()
  where
    byteLength :: SchemaName -> Int
byteLength = ByteString -> Int
BS.length (ByteString -> Int)
-> (SchemaName -> ByteString) -> SchemaName -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SchemaName -> ByteString
encodeUtf8
    generatedCollision :: Maybe SchemaName
generatedCollision =
      [(SchemaName, SchemaName)] -> Maybe SchemaName
sharedName [(SchemaName
generated, SchemaName
table) | SchemaName
table <- [SchemaName]
tables, SchemaName
generated <- SchemaName -> [SchemaName]
queueTableNames SchemaName
table]
    reservedCollision :: Maybe SchemaName
reservedCollision =
      (SchemaName -> Bool) -> [SchemaName] -> Maybe SchemaName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (SchemaName -> Set SchemaName -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` [SchemaName] -> Set SchemaName
forall a. Ord a => [a] -> Set a
Set.fromList ((SchemaName -> [SchemaName]) -> [SchemaName] -> [SchemaName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap SchemaName -> [SchemaName]
queueTableNames [SchemaName]
tables)) [SchemaName]
sharedArbiterTables
    channelCollision :: Maybe SchemaName
channelCollision =
      [(SchemaName, SchemaName)] -> Maybe SchemaName
sharedName
        [ (SchemaName -> SchemaName -> SchemaName
channel SchemaName
schemaName SchemaName
table, SchemaName
table)
        | SchemaName
table <- [SchemaName]
tables
        , SchemaName -> SchemaName -> SchemaName
channel <- [SchemaName -> SchemaName -> SchemaName
pauseNotifyChannel, SchemaName -> SchemaName -> SchemaName
cancelNotifyChannel]
        ]
    -- The first generated name more than one queue claims.
    sharedName :: [(SchemaName, SchemaName)] -> Maybe SchemaName
sharedName = [SchemaName] -> Maybe SchemaName
forall a. [a] -> Maybe a
listToMaybe ([SchemaName] -> Maybe SchemaName)
-> ([(SchemaName, SchemaName)] -> [SchemaName])
-> [(SchemaName, SchemaName)]
-> Maybe SchemaName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((SchemaName, SchemaName) -> SchemaName)
-> ((SchemaName, SchemaName) -> SchemaName)
-> [(SchemaName, SchemaName)]
-> [SchemaName]
forall params row.
Ord params =>
(row -> SchemaName) -> (row -> params) -> [row] -> [SchemaName]
conflictingPrefixes (SchemaName, SchemaName) -> SchemaName
forall a b. (a, b) -> a
fst (SchemaName, SchemaName) -> SchemaName
forall a b. (a, b) -> b
snd

-- | Migrate several tables in one schema. On success, seed the policies, then reconcile
-- durability and triggers on the same connection. The table list must be the schema's
-- whole queue set. An omitted queue loses its notify and event-streaming objects.
-- Names that fail 'validateRegistryNames' return 'MigrationError' before the
-- connection opens.
runMigrationsTrackedForTables
  :: ByteString
  -- ^ Database connection string
  -> SchemaName
  -- ^ Schema for the arbiter tables
  -> [(TableName, TableAdmission)]
  -- ^ Every queue table in the schema, with its admission trigger kinds
  -> MigrationConfig
  -- ^ Migration options
  -> AdmissionSeeds
  -- ^ Admission policy rows to seed
  -> IO (MigrationResult String)
runMigrationsTrackedForTables :: ByteString
-> SchemaName
-> [(SchemaName, TableAdmission)]
-> MigrationConfig
-> AdmissionSeeds
-> IO (MigrationResult String)
runMigrationsTrackedForTables ByteString
connStr SchemaName
schemaName [(SchemaName, TableAdmission)]
tableAdmissions MigrationConfig
config AdmissionSeeds
seeds =
  case SchemaName -> [SchemaName] -> Either SchemaName ()
validateRegistryNames SchemaName
schemaName (((SchemaName, TableAdmission) -> SchemaName)
-> [(SchemaName, TableAdmission)] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map (SchemaName, TableAdmission) -> SchemaName
forall a b. (a, b) -> a
fst [(SchemaName, TableAdmission)]
tableAdmissions) of
    Left SchemaName
err -> MigrationResult String -> IO (MigrationResult String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> MigrationResult String
forall a. a -> MigrationResult a
MigrationError (SchemaName -> String
T.unpack SchemaName
err))
    Right () -> IO Connection
-> (Connection -> IO ())
-> (Connection -> IO (MigrationResult String))
-> IO (MigrationResult String)
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (ByteString -> IO Connection
connectPostgreSQL ByteString
connStr) Connection -> IO ()
close ((Connection -> IO (MigrationResult String))
 -> IO (MigrationResult String))
-> (Connection -> IO (MigrationResult String))
-> IO (MigrationResult String)
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
      Connection
-> SchemaName
-> Maybe NominalDiffTime
-> IO (MigrationResult String)
-> IO (MigrationResult String)
withMigrationLock Connection
conn SchemaName
schemaName (MigrationConfig -> Maybe NominalDiffTime
migrationLockTimeout MigrationConfig
config) (IO (MigrationResult String) -> IO (MigrationResult String))
-> IO (MigrationResult String) -> IO (MigrationResult String)
forall a b. (a -> b) -> a -> b
$
        Connection
-> SchemaName
-> [(SchemaName, TableAdmission)]
-> MigrationConfig
-> AdmissionSeeds
-> IO (MigrationResult String)
migrateSchema Connection
conn SchemaName
schemaName [(SchemaName, TableAdmission)]
tableAdmissions MigrationConfig
config AdmissionSeeds
seeds

-- | Hold the schema's session-level migration lock for the run. Replicas that migrate at
-- the same time run one after another. A session lock spans the reconciles, which run
-- after the tracked migrations commit.
withMigrationLock
  :: PG.Connection
  -> SchemaName
  -> Maybe NominalDiffTime
  -> IO (MigrationResult String)
  -> IO (MigrationResult String)
withMigrationLock :: Connection
-> SchemaName
-> Maybe NominalDiffTime
-> IO (MigrationResult String)
-> IO (MigrationResult String)
withMigrationLock Connection
conn SchemaName
schemaName Maybe NominalDiffTime
timeout IO (MigrationResult String)
act = IO Bool
-> (Bool -> IO ())
-> (Bool -> IO (MigrationResult String))
-> IO (MigrationResult String)
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO Bool
acquire Bool -> IO ()
release Bool -> IO (MigrationResult String)
run
  where
    lockName :: Only SchemaName
lockName = SchemaName -> Only SchemaName
forall a. a -> Only a
Only (SchemaName
"arbiter.migrations:" SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
schemaName)
    acquire :: IO Bool
acquire =
      IO () -> IO (Either SqlError ())
forall e a. Exception e => IO a -> IO (Either e a)
try (Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
PG.withTransaction Connection
conn (IO ()
setTimeouts IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
lock)) IO (Either SqlError ())
-> (Either SqlError () -> IO Bool) -> IO Bool
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Right () -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
        Left SqlError
sqlError
          | Maybe NominalDiffTime -> Bool
forall a. Maybe a -> Bool
isJust Maybe NominalDiffTime
timeout, SqlError -> ByteString
PG.sqlState SqlError
sqlError ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
lockNotAvailable -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
          | Bool
otherwise -> SqlError -> IO Bool
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SqlError
sqlError
    lock :: IO ()
lock = IO [Only (Maybe SchemaName)] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection
-> Query -> Only SchemaName -> IO [Only (Maybe SchemaName)]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn Query
"SELECT pg_advisory_lock(hashtextextended(?, 0))::text" Only SchemaName
lockName :: IO [Only (Maybe Text)])
    -- Pinned for this transaction.
    setTimeouts :: IO ()
setTimeouts = ((SchemaName, SchemaName) -> IO ())
-> [(SchemaName, SchemaName)] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (SchemaName, SchemaName) -> IO ()
setLocal [(SchemaName
"lock_timeout", SchemaName
-> (NominalDiffTime -> SchemaName)
-> Maybe NominalDiffTime
-> SchemaName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe SchemaName
"0" NominalDiffTime -> SchemaName
millis Maybe NominalDiffTime
timeout), (SchemaName
"statement_timeout", SchemaName
"0")]
    setLocal :: (SchemaName, SchemaName) -> IO ()
setLocal (SchemaName
name, SchemaName
value) =
      IO [Only SchemaName] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection
-> Query -> (SchemaName, SchemaName) -> IO [Only SchemaName]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn Query
"SELECT set_config(?, ?, TRUE)" (SchemaName
name :: Text, SchemaName
value :: Text) :: IO [Only Text])
    millis :: NominalDiffTime -> SchemaName
millis NominalDiffTime
seconds = Int -> SchemaName
forall a. Show a => a -> SchemaName
showText (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
maxLockTimeoutMillis (NominalDiffTime -> Int
forall b. Integral b => NominalDiffTime -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (NominalDiffTime
seconds NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
millisPerSecond) :: Int)))
    maxLockTimeoutMillis :: Int
maxLockTimeoutMillis = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int32
forall a. Bounded a => a
maxBound :: Int32)
    millisPerSecond :: NominalDiffTime
millisPerSecond = NominalDiffTime
1000
    lockNotAvailable :: ByteString
lockNotAvailable = ByteString
"55P03"
    release :: Bool -> IO ()
release Bool
locked =
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
locked (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        IO (Either SqlError [Only Bool]) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO [Only Bool] -> IO (Either SqlError [Only Bool])
forall e a. Exception e => IO a -> IO (Either e a)
try IO [Only Bool]
unlock :: IO (Either PG.SqlError [Only Bool]))
    unlock :: IO [Only Bool]
unlock = Connection -> Query -> Only SchemaName -> IO [Only Bool]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn Query
"SELECT pg_advisory_unlock(hashtextextended(?, 0))" Only SchemaName
lockName :: IO [Only Bool]
    run :: Bool -> IO (MigrationResult String)
run Bool
locked
      | Bool
locked = IO (MigrationResult String)
act
      | Bool
otherwise =
          MigrationResult String -> IO (MigrationResult String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> MigrationResult String
forall a. a -> MigrationResult a
MigrationError (String
"Timed out waiting on the arbiter migration lock for schema " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> SchemaName -> String
T.unpack SchemaName
schemaName))

-- | Apply a schema's tracked migrations and reconcile it, under the migration lock.
migrateSchema
  :: PG.Connection
  -> SchemaName
  -> [(TableName, TableAdmission)]
  -> MigrationConfig
  -> AdmissionSeeds
  -> IO (MigrationResult String)
migrateSchema :: Connection
-> SchemaName
-> [(SchemaName, TableAdmission)]
-> MigrationConfig
-> AdmissionSeeds
-> IO (MigrationResult String)
migrateSchema Connection
conn SchemaName
schemaName [(SchemaName, TableAdmission)]
tableAdmissions MigrationConfig
config AdmissionSeeds
seeds = do
  Connection -> (Connection -> IO ()) -> IO ()
forall a. Connection -> (Connection -> IO a) -> IO a
withConnection Connection
conn Connection -> IO ()
LibPQ.disableNoticeReporting

  Connection -> SchemaName -> IO ()
ensureSchema Connection
conn SchemaName
schemaName

  Connection -> (Connection -> IO ()) -> IO ()
forall a. Connection -> (Connection -> IO a) -> IO a
withConnection Connection
conn Connection -> IO ()
LibPQ.enableNoticeReporting

  let schemaMigrations :: [MigrationCommand]
schemaMigrations = SchemaName -> [MigrationCommand]
schemaLevelMigrations SchemaName
schemaName
      tableMigrations :: [MigrationCommand]
tableMigrations = ((SchemaName, TableAdmission) -> [MigrationCommand])
-> [(SchemaName, TableAdmission)] -> [MigrationCommand]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((SchemaName -> TableAdmission -> [MigrationCommand])
-> (SchemaName, TableAdmission) -> [MigrationCommand]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (SchemaName -> SchemaName -> TableAdmission -> [MigrationCommand]
jobQueueMigrationsForTable SchemaName
schemaName)) [(SchemaName, TableAdmission)]
tableAdmissions
      migrations :: [MigrationCommand]
migrations = [MigrationCommand]
schemaMigrations [MigrationCommand] -> [MigrationCommand] -> [MigrationCommand]
forall a. Semigroup a => a -> a -> a
<> [MigrationCommand]
tableMigrations
      migrationTableName :: ByteString
migrationTableName = SchemaName -> ByteString
encodeUtf8 (SchemaName -> ByteString) -> SchemaName -> ByteString
forall a b. (a -> b) -> a -> b
$ SchemaName
schemaName SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
".schema_migrations"
      options :: MigrationOptions
options =
        MigrationOptions
defaultOptions
          { optVerbose = Quiet
          , optTableName = migrationTableName
          }

  IO (MigrationResult String) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection
-> MigrationOptions
-> [MigrationCommand]
-> IO (MigrationResult String)
runMigrations Connection
conn MigrationOptions
options [MigrationCommand
MigrationInitialization])
  migrationResult <- Connection
-> MigrationOptions
-> [MigrationCommand]
-> IO (MigrationResult String)
runMigrations Connection
conn MigrationOptions
options [MigrationCommand]
migrations
  case migrationResult of
    -- A reconcile throws on a conflicting prefix or a failed ALTER. Async exceptions propagate.
    MigrationResult String
MigrationSuccess ->
      (SomeException -> MigrationResult String)
-> (() -> MigrationResult String)
-> Either SomeException ()
-> MigrationResult String
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> MigrationResult String
forall a. a -> MigrationResult a
MigrationError (String -> MigrationResult String)
-> (SomeException -> String)
-> SomeException
-> MigrationResult String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SchemaName -> String
T.unpack (SchemaName -> String)
-> (SomeException -> SchemaName) -> SomeException -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> SchemaName
displayEx) (MigrationResult String -> () -> MigrationResult String
forall a b. a -> b -> a
const MigrationResult String
forall a. MigrationResult a
MigrationSuccess) (Either SomeException () -> MigrationResult String)
-> IO (Either SomeException ()) -> IO (MigrationResult String)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO () -> IO (Either SomeException ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny IO ()
reconcile
    MigrationResult String
other -> MigrationResult String -> IO (MigrationResult String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MigrationResult String
other
  where
    reconcile :: IO ()
reconcile = do
      Connection -> SchemaName -> [SchemaName] -> IO ()
registerQueues Connection
conn SchemaName
schemaName (((SchemaName, TableAdmission) -> SchemaName)
-> [(SchemaName, TableAdmission)] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map (SchemaName, TableAdmission) -> SchemaName
forall a b. (a, b) -> a
fst [(SchemaName, TableAdmission)]
tableAdmissions)
      Connection -> SchemaName -> [RateLimitPolicy] -> IO ()
reconcileRateLimitPolicies Connection
conn SchemaName
schemaName (AdmissionSeeds -> [RateLimitPolicy]
seedRateLimitPolicies AdmissionSeeds
seeds)
      Connection -> SchemaName -> [ConcurrencyPolicy] -> IO ()
reconcileConcurrencyPolicies Connection
conn SchemaName
schemaName (AdmissionSeeds -> [ConcurrencyPolicy]
seedConcurrencyPolicies AdmissionSeeds
seeds)
      Connection -> SchemaName -> Durability -> IO ()
reconcileRateLimitDurability Connection
conn SchemaName
schemaName (MigrationConfig -> Durability
rateLimitDurability MigrationConfig
config)
      Connection
-> SchemaName -> [SchemaName] -> MigrationConfig -> IO ()
reconcileOptionalTriggers Connection
conn SchemaName
schemaName (((SchemaName, TableAdmission) -> SchemaName)
-> [(SchemaName, TableAdmission)] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map (SchemaName, TableAdmission) -> SchemaName
forall a b. (a, b) -> a
fst [(SchemaName, TableAdmission)]
tableAdmissions) MigrationConfig
config

-- | Add an @arbiter_queues@ row for each migrated queue. Maintenance sweeps every
-- queue listed there, including queues outside the pool's registry. Existing rows
-- keep their state.
registerQueues :: PG.Connection -> SchemaName -> [TableName] -> IO ()
registerQueues :: Connection -> SchemaName -> [SchemaName] -> IO ()
registerQueues Connection
conn SchemaName
schemaName [SchemaName]
tables =
  IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection -> Query -> Only (PGArray SchemaName) -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute Connection
conn (ByteString -> Query
Query (SchemaName -> ByteString
encodeUtf8 (SchemaName -> SchemaName
registerQueuesSQL SchemaName
schemaName))) (PGArray SchemaName -> Only (PGArray SchemaName)
forall a. a -> Only a
Only ([SchemaName] -> PGArray SchemaName
forall a. [a] -> PGArray a
PGArray [SchemaName]
tables)))

-- | Create the schema. A failed CREATE is fine when the schema exists.
ensureSchema :: PG.Connection -> SchemaName -> IO ()
ensureSchema :: Connection -> SchemaName -> IO ()
ensureSchema Connection
conn SchemaName
schemaName =
  Connection -> SchemaName -> IO ()
executeSQL Connection
conn (SchemaName -> SchemaName
createSchemaSQL SchemaName
schemaName) IO () -> (SqlError -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SqlError
createError :: PG.SqlError) -> do
    present <- Connection -> SchemaName -> IO Bool
schemaExists Connection
conn SchemaName
schemaName
    unless present $
      ioError
        ( userError $
            "Failed to create schema "
              <> T.unpack schemaName
              <> " and it does not exist. Either grant CREATE privilege on the database"
              <> " or create the schema manually: CREATE SCHEMA "
              <> T.unpack schemaName
              <> ";"
              <> "\nOriginal error: "
              <> show createError
        )

-- | Upsert each policy's @default_*@ params into the policies table. Operator
-- @override_*@ values stay intact. Idempotent. Overrides and removed prefixes survive
-- a deploy.
reconcileRateLimitPolicies :: PG.Connection -> SchemaName -> [RateLimitPolicy] -> IO ()
reconcileRateLimitPolicies :: Connection -> SchemaName -> [RateLimitPolicy] -> IO ()
reconcileRateLimitPolicies =
  String
-> String
-> (RateLimitPolicy -> SchemaName)
-> (RateLimitPolicy -> String)
-> (SchemaName -> RateLimitPolicy -> SchemaName)
-> Connection
-> SchemaName
-> [RateLimitPolicy]
-> IO ()
forall params row.
Ord params =>
String
-> String
-> (row -> SchemaName)
-> (row -> params)
-> (SchemaName -> row -> SchemaName)
-> Connection
-> SchemaName
-> [row]
-> IO ()
reconcilePolicyRows String
"rate-limit policy" String
"parameters" RateLimitPolicy -> SchemaName
policyPrefix RateLimitPolicy -> String
policyParamsKey SchemaName -> RateLimitPolicy -> SchemaName
upsertRateLimitPolicyRowSQL

-- | Rate-limit policy prefixes declared with more than one parameter set.
conflictingPolicyPrefixes :: [RateLimitPolicy] -> [Text]
conflictingPolicyPrefixes :: [RateLimitPolicy] -> [SchemaName]
conflictingPolicyPrefixes = (RateLimitPolicy -> SchemaName)
-> (RateLimitPolicy -> String) -> [RateLimitPolicy] -> [SchemaName]
forall params row.
Ord params =>
(row -> SchemaName) -> (row -> params) -> [row] -> [SchemaName]
conflictingPrefixes RateLimitPolicy -> SchemaName
policyPrefix RateLimitPolicy -> String
policyParamsKey

-- | A canonical conflict key for a policy's params, rendered as text. A NaN compares
-- equal to itself.
policyParamsKey :: RateLimitPolicy -> String
policyParamsKey :: RateLimitPolicy -> String
policyParamsKey RateLimitPolicy
policy = (Double, Double, NominalDiffTime) -> String
forall a. Show a => a -> String
show (RateLimitPolicy -> Double
policyMax RateLimitPolicy
policy, RateLimitPolicy -> Double
policyRefill RateLimitPolicy
policy, RateLimitPolicy -> NominalDiffTime
policyInterval RateLimitPolicy
policy)

-- | Upsert each policy's @default_limit@, leaving operator overrides intact. Two policies
-- with the same prefix but different limits fail the migration.
reconcileConcurrencyPolicies :: PG.Connection -> SchemaName -> [ConcurrencyPolicy] -> IO ()
reconcileConcurrencyPolicies :: Connection -> SchemaName -> [ConcurrencyPolicy] -> IO ()
reconcileConcurrencyPolicies =
  String
-> String
-> (ConcurrencyPolicy -> SchemaName)
-> (ConcurrencyPolicy -> Int32)
-> (SchemaName -> ConcurrencyPolicy -> SchemaName)
-> Connection
-> SchemaName
-> [ConcurrencyPolicy]
-> IO ()
forall params row.
Ord params =>
String
-> String
-> (row -> SchemaName)
-> (row -> params)
-> (SchemaName -> row -> SchemaName)
-> Connection
-> SchemaName
-> [row]
-> IO ()
reconcilePolicyRows String
"concurrency policy" String
"limits" ConcurrencyPolicy -> SchemaName
cpPrefix ConcurrencyPolicy -> Int32
cpLimit SchemaName -> ConcurrencyPolicy -> SchemaName
upsertConcurrencyPolicyRowSQL

-- | Upsert each policy row's @default_*@ params. Fails on a prefix that contains the
-- @:@ key separator or is declared with conflicting params. Idempotent. Operator
-- overrides and removed prefixes survive a deploy.
reconcilePolicyRows
  :: (Ord params)
  => String
  -> String
  -> (row -> Text)
  -> (row -> params)
  -> (SchemaName -> row -> Text)
  -> PG.Connection
  -> SchemaName
  -> [row]
  -> IO ()
reconcilePolicyRows :: forall params row.
Ord params =>
String
-> String
-> (row -> SchemaName)
-> (row -> params)
-> (SchemaName -> row -> SchemaName)
-> Connection
-> SchemaName
-> [row]
-> IO ()
reconcilePolicyRows String
label String
noun row -> SchemaName
prefixOf row -> params
paramsOf SchemaName -> row -> SchemaName
upsertSQL Connection
conn SchemaName
schemaName [row]
rows
  | Just SchemaName
prefix <- (SchemaName -> Bool) -> [SchemaName] -> Maybe SchemaName
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (SchemaName -> SchemaName -> Bool
T.isInfixOf SchemaName
":") ((row -> SchemaName) -> [row] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map row -> SchemaName
prefixOf [row]
rows) =
      IOError -> IO ()
forall a. HasCallStack => IOError -> IO a
ioError
        ( String -> IOError
userError (String -> IOError) -> String -> IOError
forall a b. (a -> b) -> a -> b
$
            String
label
              String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" prefix "
              String -> ShowS
forall a. Semigroup a => a -> a -> a
<> SchemaName -> String
T.unpack SchemaName
prefix
              String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" contains ':', the key separator, so its keys would alias another prefix's. Use a colon-free prefix."
        )
  | (SchemaName
prefix : [SchemaName]
_) <- (row -> SchemaName) -> (row -> params) -> [row] -> [SchemaName]
forall params row.
Ord params =>
(row -> SchemaName) -> (row -> params) -> [row] -> [SchemaName]
conflictingPrefixes row -> SchemaName
prefixOf row -> params
paramsOf [row]
rows =
      IOError -> IO ()
forall a. HasCallStack => IOError -> IO a
ioError
        ( String -> IOError
userError (String -> IOError) -> String -> IOError
forall a b. (a -> b) -> a -> b
$
            String
label String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" prefix " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> SchemaName -> String
T.unpack SchemaName
prefix String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" is declared with conflicting " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
noun String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
". Give each a unique prefix."
        )
  | Bool
otherwise =
      Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
PG.withTransaction Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        (row -> IO ()) -> [row] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Connection -> SchemaName -> IO ()
executeSQL Connection
conn (SchemaName -> IO ()) -> (row -> SchemaName) -> row -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SchemaName -> row -> SchemaName
upsertSQL SchemaName
schemaName) [row]
rows

-- | Prefixes declared with more than one distinct parameter set.
conflictingPrefixes :: (Ord params) => (row -> Text) -> (row -> params) -> [row] -> [Text]
conflictingPrefixes :: forall params row.
Ord params =>
(row -> SchemaName) -> (row -> params) -> [row] -> [SchemaName]
conflictingPrefixes row -> SchemaName
prefixOf row -> params
paramsOf [row]
rows =
  Map SchemaName (Set params) -> [SchemaName]
forall k a. Map k a -> [k]
Map.keys ((Set params -> Bool)
-> Map SchemaName (Set params) -> Map SchemaName (Set params)
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (Int -> Bool) -> (Set params -> Int) -> Set params -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set params -> Int
forall a. Set a -> Int
Set.size) Map SchemaName (Set params)
paramsByPrefix)
  where
    paramsByPrefix :: Map SchemaName (Set params)
paramsByPrefix =
      (Set params -> Set params -> Set params)
-> [(SchemaName, Set params)] -> Map SchemaName (Set params)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Set params -> Set params -> Set params
forall a. Ord a => Set a -> Set a -> Set a
Set.union [(row -> SchemaName
prefixOf row
row, params -> Set params
forall a. a -> Set a
Set.singleton (row -> params
paramsOf row
row)) | row
row <- [row]
rows]

-- | Converge the bucket table's WAL persistence to the declared durability. Reads
-- the current @pg_class.relpersistence@ and issues @SET LOGGED@ or @SET UNLOGGED@ on a
-- change. The ALTER rewrites the table under @ACCESS EXCLUSIVE@. Token consumes block
-- briefly while a switch runs.
reconcileRateLimitDurability :: PG.Connection -> SchemaName -> Durability -> IO ()
reconcileRateLimitDurability :: Connection -> SchemaName -> Durability -> IO ()
reconcileRateLimitDurability Connection
conn SchemaName
schemaName Durability
durability = do
  rows <-
    Connection
-> Query -> (SchemaName, SchemaName) -> IO [Only SchemaName]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query
      Connection
conn
      Query
"SELECT relation.relpersistence::text FROM pg_class relation \
      \JOIN pg_namespace namespace ON namespace.oid = relation.relnamespace \
      \WHERE namespace.nspname = ? AND relation.relname = ?"
      (SchemaName
schemaName, SchemaName
arbiterRateLimitsTableName)
  let target = case Durability
durability of
        Durability
Durable -> SchemaName
"p" :: Text
        Durability
Unlogged -> SchemaName
"u"
  traverse_
    (\(Only SchemaName
current) -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (SchemaName
current SchemaName -> SchemaName -> Bool
forall a. Eq a => a -> a -> Bool
/= SchemaName
target) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> SchemaName -> IO ()
executeSQL Connection
conn (SchemaName -> Durability -> SchemaName
alterRateLimitsDurabilitySQL SchemaName
schemaName Durability
durability))
    (listToMaybe rows)

-- | Reconcile optional triggers schema-wide after tracked migrations. Install or drop
-- each option's objects to match the config. Drop objects for queues outside the
-- supplied table list. Migration history remains unchanged.
reconcileOptionalTriggers :: PG.Connection -> SchemaName -> [TableName] -> MigrationConfig -> IO ()
reconcileOptionalTriggers :: Connection
-> SchemaName -> [SchemaName] -> MigrationConfig -> IO ()
reconcileOptionalTriggers Connection
conn SchemaName
schemaName [SchemaName]
tables MigrationConfig
config =
  Connection -> IO () -> IO ()
forall a. Connection -> IO a -> IO a
PG.withTransaction Connection
conn (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    Connection -> SchemaName -> [SchemaName] -> IO ()
adoptUnmarkedObjects Connection
conn SchemaName
schemaName [SchemaName]
tables
    Connection -> SchemaName -> [SchemaName] -> Bool -> IO ()
reconcileNotify Connection
conn SchemaName
schemaName [SchemaName]
tables (MigrationConfig -> Bool
enableNotifications MigrationConfig
config)
    Connection -> SchemaName -> [SchemaName] -> Bool -> IO ()
reconcileEventStreaming Connection
conn SchemaName
schemaName [SchemaName]
tables (MigrationConfig -> Bool
enableEventStreaming MigrationConfig
config)

-- | Install notify objects for each queue and drop the rest, or drop them all when disabled.
reconcileNotify :: PG.Connection -> SchemaName -> [TableName] -> Bool -> IO ()
reconcileNotify :: Connection -> SchemaName -> [SchemaName] -> Bool -> IO ()
reconcileNotify Connection
conn SchemaName
schemaName [SchemaName]
tables Bool
enabled
  | Bool
enabled = do
      Connection -> SchemaName -> [SchemaName] -> IO ()
dropNotifyObjectsExcept Connection
conn SchemaName
schemaName [SchemaName]
tables
      (SchemaName -> IO ()) -> [SchemaName] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Connection -> SchemaName -> SchemaName -> IO ()
installNotify Connection
conn SchemaName
schemaName) [SchemaName]
tables
  | Bool
otherwise = Connection -> SchemaName -> [SchemaName] -> IO ()
dropNotifyObjectsExcept Connection
conn SchemaName
schemaName []

-- | Install event-stream objects for each queue and drop the rest, or drop them all when disabled.
reconcileEventStreaming :: PG.Connection -> SchemaName -> [TableName] -> Bool -> IO ()
reconcileEventStreaming :: Connection -> SchemaName -> [SchemaName] -> Bool -> IO ()
reconcileEventStreaming Connection
conn SchemaName
schemaName [SchemaName]
tables Bool
enabled
  | Bool
enabled = do
      Connection -> SchemaName -> IO ()
executeSQL Connection
conn (SchemaName -> SchemaName
createEventStreamingFunctionSQL SchemaName
schemaName)
      Connection -> SchemaName -> [SchemaName] -> IO ()
dropEventTriggersExcept Connection
conn SchemaName
schemaName ([SchemaName] -> [SchemaName]
eventTables [SchemaName]
tables)
      (SchemaName -> IO ()) -> [SchemaName] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Connection -> SchemaName -> SchemaName -> IO ()
installEventStream Connection
conn SchemaName
schemaName) [SchemaName]
tables
  | Bool
otherwise = do
      Connection -> SchemaName -> [SchemaName] -> IO ()
dropEventTriggersExcept Connection
conn SchemaName
schemaName []
      Connection -> SchemaName -> IO ()
dropEventFunctionIfUnused Connection
conn SchemaName
schemaName

-- | Install one queue's notify function and trigger. Functions are replaced in place.
-- That takes no lock on the queue table. Triggers are rebuilt when their marker is off
-- the current version.
installNotify :: PG.Connection -> SchemaName -> TableName -> IO ()
installNotify :: Connection -> SchemaName -> SchemaName -> IO ()
installNotify Connection
conn SchemaName
schemaName SchemaName
table = do
  Connection -> SchemaName -> IO ()
executeSQL Connection
conn (SchemaName -> SchemaName -> SchemaName
createNotifyFunctionSQL SchemaName
schemaName SchemaName
table)
  current <- Connection
-> SchemaName -> SchemaName -> SchemaName -> SchemaName -> IO Bool
triggerIsCurrent Connection
conn SchemaName
schemaName SchemaName
table (SchemaName -> SchemaName
notifyTriggerName SchemaName
table) SchemaName
notifyObjectComment
  unless current $ executeSQL conn (createNotifyTriggerSQL schemaName table)

-- | Install one queue's event-stream triggers on its main and DLQ tables when either
-- marker is off the current version.
installEventStream :: PG.Connection -> SchemaName -> TableName -> IO ()
installEventStream :: Connection -> SchemaName -> SchemaName -> IO ()
installEventStream Connection
conn SchemaName
schemaName SchemaName
table = do
  mainCurrent <- Connection
-> SchemaName -> SchemaName -> SchemaName -> SchemaName -> IO Bool
triggerIsCurrent Connection
conn SchemaName
schemaName SchemaName
table (SchemaName -> SchemaName
eventStreamingTriggerName SchemaName
table) SchemaName
eventStreamingObjectComment
  dlqCurrent <-
    triggerIsCurrent conn schemaName (table <> dlqSuffix) (eventStreamingDLQTriggerName table) eventStreamingObjectComment
  unless (mainCurrent && dlqCurrent) $
    executeSQL conn (createEventStreamingTriggersSQL schemaName table)

-- | Drop the shared event-stream function when arbiter marked it and no trigger uses it.
dropEventFunctionIfUnused :: PG.Connection -> SchemaName -> IO ()
dropEventFunctionIfUnused :: Connection -> SchemaName -> IO ()
dropEventFunctionIfUnused Connection
conn SchemaName
schemaName = do
  ours <- IO Bool
eventFunctionIsMarked
  -- The function stays while foreign triggers depend on it.
  depended <- eventFunctionHasTriggers
  when (ours && not depended) $ executeSQL conn (dropEventStreamingFunctionSQL schemaName)
  where
    eventFunctionIsMarked :: IO Bool
eventFunctionIsMarked =
      Connection
-> Query -> (SchemaName, SchemaName, SchemaName) -> IO Bool
forall params.
ToRow params =>
Connection -> Query -> params -> IO Bool
exists
        Connection
conn
        ( Query
"SELECT EXISTS (SELECT 1 "
            Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
functionJoins
            Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND function.proname = ? AND function.pronargs = 0 \
               \AND obj_description(function.oid, 'pg_proc') LIKE ?)"
        )
        (SchemaName
schemaName, SchemaName
eventStreamingFunctionName, SchemaName
eventStreamingObjectCommentPrefix SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
"%")

    eventFunctionHasTriggers :: IO Bool
eventFunctionHasTriggers =
      Connection -> Query -> (SchemaName, SchemaName) -> IO Bool
forall params.
ToRow params =>
Connection -> Query -> params -> IO Bool
exists
        Connection
conn
        Query
"SELECT EXISTS (SELECT 1 FROM pg_trigger trigger JOIN pg_proc function ON function.oid = trigger.tgfoid \
        \JOIN pg_namespace namespace ON namespace.oid = function.pronamespace \
        \WHERE namespace.nspname = ? AND function.proname = ? AND NOT trigger.tgisinternal)"
        (SchemaName
schemaName, SchemaName
eventStreamingFunctionName)

-- | Drop event-stream triggers outside the @keep@ tables.
dropEventTriggersExcept :: PG.Connection -> SchemaName -> [TableName] -> IO ()
dropEventTriggersExcept :: Connection -> SchemaName -> [SchemaName] -> IO ()
dropEventTriggersExcept Connection
conn SchemaName
schemaName = Connection -> SchemaName -> SchemaName -> [SchemaName] -> IO ()
dropTriggersMarked Connection
conn SchemaName
schemaName SchemaName
eventStreamingObjectCommentPrefix

-- | Drop notify triggers and functions for queues outside @keep@.
dropNotifyObjectsExcept :: PG.Connection -> SchemaName -> [TableName] -> IO ()
dropNotifyObjectsExcept :: Connection -> SchemaName -> [SchemaName] -> IO ()
dropNotifyObjectsExcept Connection
conn SchemaName
schemaName [SchemaName]
keep = do
  Connection -> SchemaName -> SchemaName -> [SchemaName] -> IO ()
dropTriggersMarked Connection
conn SchemaName
schemaName SchemaName
notifyObjectCommentPrefix [SchemaName]
keep
  -- A function that still carries a trigger stays.
  Connection
-> Query -> (SchemaName, SchemaName, PGArray SchemaName) -> IO ()
forall params.
ToRow params =>
Connection -> Query -> params -> IO ()
executeRendered
    Connection
conn
    ( Query
"SELECT format('DROP FUNCTION IF EXISTS %I.%I();', namespace.nspname, function.proname) "
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
functionJoins
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND obj_description(function.oid, 'pg_proc') LIKE ? AND function.pronargs = 0 \
           \AND NOT (function.proname::text = ANY(?::text[])) \
           \AND NOT EXISTS (SELECT 1 FROM pg_trigger trigger WHERE trigger.tgfoid = function.oid AND NOT trigger.tgisinternal)"
    )
    (SchemaName
schemaName, SchemaName
notifyObjectCommentPrefix SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
"%", [SchemaName] -> PGArray SchemaName
forall a. [a] -> PGArray a
PGArray ((SchemaName -> SchemaName) -> [SchemaName] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map SchemaName -> SchemaName
notifyFunctionName [SchemaName]
keep))

-- | Objects installed before arbiter stamped markers carry no comment. Adoption
-- stamps a marker that no install writes. The sweeps match it. The currency probe
-- rejects it. Adoption covers the names this registry generates. An unmarked
-- lookalike stays foreign.
adoptUnmarkedObjects :: PG.Connection -> SchemaName -> [TableName] -> IO ()
adoptUnmarkedObjects :: Connection -> SchemaName -> [SchemaName] -> IO ()
adoptUnmarkedObjects Connection
conn SchemaName
schemaName [SchemaName]
tables = do
  Connection
-> Query -> (SchemaName, SchemaName, PGArray SchemaName) -> IO ()
forall params.
ToRow params =>
Connection -> Query -> params -> IO ()
executeRendered
    Connection
conn
    ( Query
commentFunctionSelect
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
functionJoins
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND function.proname = ANY(?::text[]) AND function.pronargs = 0 \
           \AND obj_description(function.oid, 'pg_proc') IS NULL"
    )
    (SchemaName
notifyAdoptedObjectComment, SchemaName
schemaName, [SchemaName] -> PGArray SchemaName
forall a. [a] -> PGArray a
PGArray ((SchemaName -> SchemaName) -> [SchemaName] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map SchemaName -> SchemaName
notifyFunctionName [SchemaName]
tables))
  SchemaName -> [SchemaName] -> [SchemaName] -> IO ()
adoptTriggers SchemaName
notifyAdoptedObjectComment ((SchemaName -> SchemaName) -> [SchemaName] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map SchemaName -> SchemaName
notifyTriggerName [SchemaName]
tables) [SchemaName]
tables
  SchemaName -> [SchemaName] -> [SchemaName] -> IO ()
adoptTriggers SchemaName
eventStreamingAdoptedObjectComment [SchemaName]
eventTriggerNames ([SchemaName] -> [SchemaName]
eventTables [SchemaName]
tables)
  -- The shared function is adopted when one of its triggers is.
  Connection
-> Query
-> (SchemaName, SchemaName, SchemaName, SchemaName)
-> IO ()
forall params.
ToRow params =>
Connection -> Query -> params -> IO ()
executeRendered
    Connection
conn
    ( Query
commentFunctionSelect
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
functionJoins
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND function.proname = ? AND function.pronargs = 0 \
           \AND obj_description(function.oid, 'pg_proc') IS NULL \
           \AND EXISTS (SELECT 1 FROM pg_trigger trigger WHERE trigger.tgfoid = function.oid AND NOT trigger.tgisinternal \
           \AND obj_description(trigger.oid, 'pg_trigger') LIKE ?)"
    )
    ( SchemaName
eventStreamingAdoptedObjectComment
    , SchemaName
schemaName
    , SchemaName
eventStreamingFunctionName
    , SchemaName
eventStreamingObjectCommentPrefix SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
"%"
    )
  where
    adoptTriggers :: SchemaName -> [SchemaName] -> [SchemaName] -> IO ()
adoptTriggers SchemaName
marker [SchemaName]
triggerNames [SchemaName]
relNames =
      Connection
-> Query
-> (SchemaName, SchemaName, PGArray SchemaName, PGArray SchemaName)
-> IO ()
forall params.
ToRow params =>
Connection -> Query -> params -> IO ()
executeRendered
        Connection
conn
        ( Query
commentTriggerSelect
            Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
triggerJoins
            Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND NOT trigger.tgisinternal AND trigger.tgname = ANY(?::text[]) \
               \AND relation.relname = ANY(?::text[]) AND obj_description(trigger.oid, 'pg_trigger') IS NULL"
        )
        (SchemaName
marker, SchemaName
schemaName, [SchemaName] -> PGArray SchemaName
forall a. [a] -> PGArray a
PGArray [SchemaName]
triggerNames, [SchemaName] -> PGArray SchemaName
forall a. [a] -> PGArray a
PGArray [SchemaName]
relNames)

    eventTriggerNames :: [SchemaName]
eventTriggerNames =
      ((SchemaName, Bool) -> SchemaName)
-> [(SchemaName, Bool)] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map (SchemaName, Bool) -> SchemaName
forall a b. (a, b) -> a
fst [(SchemaName, Bool)]
legacyEventStreamingTriggers
        [SchemaName] -> [SchemaName] -> [SchemaName]
forall a. Semigroup a => a -> a -> a
<> (SchemaName -> [SchemaName]) -> [SchemaName] -> [SchemaName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\SchemaName
table -> [SchemaName -> SchemaName
eventStreamingTriggerName SchemaName
table, SchemaName -> SchemaName
eventStreamingDLQTriggerName SchemaName
table]) [SchemaName]
tables

-- | Each queue's main and DLQ tables, which carry event-stream triggers.
eventTables :: [TableName] -> [TableName]
eventTables :: [SchemaName] -> [SchemaName]
eventTables = (SchemaName -> [SchemaName]) -> [SchemaName] -> [SchemaName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\SchemaName
table -> [SchemaName
table, SchemaName
table SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
dlqSuffix])

-- | Run one rendered DDL or DML statement.
executeSQL :: PG.Connection -> Text -> IO ()
executeSQL :: Connection -> SchemaName -> IO ()
executeSQL Connection
conn = IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ())
-> (SchemaName -> IO Int64) -> SchemaName -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Connection -> Query -> IO Int64
execute_ Connection
conn (Query -> IO Int64)
-> (SchemaName -> Query) -> SchemaName -> IO Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Query
Query (ByteString -> Query)
-> (SchemaName -> ByteString) -> SchemaName -> Query
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SchemaName -> ByteString
encodeUtf8

-- | The catalog joins the trigger sweeps share.
triggerJoins :: Query
triggerJoins :: Query
triggerJoins =
  Query
"FROM pg_trigger trigger \
  \JOIN pg_class relation ON relation.oid = trigger.tgrelid \
  \JOIN pg_namespace namespace ON namespace.oid = relation.relnamespace "

-- | The catalog joins the function sweeps share.
functionJoins :: Query
functionJoins :: Query
functionJoins = Query
"FROM pg_proc function JOIN pg_namespace namespace ON namespace.oid = function.pronamespace "

-- | A @COMMENT ON FUNCTION@ statement per matched row, stamping the first parameter.
commentFunctionSelect :: Query
commentFunctionSelect :: Query
commentFunctionSelect = Query
"SELECT format('COMMENT ON FUNCTION %I.%I() IS %L;', namespace.nspname, function.proname, ?::text) "

-- | Whether a single-row boolean probe returns true.
exists :: (PG.ToRow params) => PG.Connection -> Query -> params -> IO Bool
exists :: forall params.
ToRow params =>
Connection -> Query -> params -> IO Bool
exists Connection
conn Query
sql params
params = Bool -> (Only Bool -> Bool) -> Maybe (Only Bool) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Only Bool -> Bool
forall a. Only a -> a
fromOnly (Maybe (Only Bool) -> Bool)
-> ([Only Bool] -> Maybe (Only Bool)) -> [Only Bool] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Only Bool] -> Maybe (Only Bool)
forall a. [a] -> Maybe a
listToMaybe ([Only Bool] -> Bool) -> IO [Only Bool] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> Query -> params -> IO [Only Bool]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn Query
sql params
params

-- | Whether a trigger exists on a table and carries the given marker.
triggerIsCurrent :: PG.Connection -> SchemaName -> TableName -> Text -> Text -> IO Bool
triggerIsCurrent :: Connection
-> SchemaName -> SchemaName -> SchemaName -> SchemaName -> IO Bool
triggerIsCurrent Connection
conn SchemaName
schemaName SchemaName
table SchemaName
trigger SchemaName
marker =
  Connection
-> Query
-> (SchemaName, SchemaName, SchemaName, SchemaName)
-> IO Bool
forall params.
ToRow params =>
Connection -> Query -> params -> IO Bool
exists
    Connection
conn
    ( Query
"SELECT EXISTS (SELECT 1 "
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
triggerJoins
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND relation.relname = ? AND trigger.tgname = ? \
           \AND obj_description(trigger.oid, 'pg_trigger') = ?)"
    )
    (SchemaName
schemaName, SchemaName
table, SchemaName
trigger, SchemaName
marker)

-- | Drop the triggers whose marker starts with @marker@, outside the @keep@ tables.
dropTriggersMarked :: PG.Connection -> SchemaName -> Text -> [TableName] -> IO ()
dropTriggersMarked :: Connection -> SchemaName -> SchemaName -> [SchemaName] -> IO ()
dropTriggersMarked Connection
conn SchemaName
schemaName SchemaName
marker [SchemaName]
keep =
  Connection
-> Query -> (SchemaName, SchemaName, PGArray SchemaName) -> IO ()
forall params.
ToRow params =>
Connection -> Query -> params -> IO ()
executeRendered
    Connection
conn
    ( Query
dropTriggerSelect
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
triggerJoins
        Query -> Query -> Query
forall a. Semigroup a => a -> a -> a
<> Query
"WHERE namespace.nspname = ? AND obj_description(trigger.oid, 'pg_trigger') LIKE ? \
           \AND NOT (relation.relname::text = ANY(?::text[]))"
    )
    (SchemaName
schemaName, SchemaName
marker SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
"%", [SchemaName] -> PGArray SchemaName
forall a. [a] -> PGArray a
PGArray [SchemaName]
keep)

-- | Run each statement a rendering query returns.
executeRendered :: (PG.ToRow params) => PG.Connection -> Query -> params -> IO ()
executeRendered :: forall params.
ToRow params =>
Connection -> Query -> params -> IO ()
executeRendered Connection
conn Query
sql params
params = do
  commands <- Connection -> Query -> params -> IO [Only SchemaName]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn Query
sql params
params
  traverse_ (\(Only SchemaName
command) -> Connection -> SchemaName -> IO ()
executeSQL Connection
conn SchemaName
command) commands

-- | A @DROP TRIGGER@ statement per matched row.
dropTriggerSelect :: Query
dropTriggerSelect :: Query
dropTriggerSelect = Query
"SELECT format('DROP TRIGGER IF EXISTS %I ON %I.%I;', trigger.tgname, namespace.nspname, relation.relname) "

-- | A @COMMENT ON TRIGGER@ statement per matched row, stamping the first parameter.
commentTriggerSelect :: Query
commentTriggerSelect :: Query
commentTriggerSelect =
  Query
"SELECT format('COMMENT ON TRIGGER %I ON %I.%I IS %L;', trigger.tgname, namespace.nspname, relation.relname, ?::text) "

-- | The schema-level migrations, run once per schema. The optional notification and
-- event-streaming objects are reconciled after the migrations run.
schemaLevelMigrations :: SchemaName -> [MigrationCommand]
schemaLevelMigrations :: SchemaName -> [MigrationCommand]
schemaLevelMigrations SchemaName
schemaName =
  [ String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-cron-schedules" SchemaName -> SchemaName
createCronSchedulesTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"cron-schedules-add-timezone" SchemaName -> SchemaName
addTimezoneColumnSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"cron-schedules-add-queue-name" SchemaName -> SchemaName
addQueueNameColumnSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"cron-schedules-add-run-requested" SchemaName -> SchemaName
addRunRequestedColumnSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"cron-schedules-add-last-manual-run" SchemaName -> SchemaName
addLastManualRunColumnSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-workers" SchemaName -> SchemaName
createWorkersTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-queues" SchemaName -> SchemaName
createQueuesTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"arbiter-queues-add-dlq-retention" SchemaName -> SchemaName
addDLQRetentionColumnSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-gates" SchemaName -> SchemaName
createGatesTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"arbiter-gates-add-metadata" SchemaName -> SchemaName
addGateMetadataColumnSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-rate-limit-policies" SchemaName -> SchemaName
createRateLimitPoliciesTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-rate-limits" SchemaName -> SchemaName
createRateLimitsTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-concurrency-policies" SchemaName -> SchemaName
createConcurrencyPoliciesTableSQL
  , String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
"create-arbiter-concurrency" SchemaName -> SchemaName
createConcurrencyTableSQL
  ]
  where
    script :: String -> (SchemaName -> SchemaName) -> MigrationCommand
script String
name SchemaName -> SchemaName
sql = String -> ByteString -> MigrationCommand
MigrationScript String
name (SchemaName -> ByteString
encodeUtf8 (SchemaName -> SchemaName
sql SchemaName
schemaName))

-- | One queue's tracked migrations, each under its own version identifier.
-- The optional notify and event-streaming objects are reconciled after the migrations run.
jobQueueMigrationsForTable
  :: SchemaName
  -- ^ Schema for the arbiter tables
  -> TableName
  -- ^ Queue table name
  -> TableAdmission
  -- ^ Which admission trigger kinds to install
  -> [MigrationCommand]
jobQueueMigrationsForTable :: SchemaName -> SchemaName -> TableAdmission -> [MigrationCommand]
jobQueueMigrationsForTable SchemaName
schemaName SchemaName
tableName TableAdmission
admission =
  [MigrationCommand]
coreMigrations [MigrationCommand] -> [MigrationCommand] -> [MigrationCommand]
forall a. Semigroup a => a -> a -> a
<> [MigrationCommand]
concurrencyTriggers [MigrationCommand] -> [MigrationCommand] -> [MigrationCommand]
forall a. Semigroup a => a -> a -> a
<> [MigrationCommand]
rateLimitTriggers
  where
    prefix :: String
prefix = SchemaName -> String
T.unpack SchemaName
tableName String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"-"
    script :: String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
name SchemaName -> SchemaName -> SchemaName
sql = String -> ByteString -> MigrationCommand
MigrationScript (String
prefix String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
name) (SchemaName -> ByteString
encodeUtf8 (SchemaName -> SchemaName -> SchemaName
sql SchemaName
schemaName SchemaName
tableName))

    coreMigrations :: [MigrationCommand]
coreMigrations =
      [ String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-table" SchemaName -> SchemaName -> SchemaName
createJobQueueTableSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-dlq-table" SchemaName -> SchemaName -> SchemaName
createJobQueueDLQTableSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-dlq-group-key-index" SchemaName -> SchemaName -> SchemaName
createDLQGroupKeyIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-dlq-failed-at-index" SchemaName -> SchemaName -> SchemaName
createDLQFailedAtIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-dedup-key-index" SchemaName -> SchemaName -> SchemaName
createDedupKeyIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-group-key-index" SchemaName -> SchemaName -> SchemaName
createJobQueueGroupKeyIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-parent-id-index" SchemaName -> SchemaName -> SchemaName
createParentIdIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-dlq-parent-id-index" SchemaName -> SchemaName -> SchemaName
createDLQParentIdIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-results-table" SchemaName -> SchemaName -> SchemaName
createResultsTableSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-groups-table" SchemaName -> SchemaName -> SchemaName
createGroupsTableSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-claimed-by-column" SchemaName -> SchemaName -> SchemaName
addClaimedByColumnSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"migrate-ungrouped-ready-split-indexes" SchemaName -> SchemaName -> SchemaName
migrateUngroupedReadySplitIndexesSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"migrate-groups-ready-ranking" SchemaName -> SchemaName -> SchemaName
migrateGroupsReadyRankingSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-rate-limit-columns" SchemaName -> SchemaName -> SchemaName
addRateLimitColumnsSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-throttled-index" SchemaName -> SchemaName -> SchemaName
createThrottledIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-concurrency-columns" SchemaName -> SchemaName -> SchemaName
addConcurrencyColumnsSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-concurrency-index" SchemaName -> SchemaName -> SchemaName
createConcurrencyIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-rate-limit-cost-column" SchemaName -> SchemaName -> SchemaName
addRateLimitCostColumnSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"set-job-queue-fillfactor-100" SchemaName -> SchemaName -> SchemaName
setJobQueueFillfactorSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"set-max-attempts-default" SchemaName -> SchemaName -> SchemaName
setMaxAttemptsDefaultSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-cancel-requested-at-column" SchemaName -> SchemaName -> SchemaName
addCancelRequestedAtColumnSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-archive-for-column" SchemaName -> SchemaName -> SchemaName
addArchiveForColumnSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-archive-table" SchemaName -> SchemaName -> SchemaName
createJobQueueArchiveTableSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-archive-completed-at-index" SchemaName -> SchemaName -> SchemaName
createArchiveCompletedAtIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-archive-expires-at-index" SchemaName -> SchemaName -> SchemaName
createArchiveExpiresAtIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-archive-job-id-index" SchemaName -> SchemaName -> SchemaName
createArchiveJobIdIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-archive-parent-id-index" SchemaName -> SchemaName -> SchemaName
createArchiveParentIdIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-archive-group-key-index" SchemaName -> SchemaName -> SchemaName
createArchiveGroupKeyIndexSQL
      , -- Runs after the archive table exists. It alters that table too.
        String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-trace-context-column" SchemaName -> SchemaName -> SchemaName
addTraceContextColumnSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-claim-seq-column" SchemaName -> SchemaName -> SchemaName
addClaimSeqColumnSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-groups-emptied-index" SchemaName -> SchemaName -> SchemaName
createGroupsEmptiedIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-group-retried-index" SchemaName -> SchemaName -> SchemaName
createJobQueueGroupRetriedIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-grouped-due-index" SchemaName -> SchemaName -> SchemaName
createJobQueueGroupedDueIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-group-in-flight-index" SchemaName -> SchemaName -> SchemaName
createJobQueueGroupInFlightIndexSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-groups-trigger-functions-v12" SchemaName -> SchemaName -> SchemaName
createGroupsTriggerFunctionsSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-groups-triggers" SchemaName -> SchemaName -> SchemaName
createGroupsTriggersSQL
      , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"add-kind-column" SchemaName -> SchemaName -> SchemaName
addKindColumnSQL
      ]
    concurrencyTriggers :: [MigrationCommand]
concurrencyTriggers
      | TableAdmission -> Bool
tableConcurrency TableAdmission
admission =
          [ String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-concurrency-trigger-functions" SchemaName -> SchemaName -> SchemaName
createConcurrencyTriggerFunctionsSQL
          , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-concurrency-triggers" SchemaName -> SchemaName -> SchemaName
createConcurrencyTriggersSQL
          ]
      | Bool
otherwise = []
    rateLimitTriggers :: [MigrationCommand]
rateLimitTriggers
      | TableAdmission -> Bool
tableRateLimit TableAdmission
admission =
          [ String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-rate-limit-bucket-trigger-functions" SchemaName -> SchemaName -> SchemaName
createRateLimitBucketTriggerFunctionsSQL
          , String
-> (SchemaName -> SchemaName -> SchemaName) -> MigrationCommand
script String
"create-rate-limit-bucket-triggers" SchemaName -> SchemaName -> SchemaName
createRateLimitBucketTriggersSQL
          ]
      | Bool
otherwise = []

    setJobQueueFillfactorSQL :: SchemaName -> SchemaName -> SchemaName
setJobQueueFillfactorSQL SchemaName
schema SchemaName
table = SchemaName
"ALTER TABLE " SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName -> SchemaName -> SchemaName
jobQueueTable SchemaName
schema SchemaName
table SchemaName -> SchemaName -> SchemaName
forall a. Semigroup a => a -> a -> a
<> SchemaName
" SET (fillfactor = 100);"

-- | Whether a schema exists.
schemaExists :: PG.Connection -> SchemaName -> IO Bool
schemaExists :: Connection -> SchemaName -> IO Bool
schemaExists Connection
conn SchemaName
schemaName =
  Connection -> Query -> Only SchemaName -> IO Bool
forall params.
ToRow params =>
Connection -> Query -> params -> IO Bool
exists Connection
conn Query
"SELECT EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = ?)" (SchemaName -> Only SchemaName
forall a. a -> Only a
Only SchemaName
schemaName)