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

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Groups SQL templates.
module Arbiter.Core.Sql.Groups
  ( groupsWindowSQL
  , emptiedWindowSQL
  , lockGroupsSQL
  , refreshGroupsSQL
  , insertMissingGroupsSQL
  , settleGroupLocksCte
  , lockGroupKeysSQL
  ) where

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

import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueGroupsTable, jobQueueTable)
import Arbiter.Core.Job.Schema.Groups (groupAggregates, inFlightPredicate)
import Arbiter.Core.Sql.QQ (sql)
import Arbiter.Core.Sql.Query (Query)
import Arbiter.Core.SqlLiterals (lockGroupKeyOnConflictSQL)

-- | A @group_locks@ CTE over the group rows of the @deleted@ and @updated@ CTEs, in key
-- order, taken only when an @updated@ CTE changed a row. The DELETE trigger and the
-- UPDATE trigger each lock in key order, but the DELETE trigger runs first.
settleGroupLocksCte :: SchemaName -> TableName -> Text -> [Text] -> Text
settleGroupLocksCte :: Text -> Text -> Text -> [Text] -> Text
settleGroupLocksCte Text
schemaName Text
tableName Text
deleted [Text]
updated =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      anyUpdated :: Text
anyUpdated = Text -> [Text] -> Text
T.intercalate Text
" OR " [Text
"EXISTS (SELECT 1 FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
cte Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")" | Text
cte <- [Text]
updated]
      touched :: Text
touched = Text -> [Text] -> Text
T.intercalate Text
" UNION " [Text
"SELECT group_key FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
cte | Text
cte <- Text
deleted Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: [Text]
updated]
   in [text|
        group_locks AS (
          SELECT 1 FROM ${groupsTbl} locked_group
          WHERE (${anyUpdated})
            AND locked_group.group_key IN (${touched})
          ORDER BY locked_group.group_key
          FOR UPDATE
        )
      |]

-- | Create the missing group summaries among @keys@ and lock them all, in key order.
lockGroupKeysSQL :: SchemaName -> TableName -> [Text] -> Query ()
lockGroupKeysSQL :: Text -> Text -> [Text] -> Query ()
lockGroupKeysSQL Text
schemaName Text
tableName [Text]
keys =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
   in [sql|INSERT INTO ${groupsTbl} (group_key) SELECT DISTINCT unnest(#{keys :: [CText]}::text[]) AS group_key ORDER BY group_key ${lockGroupKeyOnConflictSQL}|]

-- | At most @limit@ group keys past @cursor@, in the database's key order, unlocked.
-- Its last key is the caller's resume cursor.
groupsWindowSQL :: SchemaName -> TableName -> Int -> Maybe Text -> Query Text
groupsWindowSQL :: Text -> Text -> Int -> Maybe Text -> Query Text
groupsWindowSQL Text
schemaName Text
tableName Int
limit Maybe Text
cursor =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      lim :: Int64
lim = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit :: Int64
      after :: Query ()
after = (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Text
key -> [sql|WHERE group_key > #{key :: CText}|]) Maybe Text
cursor
   in [sql|SELECT @{group_key :: CText} FROM ${groupsTbl} ${after} ORDER BY group_key LIMIT #{lim :: CInt8}|]

-- | At most @limit@ keys past @cursor@ the maintenance triggers emptied in place, in the
-- database's key order. Its last key is the caller's resume cursor.
emptiedWindowSQL :: SchemaName -> TableName -> Int -> Maybe Text -> Query Text
emptiedWindowSQL :: Text -> Text -> Int -> Maybe Text -> Query Text
emptiedWindowSQL Text
schemaName Text
tableName Int
limit Maybe Text
cursor =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      lim :: Int64
lim = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit :: Int64
      after :: Query ()
after = (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Text
key -> [sql|AND group_key > #{key :: CText}|]) Maybe Text
cursor
   in [sql|
        SELECT @{group_key :: CText} FROM ${groupsTbl}
        WHERE job_count = 0 ${after}
        ORDER BY group_key LIMIT #{lim :: CInt8}
      |]

-- | @FOR UPDATE SKIP LOCKED@ over the window 'groupsWindowSQL' returned plus the keys
-- 'emptiedWindowSQL' returned, one ascending pass in the key order the maintenance
-- triggers use. Returns the keys this transaction holds.
lockGroupsSQL :: SchemaName -> TableName -> Maybe Text -> Maybe Text -> [Text] -> Query Text
lockGroupsSQL :: Text -> Text -> Maybe Text -> Maybe Text -> [Text] -> Query Text
lockGroupsSQL Text
schemaName Text
tableName Maybe Text
cursor Maybe Text
upper [Text]
emptied =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      after :: Query ()
after = (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Text
key -> [sql|AND group_key > #{key :: CText}|]) Maybe Text
cursor
      window :: Query ()
window =
        (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap
          ( \Text
upperKey ->
              [sql|SELECT group_key FROM ${groupsTbl} WHERE group_key <= #{upperKey :: CText} ${after} UNION |]
          )
          Maybe Text
upper
   in [sql|
        WITH targets AS (
          ${window}SELECT unnest(#{emptied :: [CText]}::text[]) AS group_key
        )
        SELECT @{group_key :: CText} FROM ${groupsTbl}
        WHERE group_key IN (SELECT group_key FROM targets)
        ORDER BY group_key
        FOR UPDATE SKIP LOCKED
      |]

-- | Summary columns, shared by the refresh and the insert.
summaryAggregates :: Text
summaryAggregates :: Text
summaryAggregates =
  let aggs :: Text
aggs = Text -> Text
groupAggregates Text
""
      inFlight :: Text
inFlight = Text -> Text
inFlightPredicate Text
""
   in [text|${aggs}, MAX(not_visible_until) FILTER (WHERE ${inFlight}) AS in_flight_until|]

-- | Recompute the groups table, scoped to the locked keys from 'lockGroupsSQL'.
-- Deletes the summary of a key with no jobs left. Returns the count of rows updated or
-- deleted. A separate statement whose snapshot post-dates the lock.
refreshGroupsSQL :: SchemaName -> TableName -> [Text] -> Query Int64
refreshGroupsSQL :: Text -> Text -> [Text] -> Query Int64
refreshGroupsSQL Text
schemaName Text
tableName [Text]
keys =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      aggs :: Text
aggs = Text
summaryAggregates
   in [sql|
        WITH params AS (
          SELECT unnest(#{keys :: [CText]}::text[]) AS group_key
        ),
        current AS (
          SELECT group_key, ${aggs}
          FROM ${tbl}
          WHERE group_key IN (SELECT group_key FROM params)
          GROUP BY group_key
        ),
        deleted AS (
          DELETE FROM ${groupsTbl} summary
          WHERE summary.group_key IN (SELECT group_key FROM params)
            AND NOT EXISTS (SELECT 1 FROM current fresh WHERE fresh.group_key = summary.group_key)
          RETURNING 1
        ),
        updated AS (
          UPDATE ${groupsTbl} summary
          SET min_priority = fresh.min_priority,
              min_id = fresh.min_id,
              job_count = fresh.job_count,
              ready_count = fresh.ready_count,
              next_due = fresh.next_due,
              in_flight_until = fresh.in_flight_until
          FROM current fresh
          WHERE summary.group_key = fresh.group_key
            AND (summary.min_priority <> fresh.min_priority OR summary.min_id <> fresh.min_id
                 OR summary.job_count <> fresh.job_count
                 OR summary.ready_count <> fresh.ready_count
                 OR summary.next_due IS DISTINCT FROM fresh.next_due
                 OR summary.in_flight_until IS DISTINCT FROM fresh.in_flight_until)
          RETURNING 1
        )
        SELECT (SELECT count(*) FROM deleted) + (SELECT count(*) FROM updated) AS @{rewritten :: CInt8}
      |]

-- | Insert summary rows for grouped jobs that have none, at most @limit@ keys past
-- @lower@ up to @upper@, in key order. Returns each key and whether its insert landed.
insertMissingGroupsSQL :: SchemaName -> TableName -> Int -> Maybe Text -> Maybe Text -> Query (Text, Bool)
insertMissingGroupsSQL :: Text
-> Text -> Int -> Maybe Text -> Maybe Text -> Query (Text, Bool)
insertMissingGroupsSQL Text
schemaName Text
tableName Int
limit Maybe Text
lower Maybe Text
upper =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      aggs :: Text
aggs = Text
summaryAggregates
      lim :: Int64
lim = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
limit :: Int64
      after :: Query ()
after = (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Text
key -> [sql|AND job.group_key > #{key :: CText}|]) Maybe Text
lower
      upTo :: Query ()
upTo = (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Text
key -> [sql|AND job.group_key <= #{key :: CText}|]) Maybe Text
upper
   in [sql|
        WITH missing_keys AS (
          SELECT DISTINCT job.group_key
          FROM ${tbl} job
          WHERE job.group_key IS NOT NULL
            ${after}
            ${upTo}
            AND NOT EXISTS (SELECT 1 FROM ${groupsTbl} summary WHERE summary.group_key = job.group_key)
          ORDER BY job.group_key
          LIMIT #{lim :: CInt8}
        ),
        missing AS (
          SELECT group_key, ${aggs}
          FROM ${tbl}
          WHERE group_key IN (SELECT group_key FROM missing_keys)
          GROUP BY group_key
        ),
        inserted AS (
          INSERT INTO ${groupsTbl} (group_key, min_priority, min_id, job_count, ready_count, next_due, in_flight_until)
          SELECT group_key, min_priority, min_id, job_count, ready_count, next_due, in_flight_until
          FROM missing
          ORDER BY group_key
          ON CONFLICT (group_key) DO NOTHING
          RETURNING group_key
        )
        SELECT missing_key.group_key AS @{group_key :: CText}, (landed.group_key IS NOT NULL) AS @{inserted :: CBool}
        FROM missing_keys missing_key
        LEFT JOIN inserted landed ON landed.group_key = missing_key.group_key
        ORDER BY missing_key.group_key
      |]