{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- DDL and SQL fragments for the rate-limit feature. No database execution here.
module Arbiter.Core.RateLimit.Schema
  ( -- * Table names
    arbiterRateLimitPoliciesTable
  , arbiterRateLimitPoliciesTableName
  , arbiterRateLimitsTable
  , arbiterRateLimitsTableName

    -- * SQL fragments
  , bucketSeedInsert

    -- * DDL
  , createRateLimitPoliciesTableSQL
  , createRateLimitsTableSQL
  , alterRateLimitsDurabilitySQL
  , upsertRateLimitPolicyRowSQL
  , addRateLimitColumnsSQL
  , addRateLimitCostColumnSQL
  , createThrottledIndexSQL
  , createRateLimitBucketTriggerFunctionsSQL
  , createRateLimitBucketTriggersSQL
  ) 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
  , indexSQL
  , jobQueueDLQTable
  , jobQueueTable
  , maintenanceFunctionNames
  , statementTriggerSQL
  )
import Arbiter.Core.RateLimit.Spec (Durability (..), RateLimitPolicy (..))
import Arbiter.Core.SqlLiterals (dollarQuote, doubleLiteral, sharedTable, textLiteral)

-- | Qualified name of the app-global policies table.
arbiterRateLimitPoliciesTable :: SchemaName -> Text
arbiterRateLimitPoliciesTable :: Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName = Text -> Text -> Text
sharedTable Text
schemaName Text
arbiterRateLimitPoliciesTableName

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

-- | Bare 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.
arbiterRateLimitsTable :: SchemaName -> Text
arbiterRateLimitsTable :: Text -> Text
arbiterRateLimitsTable Text
schemaName = Text -> Text -> Text
sharedTable Text
schemaName 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);"
    ]

-- | Set the bucket table's WAL persistence. Rewrites the table under @ACCESS EXCLUSIVE@.
alterRateLimitsDurabilitySQL :: SchemaName -> Durability -> Text
alterRateLimitsDurabilitySQL :: Text -> Durability -> Text
alterRateLimitsDurabilitySQL Text
schemaName Durability
dur =
  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.
upsertRateLimitPolicyRowSQL :: SchemaName -> RateLimitPolicy -> Text
upsertRateLimitPolicyRowSQL :: Text -> RateLimitPolicy -> Text
upsertRateLimitPolicyRowSQL Text
schemaName RateLimitPolicy
policy =
  Text -> Text -> [(Text, Text)] -> Text
policyUpsertSQL
    (Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName)
    (Text -> Text
textLiteral (RateLimitPolicy -> Text
policyPrefix RateLimitPolicy
policy))
    [ (Text
"default_max_tokens", Double -> Text
doubleLiteral (RateLimitPolicy -> Double
policyMax RateLimitPolicy
policy))
    , (Text
"default_refill_amount", Double -> Text
doubleLiteral (RateLimitPolicy -> Double
policyRefill RateLimitPolicy
policy))
    , (Text
"default_interval", Double -> Text
doubleLiteral (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (RateLimitPolicy -> NominalDiffTime
policyInterval RateLimitPolicy
policy)))
    ]

-- | 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 job. A
-- deferred grouped job holds its group's in-flight slot. A deferral 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;"
    ]

-- | Per-queue trigger functions that seed a full token bucket row for each
-- rate-limited job's key whose prefix has a policy. Tokens are spent at claim and
-- refill over time. There is no delete function. 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 -> Text
bucketTriggerBase Text
tableName
      (Text
funcInsert, Text
_funcDelete, Text
funcUpdate) = Text -> Text -> (Text, Text, Text)
maintenanceFunctionNames Text
schemaName Text
baseName
   in [Text] -> Text
T.unlines
        [ Text -> Text -> Text -> Text
bucketInsertFunction Text
funcInsert Text
buckets Text
policies
        , Text -> Text -> Text -> Text
bucketUpdateFunction Text
funcUpdate Text
buckets Text
policies
        ]

-- | 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
bucketInsertFunction :: Text -> Text -> Text -> Text
bucketInsertFunction Text
funcName Text
buckets Text
policies =
  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
bucketUpdateFunction :: Text -> Text -> Text -> Text
bucketUpdateFunction Text
funcName Text
buckets Text
policies =
  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 -> Text
bucketTriggerBase Text
tableName
   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. The
-- prefix leads and @rate_limit_key@ trails.
createThrottledIndexSQL :: SchemaName -> TableName -> Text
createThrottledIndexSQL :: Text -> Text -> Text
createThrottledIndexSQL Text
schemaName Text
tableName =
  Text -> Text -> Text -> Maybe Text -> Text
indexSQL
    (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 -> Text -> Text
jobQueueTable Text
schemaName Text
tableName)
    Text
"rate_limit_prefix, rate_limit_key"
    (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"throttled_until IS NOT NULL")

-- | Base name of a queue's bucket-seeding trigger functions and triggers.
bucketTriggerBase :: TableName -> Text
bucketTriggerBase :: Text -> Text
bucketTriggerBase Text
tableName = 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"