-- | Reading, combining, and storing the results of a rollup job's children.
module Arbiter.Worker.Results
  ( childResults
  , mergedChildResults
  , mergeChildResults
  , storeJobResult
  , storeEncodedResult
  , storeEncodedResults
  ) where

import Arbiter.Core.Job.Types (JobRead, parentId, primaryKey, queueName)
import Arbiter.Core.JobResult (EncodeJobResult, decodeJobResult, encodeJobResult)
import Arbiter.Core.MonadArbiter (MonadArbiter, ResultOf, getSchema)
import Arbiter.Core.Operations qualified as Ops
import Control.Monad (void)
import Data.Aeson (FromJSON, Value)
import Data.Either (partitionEithers)
import Data.Foldable (fold, foldMap', for_)
import Data.Int (Int64)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (mapMaybe)
import Data.Text (Text)

-- | A rollup parent's immediate child results, keyed by child id, and its DLQ
-- errors, keyed by DLQ row id for 'Arbiter.Core.HighLevel.retryFromDLQ'. A 'Left'
-- is a DLQ child's error or a decode failure.
childResults
  :: (FromJSON (ResultOf m payload), MonadArbiter m)
  => JobRead payload
  -> m (Map Int64 (Either Text (ResultOf m payload)), Map Int64 Text)
childResults :: forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m) =>
JobRead payload
-> m (Map Int64 (Either Text (ResultOf m payload)), Map Int64 Text)
childResults JobRead payload
job = do
  schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
  (results, failures, snapshot, dlqFailures) <-
    Ops.readChildResultsRaw schema (queueName job) (primaryKey job)
  let raw = Map Int64 Value
-> Map Int64 Text -> Maybe Value -> Map Int64 (Either Text Value)
Ops.mergeRawChildResults Map Int64 Value
results Map Int64 Text
failures Maybe Value
snapshot
  pure (Map.map (>>= decodeJobResult) raw, dlqFailures)

-- | 'childResults' with successfully decoded values combined through 'Monoid'.
mergedChildResults
  :: ( FromJSON (ResultOf m payload)
     , MonadArbiter m
     , Monoid (ResultOf m payload)
     )
  => JobRead payload
  -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults :: forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job = do
  (results, dlqFailures) <- JobRead payload
-> m (Map Int64 (Either Text (ResultOf m payload)), Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m) =>
JobRead payload
-> m (Map Int64 (Either Text (ResultOf m payload)), Map Int64 Text)
childResults JobRead payload
job
  pure (mergeChildResults results, dlqFailures)

-- | Combine successful child results. A 'Left' counts as 'mempty'.
mergeChildResults :: (Monoid a) => Map Int64 (Either Text a) -> a
mergeChildResults :: forall a. Monoid a => Map Int64 (Either Text a) -> a
mergeChildResults = (Either Text a -> a) -> Map Int64 (Either Text a) -> a
forall m a. Monoid m => (a -> m) -> Map Int64 a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' Either Text a -> a
forall m. Monoid m => Either Text m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold

-- | Store a job's result for its parent rollup. A root job's result goes on its
-- archive entry only when its 'Arbiter.Core.Job.Types.archiveFor' is positive.
-- The archive entry exists only after the job's ack.
storeJobResult
  :: (EncodeJobResult result, MonadArbiter m)
  => JobRead payload
  -> result
  -> m ()
storeJobResult :: forall result (m :: * -> *) payload.
(EncodeJobResult result, MonadArbiter m) =>
JobRead payload -> result -> m ()
storeJobResult JobRead payload
job = JobRead payload -> Maybe Value -> m ()
forall (m :: * -> *) payload.
MonadArbiter m =>
JobRead payload -> Maybe Value -> m ()
storeEncodedResult JobRead payload
job (Maybe Value -> m ()) -> (result -> Maybe Value) -> result -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. result -> Maybe Value
forall a. EncodeJobResult a => a -> Maybe Value
encodeJobResult

-- | 'storeJobResult' on an already-encoded result. 'Nothing' stores nothing.
storeEncodedResult
  :: (MonadArbiter m)
  => JobRead payload
  -> Maybe Value
  -> m ()
storeEncodedResult :: forall (m :: * -> *) payload.
MonadArbiter m =>
JobRead payload -> Maybe Value -> m ()
storeEncodedResult JobRead payload
job Maybe Value
mVal = do
  schemaName <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
  for_ (resultRow (job, mVal)) $
    either
      (\(Int64
pid, Int64
cid, Value
val) -> m Int64 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int64 -> m ()) -> m Int64 -> m ()
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> Int64 -> Value -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> Int64 -> Value -> m Int64
Ops.insertResult Text
schemaName (JobRead payload -> Text
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> q
queueName JobRead payload
job) Int64
pid Int64
cid Value
val)
      (\(Int64
cid, Value
val) -> m Int64 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int64 -> m ()) -> m Int64 -> m ()
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> Value -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> Value -> m Int64
Ops.updateArchiveResult Text
schemaName (JobRead payload -> Text
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> q
queueName JobRead payload
job) Int64
cid Value
val)

-- | 'storeEncodedResult' over a batch from one queue. One statement stores the
-- child results and one the archived roots.
storeEncodedResults
  :: (MonadArbiter m)
  => [(JobRead payload, Maybe Value)]
  -> m ()
storeEncodedResults :: forall (m :: * -> *) payload.
MonadArbiter m =>
[(JobRead payload, Maybe Value)] -> m ()
storeEncodedResults [] = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
storeEncodedResults pairs :: [(JobRead payload, Maybe Value)]
pairs@((JobRead payload
firstJob, Maybe Value
_) : [(JobRead payload, Maybe Value)]
_) = do
  schemaName <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
  let (childRows, rootRows) = partitionEithers (mapMaybe resultRow pairs)
      queue = JobRead payload -> Text
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> q
queueName JobRead payload
firstJob
  void $ Ops.insertResultsBatch schemaName queue childRows
  void $ Ops.updateArchiveResultsBatch schemaName queue rootRows

-- | Where a result goes: the results table for a child, the archive entry for an
-- archived root, nowhere otherwise.
resultRow :: (JobRead payload, Maybe Value) -> Maybe (Either (Int64, Int64, Value) (Int64, Value))
resultRow :: forall payload.
(JobRead payload, Maybe Value)
-> Maybe (Either (Int64, Int64, Value) (Int64, Value))
resultRow (JobRead payload
job, Maybe Value
mVal) = do
  val <- Maybe Value
mVal
  case parentId job of
    Just Int64
pid -> Either (Int64, Int64, Value) (Int64, Value)
-> Maybe (Either (Int64, Int64, Value) (Int64, Value))
forall a. a -> Maybe a
Just ((Int64, Int64, Value)
-> Either (Int64, Int64, Value) (Int64, Value)
forall a b. a -> Either a b
Left (Int64
pid, JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job, Value
val))
    Maybe Int64
Nothing
      | JobRead payload -> Bool
forall payload. JobRead payload -> Bool
Ops.archivesOnAck JobRead payload
job -> Either (Int64, Int64, Value) (Int64, Value)
-> Maybe (Either (Int64, Int64, Value) (Int64, Value))
forall a. a -> Maybe a
Just ((Int64, Value) -> Either (Int64, Int64, Value) (Int64, Value)
forall a b. b -> Either a b
Right (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job, Value
val))
      | Bool
otherwise -> Maybe (Either (Int64, Int64, Value) (Int64, Value))
forall a. Maybe a
Nothing