{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

-- | Conversion of declared 'Policy' values to upsertable rows, plus DDL for the
-- policies table, bucket table, and job columns. No database execution here.
module Arbiter.Core.RateLimit.Schema
  ( -- * Policy rows
    PolicyRow (..)
  , toPolicyRow

    -- * Table name helpers
  , arbiterRateLimitPoliciesTable
  , arbiterRateLimitPoliciesTableName
  , arbiterRateLimitsTable
  , arbiterRateLimitsTableName

    -- * DDL
  , createRateLimitPoliciesTableSQL
  , createRateLimitsTableSQL
  , alterRateLimitsDurabilitySQL
  , upsertPolicyRowSQL
  , addRateLimitColumnsSQL
  , addRateLimitCostColumnSQL
  , createThrottledIndexSQL
  , createRateLimitBucketTriggerFunctionsSQL
  , createRateLimitBucketTriggersSQL
  , bucketSeedInsert
  ) where

import Data.Text (Text)
import Data.Text qualified as T
import NeatInterpolation (text)

import Arbiter.Core.Admission (effectivePolicyCol, policyUpsertSQL)
import Arbiter.Core.Job.Schema
  ( SchemaName
  , TableName
  , jobQueueDLQTable
  , jobQueueTable
  , maintenanceFunctionNames
  , statementTriggerSQL
  )
import Arbiter.Core.RateLimit.Spec (Durability (..), Policy (..))
import Arbiter.Core.SqlLiterals (doubleLiteral, quoteIdentifier, textLiteral)

-- | A token-bucket policy as upsertable row fields.
data PolicyRow = PolicyRow
  { PolicyRow -> Text
prefixId :: Text
  , PolicyRow -> Double
maxTokens :: Double
  , PolicyRow -> Double
refillAmt :: Double
  , PolicyRow -> Double
interval :: Double
  }
  deriving stock (PolicyRow -> PolicyRow -> Bool
(PolicyRow -> PolicyRow -> Bool)
-> (PolicyRow -> PolicyRow -> Bool) -> Eq PolicyRow
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PolicyRow -> PolicyRow -> Bool
== :: PolicyRow -> PolicyRow -> Bool
$c/= :: PolicyRow -> PolicyRow -> Bool
/= :: PolicyRow -> PolicyRow -> Bool
Eq, Int -> PolicyRow -> ShowS
[PolicyRow] -> ShowS
PolicyRow -> String
(Int -> PolicyRow -> ShowS)
-> (PolicyRow -> String)
-> ([PolicyRow] -> ShowS)
-> Show PolicyRow
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PolicyRow -> ShowS
showsPrec :: Int -> PolicyRow -> ShowS
$cshow :: PolicyRow -> String
show :: PolicyRow -> String
$cshowList :: [PolicyRow] -> ShowS
showList :: [PolicyRow] -> ShowS
Show)

-- | A policy in its stored form.
toPolicyRow :: Policy -> PolicyRow
toPolicyRow :: Policy -> PolicyRow
toPolicyRow (Policy Text
prefix Double
burst Double
refill NominalDiffTime
period) =
  PolicyRow {prefixId :: Text
prefixId = Text
prefix, maxTokens :: Double
maxTokens = Double
burst, refillAmt :: Double
refillAmt = Double
refill, interval :: Double
interval = NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
period}

-- | Qualified name of the app-global policies table.
arbiterRateLimitPoliciesTable :: SchemaName -> Text
arbiterRateLimitPoliciesTable :: Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName =
  Text -> Text
quoteIdentifier Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
arbiterRateLimitPoliciesTableName

-- | Bare name of the policies table, for catalog lookups by relname.
arbiterRateLimitPoliciesTableName :: Text
arbiterRateLimitPoliciesTableName :: Text
arbiterRateLimitPoliciesTableName = Text
"arbiter_rate_limit_policies"

-- | Bare (unqualified) name of the bucket table, for catalog lookups by relname.
arbiterRateLimitsTableName :: Text
arbiterRateLimitsTableName :: Text
arbiterRateLimitsTableName = Text
"arbiter_rate_limits"

-- | Qualified name of the single bucket table. Its WAL durability is a table-level
-- property reconciled by the migration.
arbiterRateLimitsTable :: SchemaName -> Text
arbiterRateLimitsTable :: Text -> Text
arbiterRateLimitsTable Text
schemaName =
  Text -> Text
quoteIdentifier Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
arbiterRateLimitsTableName

-- | DDL for the policies table. @default_*@ is migration-owned. @override_*@ is
-- management-owned. The effective params are @COALESCE(override, default)@.
createRateLimitPoliciesTableSQL :: SchemaName -> Text
createRateLimitPoliciesTableSQL :: Text -> Text
createRateLimitPoliciesTableSQL Text
schemaName =
  [Text] -> Text
T.unlines
    [ Text
"CREATE TABLE IF NOT EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ("
    , Text
"  prefix_id TEXT PRIMARY KEY,"
    , Text
"  default_max_tokens DOUBLE PRECISION NOT NULL CHECK (default_max_tokens >= 0),"
    , Text
"  default_refill_amount DOUBLE PRECISION NOT NULL CHECK (default_refill_amount >= 0),"
    , Text
"  default_interval DOUBLE PRECISION NOT NULL CHECK (default_interval > 0),"
    , Text
"  override_max_tokens DOUBLE PRECISION CHECK (override_max_tokens >= 0),"
    , Text
"  override_refill_amount DOUBLE PRECISION CHECK (override_refill_amount >= 0),"
    , Text
"  override_interval DOUBLE PRECISION CHECK (override_interval > 0)"
    , Text
");"
    ]

-- | DDL for the bucket table, always created @UNLOGGED@.
createRateLimitsTableSQL :: SchemaName -> Text
createRateLimitsTableSQL :: Text -> Text
createRateLimitsTableSQL Text
schemaName =
  [Text] -> Text
T.unlines
    [ Text
"CREATE UNLOGGED TABLE IF NOT EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
arbiterRateLimitsTable Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ("
    , Text
"  rate_limit_key TEXT PRIMARY KEY,"
    , Text
"  policy_prefix TEXT NOT NULL,"
    , Text
"  tokens DOUBLE PRECISION NOT NULL,"
    , Text
"  last_refill TIMESTAMPTZ NOT NULL"
    , Text
") WITH (fillfactor = 80);"
    ]

-- | Converge the bucket table's WAL persistence to a durability. Rewrites the
-- table under @ACCESS EXCLUSIVE@. Callers issue it when the durability differs
-- from the current state.
alterRateLimitsDurabilitySQL :: Durability -> SchemaName -> Text
alterRateLimitsDurabilitySQL :: Durability -> Text -> Text
alterRateLimitsDurabilitySQL Durability
dur Text
schemaName =
  Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
arbiterRateLimitsTable Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
set Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
";"
  where
    set :: Text
set = case Durability
dur of
      Durability
Durable -> Text
" SET LOGGED"
      Durability
Unlogged -> Text
" SET UNLOGGED"

-- | Upsert a policy's @default_*@ params. Any operator @override_*@ is left untouched.
upsertPolicyRowSQL :: SchemaName -> PolicyRow -> Text
upsertPolicyRowSQL :: Text -> PolicyRow -> Text
upsertPolicyRowSQL Text
schemaName PolicyRow
row =
  Text -> Text -> [(Text, Text)] -> Text
policyUpsertSQL
    (Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName)
    (Text -> Text
textLiteral (PolicyRow -> Text
prefixId PolicyRow
row))
    [ (Text
"default_max_tokens", Double -> Text
doubleLiteral (PolicyRow -> Double
maxTokens PolicyRow
row))
    , (Text
"default_refill_amount", Double -> Text
doubleLiteral (PolicyRow -> Double
refillAmt PolicyRow
row))
    , (Text
"default_interval", Double -> Text
doubleLiteral (PolicyRow -> Double
interval PolicyRow
row))
    ]

-- | Migration adding the rate-limit columns to a queue's job and DLQ tables. All
-- nullable. @throttled_until@ (job table only) marks a throttle-deferred grouped
-- head in-flight. Its group stays stalled and spends no attempt.
addRateLimitColumnsSQL :: SchemaName -> TableName -> Text
addRateLimitColumnsSQL :: Text -> Text -> Text
addRateLimitColumnsSQL Text
schemaName Text
tableName =
  [Text] -> Text
T.unlines
    [ Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS rate_limit_key TEXT;"
    , Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS rate_limit_prefix TEXT;"
    , Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS throttled_until TIMESTAMPTZ;"
    , Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueDLQTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS rate_limit_key TEXT;"
    , Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueDLQTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS rate_limit_prefix TEXT;"
    ]

-- | Add the per-job token cost column to a queue's job and DLQ tables. Defaulted to
-- a unit cost. Existing rows backfill to it. The DLQ stores it and a retried job
-- retains its cost.
addRateLimitCostColumnSQL :: SchemaName -> TableName -> Text
addRateLimitCostColumnSQL :: Text -> Text -> Text
addRateLimitCostColumnSQL Text
schemaName Text
tableName =
  [Text] -> Text
T.unlines
    [ Text
"ALTER TABLE "
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS rate_limit_cost DOUBLE PRECISION NOT NULL DEFAULT 1;"
    , Text
"ALTER TABLE "
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueDLQTable Text
schemaName Text
tableName
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS rate_limit_cost DOUBLE PRECISION NOT NULL DEFAULT 1;"
    ]

-- | Statement-level triggers that ensure a full token bucket row exists for every
-- rate-limited job's key whose prefix has a policy. Tokens are spent at claim and
-- refill over time. There is no delete trigger. Key creation or a dedup-replace key
-- move seeds a row.
createRateLimitBucketTriggerFunctionsSQL :: SchemaName -> TableName -> Text
createRateLimitBucketTriggerFunctionsSQL :: Text -> Text -> Text
createRateLimitBucketTriggerFunctionsSQL Text
schemaName Text
tableName =
  let buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schemaName
      policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName
      baseName :: Text
baseName = Text
"ensure_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_rate_limit_buckets"
      (Text
funcInsert, Text
_funcDelete, Text
funcUpdate) = Text -> Text -> (Text, Text, Text)
maintenanceFunctionNames Text
schemaName Text
baseName
      dollarQuote :: Text
dollarQuote = Text
"$$"
   in [Text] -> Text
T.unlines
        [ Text -> Text -> Text -> Text -> Text
bucketInsertFunction Text
funcInsert Text
buckets Text
policies Text
dollarQuote
        , Text -> Text -> Text -> Text -> Text
bucketUpdateFunction Text
funcUpdate Text
buckets Text
policies Text
dollarQuote
        ]

-- | The seed INSERT shared by the trigger functions and the claim's @rl_seed@ CTE,
-- parameterized over the source rows (aliased @n@) and key filter.
bucketSeedInsert :: Text -> Text -> Text -> Text -> Text
bucketSeedInsert :: Text -> Text -> Text -> Text -> Text
bucketSeedInsert Text
buckets Text
policies Text
source Text
match =
  let effMax :: Text
effMax = Text -> Text -> Text
effectivePolicyCol Text
"p" Text
"max_tokens"
   in [text|
        INSERT INTO ${buckets} (rate_limit_key, policy_prefix, tokens, last_refill)
        SELECT n.rate_limit_key, MAX(n.rate_limit_prefix),
               MAX(${effMax}), NOW()
        FROM ${source}
        JOIN ${policies} p ON p.prefix_id = n.rate_limit_prefix
        WHERE ${match}
          AND NOT EXISTS (SELECT 1 FROM ${buckets} b WHERE b.rate_limit_key = n.rate_limit_key)
        GROUP BY n.rate_limit_key
        ORDER BY n.rate_limit_key
        ON CONFLICT (rate_limit_key) DO NOTHING
      |]

-- | Seed a full bucket for each fresh inserted key whose prefix has a policy. The NOT
-- EXISTS guard skips the insert for an already-present key.
bucketInsertFunction :: Text -> Text -> Text -> Text -> Text
bucketInsertFunction :: Text -> Text -> Text -> Text -> Text
bucketInsertFunction Text
funcName Text
buckets Text
policies Text
dollarQuote =
  let seed :: Text
seed = Text -> Text -> Text -> Text -> Text
bucketSeedInsert Text
buckets Text
policies Text
"new_table n" Text
"n.rate_limit_key IS NOT NULL"
   in [text|
    CREATE OR REPLACE FUNCTION ${funcName}()
    RETURNS TRIGGER AS ${dollarQuote}
    BEGIN
      IF NOT EXISTS (SELECT 1 FROM new_table WHERE rate_limit_key IS NOT NULL LIMIT 1) THEN
        RETURN NULL;
      END IF;

      ${seed};

      RETURN NULL;
    END;
    ${dollarQuote} LANGUAGE plpgsql;
  |]

-- | Seed a bucket for a key a dedup-replace moved onto. A claim or heartbeat
-- leaves the key unchanged and returns early.
bucketUpdateFunction :: Text -> Text -> Text -> Text -> Text
bucketUpdateFunction :: Text -> Text -> Text -> Text -> Text
bucketUpdateFunction Text
funcName Text
buckets Text
policies Text
dollarQuote =
  let seed :: Text
seed =
        Text -> Text -> Text -> Text -> Text
bucketSeedInsert
          Text
buckets
          Text
policies
          Text
"new_table n JOIN old_table o ON o.id = n.id"
          Text
"n.rate_limit_key IS NOT NULL AND n.rate_limit_key IS DISTINCT FROM o.rate_limit_key"
   in [text|
    CREATE OR REPLACE FUNCTION ${funcName}()
    RETURNS TRIGGER AS ${dollarQuote}
    BEGIN
      IF NOT EXISTS (
        SELECT 1 FROM new_table n JOIN old_table o ON o.id = n.id
        WHERE n.rate_limit_key IS NOT NULL AND n.rate_limit_key IS DISTINCT FROM o.rate_limit_key
        LIMIT 1
      ) THEN
        RETURN NULL;
      END IF;

      ${seed};

      RETURN NULL;
    END;
    ${dollarQuote} LANGUAGE plpgsql;
  |]

-- | The statement-level AFTER INSERT and AFTER UPDATE triggers backing
-- 'createRateLimitBucketTriggerFunctionsSQL'. There is no delete trigger.
createRateLimitBucketTriggersSQL :: SchemaName -> TableName -> Text
createRateLimitBucketTriggersSQL :: Text -> Text -> Text
createRateLimitBucketTriggersSQL Text
schemaName Text
tableName =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      baseName :: Text
baseName = Text
"ensure_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_rate_limit_buckets"
   in Text -> [Text] -> Text
T.intercalate
        Text
"\n\n"
        [ Text -> Text -> Text -> Text -> Text -> Text -> Text
statementTriggerSQL Text
schemaName Text
tbl Text
baseName Text
"_insert" Text
"INSERT" Text
"NEW TABLE AS new_table"
        , Text -> Text -> Text -> Text -> Text -> Text -> Text
statementTriggerSQL Text
schemaName Text
tbl Text
baseName Text
"_update" Text
"UPDATE" Text
"OLD TABLE AS old_table NEW TABLE AS new_table"
        ]
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n"

-- | Index backing the throttle wake and per-prefix count, in its own migration. The
-- prefix leads and @rate_limit_key@ trails.
createThrottledIndexSQL :: SchemaName -> TableName -> Text
createThrottledIndexSQL :: Text -> Text -> Text
createThrottledIndexSQL Text
schemaName Text
tableName =
  [Text] -> Text
T.unlines
    [ Text
"CREATE INDEX IF NOT EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
quoteIdentifier (Text
"idx_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_throttled")
    , Text
"ON " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" (rate_limit_prefix, rate_limit_key)"
    , Text
"WHERE throttled_until IS NOT NULL;"
    ]