{-# LANGUAGE OverloadedStrings #-}

-- | Job row edits from a side connection, as another worker or an operator makes them.
module Arbiter.Test.Worker.Rows
  ( updateJob
  , reclaimJob
  , flagCancelled
  , releaseRow
  , takeClaimHolder
  , holdRowLock
  , rowCount
  ) where

import Arbiter.Core.Job.Schema (jobQueueTable)
import Control.Concurrent (threadDelay)
import Control.Monad (void)
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.Text (Text)
import Database.PostgreSQL.Simple qualified as PG

import Arbiter.Test.Setup (execute_, toQuery, withConn)

-- | Run one UPDATE on a job row over a fresh connection.
updateJob :: ByteString -> Text -> Text -> Text -> Int64 -> IO ()
updateJob :: ByteString -> Text -> Text -> Text -> Int64 -> IO ()
updateJob ByteString
connStr Text
schema Text
table Text
setClause Int64
jobId = ByteString -> (Connection -> IO ()) -> IO ()
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO ()) -> IO ()) -> (Connection -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
  IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$
    Connection -> Query -> Only Int64 -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute
      Connection
conn
      (Text -> Query
toQuery (Text
"UPDATE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" SET " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
setClause Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE id = ?"))
      (Int64 -> Only Int64
forall a. a -> Only a
PG.Only Int64
jobId)

-- | Take the claim under a new token, as another worker's claim does.
reclaimJob :: ByteString -> Text -> Text -> Int64 -> IO ()
reclaimJob :: ByteString -> Text -> Text -> Int64 -> IO ()
reclaimJob ByteString
connStr Text
schema Text
table = ByteString -> Text -> Text -> Text -> Int64 -> IO ()
updateJob ByteString
connStr Text
schema Text
table Text
"attempts = attempts + 1, claim_seq = claim_seq + 1"

-- | Flag a job cancelled under its lease, as a force-cancel does, without the NOTIFY.
flagCancelled :: ByteString -> Text -> Text -> Int64 -> IO ()
flagCancelled :: ByteString -> Text -> Text -> Int64 -> IO ()
flagCancelled ByteString
connStr Text
schema Text
table = ByteString -> Text -> Text -> Text -> Int64 -> IO ()
updateJob ByteString
connStr Text
schema Text
table Text
"cancel_requested_at = NOW(), claim_seq = claim_seq + 1"

-- | Release a claim and make the row claimable now.
releaseRow :: ByteString -> Text -> Text -> Int64 -> IO ()
releaseRow :: ByteString -> Text -> Text -> Int64 -> IO ()
releaseRow ByteString
connStr Text
schema Text
table = ByteString -> Text -> Text -> Text -> Int64 -> IO ()
updateJob ByteString
connStr Text
schema Text
table Text
"claimed_by = NULL, not_visible_until = NULL"

-- | Take every open claim without bumping its token. The extend then reports 'Arbiter.Core.HighLevel.VisibilityUnchanged'.
takeClaimHolder :: ByteString -> Text -> Text -> IO ()
takeClaimHolder :: ByteString -> Text -> Text -> IO ()
takeClaimHolder ByteString
connStr Text
schema Text
table = ByteString -> (Connection -> IO ()) -> IO ()
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO ()) -> IO ()) -> (Connection -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
  Connection -> Text -> IO ()
execute_
    Connection
conn
    ( Text
"UPDATE "
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schema Text
table
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" SET claimed_by = '00000000-0000-0000-0000-000000000009'::uuid WHERE claimed_by IS NOT NULL"
    )

-- | Hold a row lock on one job for @micros@, as a transaction touching it would.
holdRowLock :: ByteString -> Text -> Text -> Int64 -> Int -> IO ()
holdRowLock :: ByteString -> Text -> Text -> Int64 -> Int -> IO ()
holdRowLock ByteString
connStr Text
schema Text
table Int64
jobId Int
micros = ByteString -> (Connection -> IO ()) -> IO ()
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO ()) -> IO ()) -> (Connection -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> do
  Connection -> IO ()
PG.begin Connection
conn
  _ <-
    Connection -> Query -> Only Int64 -> IO [Only Int64]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
PG.query
      Connection
conn
      (Text -> Query
toQuery (Text
"SELECT id FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE id = ? FOR UPDATE"))
      (Int64 -> Only Int64
forall a. a -> Only a
PG.Only Int64
jobId)
      :: IO [PG.Only Int64]
  threadDelay micros
  PG.commit conn

-- | How many rows carry the job id.
rowCount :: ByteString -> Text -> Text -> Int64 -> IO Int
rowCount :: ByteString -> Text -> Text -> Int64 -> IO Int
rowCount ByteString
connStr Text
schema Text
table Int64
jobId = ByteString -> (Connection -> IO Int) -> IO Int
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO Int) -> IO Int)
-> (Connection -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> do
  [PG.Only count] <-
    Connection -> Query -> Only Int64 -> IO [Only Int]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
PG.query
      Connection
conn
      (Text -> Query
toQuery (Text
"SELECT count(*)::int FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE id = ?"))
      (Int64 -> Only Int64
forall a. a -> Only a
PG.Only Int64
jobId)
  pure count