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

-- | Lifecycle SQL templates.
module Arbiter.Core.Sql.Lifecycle
  ( smartAckJobSQL
  , smartAckJobsBatchSQL
  , setVisibilityTimeoutSQL
  , setVisibilityTimeoutBatchSQL
  , updateJobForRetrySQL
  , nackJobSQL
  , nackJobsBatchSQL
  , promoteJobSQL
  ) where

import Data.Int (Int32, Int64)
import Data.Text (Text)

import Arbiter.Core.Job.Schema (jobQueueTable)
import Arbiter.Core.Sql.Archive (archiveAckCte)
import Arbiter.Core.Sql.QQ (sql)
import Arbiter.Core.Sql.Query (Query, mwhen)
import Arbiter.Core.Sql.Tree (lockedByIdsCte)

-- | Parent-aware ack. Deletes a childless job, suspends one whose children are still
-- running, and wakes a suspended parent whose last child left the queue. Returns 1,
-- or 0 for a job gone, reclaimed or cancelled. When @archiveEnabled@, the deleted row
-- is teed into the archive per-row on @archive_for@.
smartAckJobSQL :: Bool -> Text -> Text -> Int64 -> Int64 -> Query Int64
smartAckJobSQL :: Bool -> Text -> Text -> Int64 -> Int64 -> Query Int64
smartAckJobSQL Bool
archiveEnabled Text
schema Text
tableName Int64
jobId Int64
cseq =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
      returning :: Text
returning = if Bool
archiveEnabled then Text
"*" else Text
"id, parent_id" :: Text
      archived :: Text
archived = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
archiveEnabled (Text -> Text -> Text -> Text
archiveAckCte Text
schema Text
tableName Text
"ack")
   in [sql|
        WITH ack AS (
          DELETE FROM ${tbl}
          WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8}
            AND NOT EXISTS (SELECT 1 FROM ${tbl} WHERE parent_id = #{jobId :: CInt8})
          RETURNING ${returning}
        ),
        ${archived}
        suspend AS (
          UPDATE ${tbl}
          SET suspended = TRUE, not_visible_until = NULL, claimed_by = NULL, updated_at = NOW()
          WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8}
            AND NOT EXISTS (SELECT 1 FROM ack)
            AND EXISTS (SELECT 1 FROM ${tbl} WHERE parent_id = #{jobId :: CInt8})
          RETURNING id
        ),
        wake_parent AS (
          UPDATE ${tbl}
          SET suspended = FALSE, updated_at = NOW()
          WHERE id = (SELECT parent_id FROM ack WHERE parent_id IS NOT NULL)
            AND suspended = TRUE
            AND NOT EXISTS (
              SELECT 1 FROM ${tbl} child
              WHERE child.parent_id = (SELECT parent_id FROM ack WHERE parent_id IS NOT NULL)
                AND child.id NOT IN (SELECT id FROM ack)
            )
          RETURNING id
        )
        SELECT
          (SELECT count(*) FROM ack) + (SELECT count(*) FROM suspend) AS @{result :: CInt8}
      |]

-- | Set-based smart ack over @unnest@ed @(id, claim_seq)@ arrays. Deletes leaves,
-- suspends finalizers that still have children, and wakes parents whose last
-- child completed. The wake check excludes acked children explicitly. Returns the
-- acked ids. Reclaimed jobs are absent. Locks children-first to match nack and
-- force-cancel. The caller holds the parent locks.
smartAckJobsBatchSQL :: Bool -> Text -> Text -> [Int64] -> [Int64] -> Query Int64
smartAckJobsBatchSQL :: Bool -> Text -> Text -> [Int64] -> [Int64] -> Query Int64
smartAckJobsBatchSQL Bool
archiveEnabled Text
schema Text
tableName [Int64]
ids [Int64]
cseqs =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
      returning :: Text
returning = if Bool
archiveEnabled then Text
"job.*" else Text
"job.id, job.parent_id" :: Text
      archived :: Text
archived = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
archiveEnabled (Text -> Text -> Text -> Text
archiveAckCte Text
schema Text
tableName Text
"ack")
      locked :: Query ()
locked = Text -> [Int64] -> Query ()
lockedByIdsCte Text
tbl [Int64]
ids
   in [sql|
        WITH input AS (
          SELECT unnest(#{ids :: [CInt8]}::bigint[]) AS id, unnest(#{cseqs :: [CInt8]}::bigint[]) AS cseq
        ),
        ${locked},
        ack AS (
          DELETE FROM ${tbl} job
          USING input input_row
          WHERE job.id = input_row.id AND job.claim_seq = input_row.cseq
            AND job.id IN (SELECT id FROM locked)
            AND NOT EXISTS (SELECT 1 FROM ${tbl} child WHERE child.parent_id = job.id)
          RETURNING ${returning}
        ),
        ${archived}
        suspend AS (
          UPDATE ${tbl} job
          SET suspended = TRUE, not_visible_until = NULL, claimed_by = NULL, updated_at = NOW()
          FROM input input_row
          WHERE job.id = input_row.id AND job.claim_seq = input_row.cseq
            AND job.id IN (SELECT id FROM locked)
            AND NOT EXISTS (SELECT 1 FROM ack acked WHERE acked.id = job.id)
            AND EXISTS (SELECT 1 FROM ${tbl} child WHERE child.parent_id = job.id)
          RETURNING job.id
        ),
        wake_parent AS (
          UPDATE ${tbl} parent
          SET suspended = FALSE, updated_at = NOW()
          WHERE parent.id IN (SELECT DISTINCT parent_id FROM ack WHERE parent_id IS NOT NULL)
            AND parent.suspended = TRUE
            AND NOT EXISTS (
              SELECT 1 FROM ${tbl} child
              WHERE child.parent_id = parent.id
                AND NOT EXISTS (SELECT 1 FROM ack acked WHERE acked.id = child.id)
            )
          RETURNING parent.id
        )
        SELECT @{id :: CInt8} FROM ack
        UNION
        SELECT id FROM suspend
      |]

-- | Extend a job's visibility timeout. Matches on the claim token. Suspended rows
-- hold no lease.
setVisibilityTimeoutSQL :: Text -> Text -> Double -> Int64 -> Int64 -> Query ()
setVisibilityTimeoutSQL :: Text -> Text -> Double -> Int64 -> Int64 -> Query ()
setVisibilityTimeoutSQL Text
schema Text
tableName Double
secs Int64
jobId Int64
cseq =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = CASE WHEN #{secs :: CFloat8}::double precision <= 0
                                     THEN NULL
                                     ELSE NOW() + (#{secs :: CFloat8}::double precision * interval '1 second') END,
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8} AND NOT suspended
      |]

-- | 'setVisibilityTimeoutSQL' over a batch, for the heartbeat. Extends every job still
-- under this claim, held by the same worker and unsuspended, and reports per row whether
-- the update landed alongside its claim token, cancel flag and suspension. @valuesFrag@
-- carries the input @(id, claim_seq, claimed_by)@ rows.
setVisibilityTimeoutBatchSQL :: Text -> Text -> Query () -> [Int64] -> Double -> Query ()
setVisibilityTimeoutBatchSQL :: Text -> Text -> Query () -> [Int64] -> Double -> Query ()
setVisibilityTimeoutBatchSQL Text
schema Text
tableName Query ()
valuesFrag [Int64]
ids Double
secs =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
      locked :: Query ()
locked = Text -> [Int64] -> Query ()
lockedByIdsCte Text
tbl [Int64]
ids
   in [sql|
        WITH input_jobs AS (
          SELECT input_values.id::bigint AS id,
                 input_values.expected_claim_seq::bigint AS expected_claim_seq,
                 input_values.expected_claimed_by::uuid AS expected_claimed_by
          FROM (VALUES ${valuesFrag}) AS input_values(id, expected_claim_seq, expected_claimed_by)
        ),
        ${locked},
        updated AS (
          UPDATE ${tbl} job
          SET not_visible_until = CASE WHEN #{secs :: CFloat8}::double precision <= 0
                                       THEN NULL
                                       ELSE NOW() + (#{secs :: CFloat8}::double precision * interval '1 second') END,
              updated_at = NOW()
          FROM input_jobs input_job
          WHERE job.id = input_job.id AND job.id IN (SELECT id FROM locked)
            AND job.claim_seq = input_job.expected_claim_seq AND NOT job.suspended
            AND job.claimed_by IS NOT DISTINCT FROM input_job.expected_claimed_by
          RETURNING job.id
        )
        SELECT
          input_job.id,
          (heartbeated.id IS NOT NULL) as was_heartbeated,
          job.claim_seq as current_db_claim_seq,
          (job.cancel_requested_at IS NOT NULL) as cancel_requested,
          (job.suspended IS TRUE) as suspended,
          job.claimed_by as claimed_by
        FROM input_jobs input_job
        LEFT JOIN updated heartbeated ON heartbeated.id = input_job.id
        LEFT JOIN ${tbl} job ON job.id = input_job.id
      |]

-- | Park a failed job for its retry backoff. Matches on the claim token.
updateJobForRetrySQL :: Text -> Text -> Int64 -> Text -> Int64 -> Int64 -> Query ()
updateJobForRetrySQL :: Text -> Text -> Int64 -> Text -> Int64 -> Int64 -> Query ()
updateJobForRetrySQL Text
schema Text
tableName Int64
backoff Text
errorMsg Int64
jobId Int64
cseq =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = NOW() + (#{backoff :: CInt8} * interval '1 second'),
            last_error = #{errorMsg :: CText},
            updated_at = NOW(),
            claimed_by = NULL
        WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8} AND NOT suspended
      |]

-- | Soft nack. Releases the claim and hands back the attempt it consumed, recording no
-- failure. Leaves @not_visible_until@ as it stands.
nackJobSQL :: Text -> Text -> Int64 -> Int64 -> Int32 -> Query ()
nackJobSQL :: Text -> Text -> Int64 -> Int64 -> Int32 -> Query ()
nackJobSQL Text
schema Text
tableName Int64
jobId Int64
cseq Int32
att =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET attempts = LEAST(GREATEST(#{att :: CInt4} - 1, attempts - 1, 0), attempts),
            claimed_by = NULL,
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8} AND NOT suspended
          AND claimed_by IS NOT NULL
      |]

-- | 'nackJobSQL' over @unnest@ed @(id, claim_seq, attempts)@ arrays, locking
-- children-first to match ack and force-cancel. Returns the ids nacked.
nackJobsBatchSQL :: Text -> Text -> [Int64] -> [Int64] -> [Int32] -> Query Int64
nackJobsBatchSQL :: Text -> Text -> [Int64] -> [Int64] -> [Int32] -> Query Int64
nackJobsBatchSQL Text
schema Text
tableName [Int64]
ids [Int64]
cseqs [Int32]
atts =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
      locked :: Query ()
locked = Text -> [Int64] -> Query ()
lockedByIdsCte Text
tbl [Int64]
ids
   in [sql|
        WITH input AS (
          SELECT unnest(#{ids :: [CInt8]}::bigint[]) AS in_id,
                 unnest(#{cseqs :: [CInt8]}::bigint[]) AS cseq,
                 unnest(#{atts :: [CInt4]}::int[]) AS att
        ),
        ${locked}
        UPDATE ${tbl} job
        SET attempts = LEAST(GREATEST(input_row.att - 1, job.attempts - 1, 0), job.attempts),
            claimed_by = NULL,
            updated_at = NOW()
        FROM input input_row
        WHERE job.id = input_row.in_id AND job.id IN (SELECT id FROM locked)
          AND job.claim_seq = input_row.cseq AND NOT job.suspended
          AND job.claimed_by IS NOT NULL
        RETURNING @{id :: CInt8}
      |]

-- | Make a delayed or retrying job immediately visible. Refuses an in-flight job.
promoteJobSQL :: Text -> Text -> Int64 -> Query ()
promoteJobSQL :: Text -> Text -> Int64 -> Query ()
promoteJobSQL Text
schema Text
tableName Int64
jobId =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = NULL,
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8}
          AND NOT suspended
          AND not_visible_until IS NOT NULL
          AND not_visible_until > NOW()
          AND claimed_by IS NULL
      |]