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

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Completed-job archive SQL templates.
module Arbiter.Core.Sql.Archive
  ( archiveAckCte
  , updateArchiveResultSQL
  , updateArchiveResultsBatchSQL
  , purgeArchiveSQL
  , listArchiveFilteredSQL
  , countArchiveFilteredSQL
  , deleteArchiveJobsBatchSQL
  , reEnqueueFromArchiveSQL
  ) where

import Data.Aeson (Value)
import Data.Int (Int64)
import Data.Text (Text)
import Data.Time (UTCTime)
import NeatInterpolation (text)

import Arbiter.Core.Codec (archiveRowCodec, codecColumns, jobRowCodec, joinColumns)
import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueArchiveTable, jobQueueTable)
import Arbiter.Core.Job.Types (JobRead, Stored)
import Arbiter.Core.Sql.Insert (RowEdit (..), editJoin)
import Arbiter.Core.Sql.Jobs
  ( countWhereSQL
  , enqueuedAgainCols
  , enqueuedAgainValsEditing
  , jobColsExceptId
  , jobReadColumns
  )
import Arbiter.Core.Sql.QQ (sql)
import Arbiter.Core.Sql.Query (Query, rows)
import Arbiter.Core.SqlLiterals (intLiteral)

-- | The archive read columns, in codec order. The archive uses @job_id@ for the
-- main-table @id@.
allArchiveColumns :: Text
allArchiveColumns :: Text
allArchiveColumns = [Text] -> Text
joinColumns (RowCodec
  (Int64, UTCTime, JobRead (Stored (ZonkAny 0)), Maybe Value)
-> [Text]
forall a. RowCodec a -> [Text]
codecColumns (Text
-> RowCodec
     (Int64, UTCTime, JobRead (Stored (ZonkAny 0)), Maybe Value)
forall payload.
Text
-> RowCodec (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
archiveRowCodec Text
""))

-- | The @archived@ CTE teeing rows from the named @ack@ CTE into the archive, per-row
-- on @archive_for@. @archive_expires_at@ is precomputed.
-- The fragment ends with a comma, so another CTE must follow it.
archiveAckCte :: SchemaName -> TableName -> Text -> Text
archiveAckCte :: Text -> Text -> Text -> Text
archiveAckCte Text
schemaName Text
tableName Text
ackCte =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
   in [text|
        archived AS (
          INSERT INTO ${archiveTbl} (job_id, ${jobColsExceptId}, rate_limit_cost, completed_at, archive_expires_at)
          SELECT id, ${jobColsExceptId}, rate_limit_cost, NOW(), NOW() + (archive_for * interval '1 second')
          FROM ${ackCte}
          WHERE archive_for > 0
        ),
      |]

-- | Set a completed root job's stored @result@ on its archive row. No-ops without
-- an archive row.
updateArchiveResultSQL :: SchemaName -> TableName -> Value -> Int64 -> Query ()
updateArchiveResultSQL :: Text -> Text -> Value -> Int64 -> Query ()
updateArchiveResultSQL Text
schemaName Text
tableName Value
result Int64
jobId =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
   in [sql|UPDATE ${archiveTbl} SET result = #{result :: CJsonb} WHERE job_id = #{jobId :: CInt8}|]

-- | 'updateArchiveResultSQL' for several jobs in one statement.
updateArchiveResultsBatchSQL :: SchemaName -> TableName -> [Int64] -> [Value] -> Query ()
updateArchiveResultsBatchSQL :: Text -> Text -> [Int64] -> [Value] -> Query ()
updateArchiveResultsBatchSQL Text
schemaName Text
tableName [Int64]
jobIds [Value]
results =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
   in [sql|
        UPDATE ${archiveTbl} archive_row SET result = src.result
        FROM (
          SELECT unnest(#{jobIds :: [CInt8]}::bigint[]) AS job_id,
                 unnest(#{results :: [CJsonb]}::jsonb[]) AS result
        ) src
        WHERE archive_row.job_id = src.job_id
      |]

-- | Delete a bounded batch of archived jobs whose per-row @archive_expires_at@ has passed.
purgeArchiveSQL :: SchemaName -> TableName -> Int -> Query ()
purgeArchiveSQL :: Text -> Text -> Int -> Query ()
purgeArchiveSQL Text
schemaName Text
tableName Int
limit =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
      lim :: Text
lim = Int -> Text
forall a. Integral a => a -> Text
intLiteral Int
limit
   in [sql|
        DELETE FROM ${archiveTbl}
        WHERE ctid IN (
          SELECT ctid FROM ${archiveTbl}
          WHERE archive_expires_at < NOW()
          LIMIT ${lim}
        )
      |]

-- | List archived jobs under a dynamic WHERE.
listArchiveFilteredSQL
  :: SchemaName
  -> TableName
  -> Query ()
  -> Text
  -> Int64
  -> Int64
  -> Query (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
listArchiveFilteredSQL :: forall payload.
Text
-> Text
-> Query ()
-> Text
-> Int64
-> Int64
-> Query (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
listArchiveFilteredSQL Text
schemaName Text
tableName Query ()
whereFrag Text
orderBy Int64
limit Int64
offset =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
   in RowCodec (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
-> Query ()
-> Query (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
forall a. RowCodec a -> Query () -> Query a
rows
        (Text
-> RowCodec (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
forall payload.
Text
-> RowCodec (Int64, UTCTime, JobRead (Stored payload), Maybe Value)
archiveRowCodec Text
tableName)
        [sql|
          SELECT ${allArchiveColumns}
          FROM ${archiveTbl}
          ${whereFrag}
          ORDER BY ${orderBy}
          LIMIT #{limit :: CInt8} OFFSET #{offset :: CInt8}
        |]

-- | Count archived jobs under a dynamic WHERE.
countArchiveFilteredSQL :: SchemaName -> TableName -> Query () -> Query Int64
countArchiveFilteredSQL :: Text -> Text -> Query () -> Query Int64
countArchiveFilteredSQL Text
schemaName Text
tableName Query ()
whereFrag =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
   in Text -> Query () -> Query Int64
countWhereSQL Text
archiveTbl Query ()
whereFrag

-- | Delete archived jobs by archive primary key.
deleteArchiveJobsBatchSQL :: SchemaName -> TableName -> [Int64] -> Query ()
deleteArchiveJobsBatchSQL :: Text -> Text -> [Int64] -> Query ()
deleteArchiveJobsBatchSQL Text
schemaName Text
tableName [Int64]
archiveIds =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
   in [sql|DELETE FROM ${archiveTbl} WHERE id = ANY(#{archiveIds :: [CInt8]})|]

-- | Re-enqueue an archived job as a fresh standalone job, keeping the archive
-- row. Carries 'enqueuedAgainCols' and resets the other columns to their defaults.
-- An @edit@ replaces its columns.
reEnqueueFromArchiveSQL :: SchemaName -> TableName -> Int64 -> Maybe RowEdit -> Query (JobRead (Stored payload))
reEnqueueFromArchiveSQL :: forall payload.
Text
-> Text
-> Int64
-> Maybe RowEdit
-> Query (JobRead (Stored payload))
reEnqueueFromArchiveSQL Text
schemaName Text
tableName Int64
archiveId Maybe RowEdit
edit =
  let archiveTbl :: Text
archiveTbl = Text -> Text -> Text
jobQueueArchiveTable Text
schemaName Text
tableName
      tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      joinedEdit :: Query ()
joinedEdit = Maybe RowEdit -> Query ()
editJoin Maybe RowEdit
edit
      vals :: Text
vals = [Text] -> (Text -> Text) -> Text
enqueuedAgainValsEditing ((RowEdit -> [Text]) -> Maybe RowEdit -> [Text]
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap RowEdit -> [Text]
editColumns Maybe RowEdit
edit) (Text
"edit." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>)
   in RowCodec (JobRead (Stored payload))
-> Query () -> Query (JobRead (Stored payload))
forall a. RowCodec a -> Query () -> Query a
rows
        (Text -> RowCodec (JobRead (Stored payload))
forall payload. Text -> RowCodec (JobRead (Stored payload))
jobRowCodec Text
tableName)
        [sql|
          INSERT INTO ${tbl} (${enqueuedAgainCols})
          SELECT ${vals}
          FROM ${archiveTbl} ${joinedEdit}
          WHERE id = #{archiveId :: CInt8}
          RETURNING ${jobReadColumns}
        |]