Schema.hs

{-# LANGUAGE ApplicativeDo #-}

{- | Hasql-level storage for the idempotent integration-event inbox.

This module owns the SQL surface that the inbox wrapper
('Keiro.Inbox.runInboxTransaction') and the optional Kafka decoder
('Keiro.Inbox.Kafka') consume.
-}
module Keiro.Inbox.Schema (
    tryInsertCompletedTx,
    markCompletedTx,
    markFailedTx,
    recordFailedAttemptTx,
    lookupInbox,
    listInbox,
    garbageCollectCompleted,
    countInboxBacklog,
)
where

import Contravariant.Extras (contrazip2, contrazip3, contrazip4)
import Data.ByteString (ByteString)
import Data.Functor.Contravariant ((>$<))
import Data.Time.Clock (NominalDiffTime, addUTCTime)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
import Data.UUID (UUID)
import Effectful (Eff, (:>))
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Keiro.Inbox.Types
import Keiro.Integration.Event (
    IntegrationEvent (..),
    SchemaReference (..),
    TraceContext (..),
    contentTypeText,
    parseContentType,
 )
import Keiro.Prelude
import Kiroku.Store.Effect (Store)
import Kiroku.Store.Transaction (runTransaction)
import Kiroku.Store.Types (EventId (..), GlobalPosition (..))
import "hasql-transaction" Hasql.Transaction qualified as Tx

{- | Attempt to insert a new completed row.

Returns 'Left' carrying the existing row if @(source, dedupe_key)@
already exists; returns 'Right ()' if the insert created a new row. The
caller — typically 'Keiro.Inbox.runInboxTransaction' — then either runs
the handler (new) or branches on the existing row's status (duplicate).

The row is inserted as @completed@ before the handler runs, but remains
uncommitted until the handler transaction succeeds. On handler failure the
whole transaction rolls back, so the completed row never becomes visible.

If a concurrent retention job deletes the conflicting row between the
insert attempt and the lookup, this returns 'Right ()'. The handler's effects
then commit without a replacement deduplication row, and a later redelivery
can run the handler again. This preserves at-least-once delivery but not
permanent exactly-once processing. Retention must exceed the maximum tolerated
redelivery delay, and handlers must remain idempotent.
-}
tryInsertCompletedTx ::
    InboxPersistence ->
    Text ->
    Text ->
    IntegrationEvent ->
    Maybe KafkaDeliveryRef ->
    UTCTime ->
    Tx.Transaction (Either InboxRow ())
tryInsertCompletedTx persistence src dedupe event kafka now = do
    inserted <-
        Tx.statement (toEncodedInsert persistence src dedupe event kafka now) tryInsertStmt
    if inserted
        then pure (Right ())
        else do
            existing <- Tx.statement (src, dedupe) selectByKeyStmt
            case existing of
                Just row -> pure (Left row)
                Nothing -> pure (Right ())

-- | Mark an inbox row completed inside the same transaction as the handler.
markCompletedTx :: Text -> Text -> UTCTime -> Tx.Transaction ()
markCompletedTx src dedupe now =
    Tx.statement (src, dedupe, now) markCompletedStmt

-- | Mark an inbox row failed inside the same transaction as the handler.
markFailedTx :: Text -> Text -> Text -> UTCTime -> Tx.Transaction ()
markFailedTx src dedupe errMsg now =
    Tx.statement (src, dedupe, errMsg, now) markFailedStmt

{- | Record one failed handler attempt for @(source, dedupe_key)@.

Creates a failed row when the handler transaction rolled back the initial
processing insert, or increments the existing failed row's attempt count.
Returns the new attempt count.
-}
recordFailedAttemptTx ::
    Text ->
    Text ->
    IntegrationEvent ->
    Maybe KafkaDeliveryRef ->
    Text ->
    UTCTime ->
    Tx.Transaction Int
recordFailedAttemptTx src dedupe event kafka errMsg now =
    Tx.statement (toEncodedFailedInsert src dedupe event kafka errMsg now) recordFailedAttemptStmt

-- | Read one inbox row.
lookupInbox :: (Store :> es) => Text -> Text -> Eff es (Maybe InboxRow)
lookupInbox src dedupe =
    runTransaction $
        Tx.statement (src, dedupe) selectByKeyStmt

-- | List inbox rows for a source, ordered by @received_at@. Test helper.
listInbox :: (Store :> es) => Text -> Eff es [InboxRow]
listInbox src =
    runTransaction $
        Tx.statement src listBySourceStmt

{- | Count inbox rows in a non-terminal state (backlog gauge source).

Backlog = rows still @processing@ (in flight) or @failed@ (awaiting a
retry decision). Completed rows are terminal and excluded.
-}
countInboxBacklog :: (Store :> es) => Eff es Int
countInboxBacklog =
    runTransaction (Tx.statement () countInboxBacklogStmt)

{- | Delete completed inbox rows older than @keepFor@ from @now@.

Returns the number of rows deleted. The retention window defines the
duplicate-detection window: a redelivery that arrives after retention
GC has run will be processed again, so the window must exceed the
maximum delivery delay tolerated by the operator. The user guide
recommends 30 days as a default. See 'tryInsertCompletedTx' for the related
concurrent-GC race and its at-least-once consequence.
-}
garbageCollectCompleted ::
    (Store :> es) =>
    NominalDiffTime ->
    UTCTime ->
    Eff es Int
garbageCollectCompleted keepFor now = do
    let cutoff = addUTCTime (negate keepFor) now
    result <-
        runTransaction $
            Tx.statement cutoff gcStmt
    pure (fromIntegral result)

-- ---------------------------------------------------------------------------
-- Encoder support
-- ---------------------------------------------------------------------------

data EncodedInsert = EncodedInsert
    { source :: !Text
    , dedupeKey :: !Text
    , messageId :: !(Maybe Text)
    , sourceEventId :: !(Maybe UUID)
    , sourceGlobalPosition :: !(Maybe Int64)
    , destination :: !(Maybe Text)
    , eventType :: !(Maybe Text)
    , schemaVersion :: !(Maybe Int64)
    , contentType :: !Text
    , schemaRegistry :: !(Maybe Text)
    , schemaSubject :: !(Maybe Text)
    , schemaVersionRef :: !(Maybe Int64)
    , schemaId :: !(Maybe Int64)
    , schemaFingerprint :: !(Maybe Text)
    , causationId :: !(Maybe UUID)
    , correlationId :: !(Maybe UUID)
    , traceparent :: !(Maybe Text)
    , tracestate :: !(Maybe Text)
    , kafkaTopic :: !(Maybe Text)
    , kafkaPartition :: !(Maybe Int64)
    , kafkaOffset :: !(Maybe Int64)
    , payloadBytes :: !ByteString
    , attributes :: !(Maybe Value)
    , occurredAt :: !(Maybe UTCTime)
    , receivedAt :: !UTCTime
    }
    deriving stock (Generic)

data EncodedFailedInsert = EncodedFailedInsert
    { insert :: !EncodedInsert
    , lastError :: !Text
    }
    deriving stock (Generic)

toEncodedInsert ::
    InboxPersistence ->
    Text ->
    Text ->
    IntegrationEvent ->
    Maybe KafkaDeliveryRef ->
    UTCTime ->
    EncodedInsert
toEncodedInsert persistence src dedupe event kafka now =
    let mref = event ^. #schemaReference
        mtrace = event ^. #traceContext
     in EncodedInsert
            { source = src
            , dedupeKey = dedupe
            , messageId = nullIfEmpty (event ^. #messageId)
            , sourceEventId = fmap unEventId (event ^. #sourceEventId)
            , sourceGlobalPosition = fmap unGlobalPosition (event ^. #sourceGlobalPosition)
            , destination = nullIfEmpty (event ^. #destination)
            , eventType = nullIfEmpty (event ^. #eventType)
            , schemaVersion = persistedSchema (Just (fromIntegral (event ^. #schemaVersion)))
            , contentType = contentTypeText (event ^. #contentType)
            , schemaRegistry = persistedSchema (mref >>= (^. #registry))
            , schemaSubject = persistedSchema (mref >>= (^. #subject))
            , schemaVersionRef = persistedSchema (fmap fromIntegral (mref >>= (^. #version)))
            , schemaId = persistedSchema (mref >>= (^. #schemaId))
            , schemaFingerprint = persistedSchema (mref >>= (^. #fingerprint))
            , causationId = fmap unEventId (event ^. #causationId)
            , correlationId = fmap unEventId (event ^. #correlationId)
            , traceparent = persistedEnvelope (fmap (^. #traceparent) mtrace)
            , tracestate = persistedEnvelope (mtrace >>= (^. #tracestate))
            , kafkaTopic = fmap (^. #topic) kafka
            , kafkaPartition = fmap (^. #partition) kafka
            , kafkaOffset = fmap (^. #offset) kafka
            , payloadBytes = case persistence of
                PersistFullEnvelope -> event ^. #payloadBytes
                PersistDedupeOnly -> mempty
            , attributes = persistedEnvelope (event ^. #attributes)
            , occurredAt = Just (event ^. #occurredAt)
            , receivedAt = now
            }
  where
    persistedEnvelope :: Maybe x -> Maybe x
    persistedEnvelope value = case persistence of
        PersistFullEnvelope -> value
        PersistDedupeOnly -> Nothing

    persistedSchema :: Maybe x -> Maybe x
    persistedSchema = persistedEnvelope

toEncodedFailedInsert ::
    Text ->
    Text ->
    IntegrationEvent ->
    Maybe KafkaDeliveryRef ->
    Text ->
    UTCTime ->
    EncodedFailedInsert
toEncodedFailedInsert src dedupe event kafka errMsg now =
    EncodedFailedInsert
        { insert = toEncodedInsert PersistFullEnvelope src dedupe event kafka now
        , lastError = errMsg
        }

nullIfEmpty :: Text -> Maybe Text
nullIfEmpty t = if t == mempty then Nothing else Just t

unEventId :: EventId -> UUID
unEventId (EventId u) = u

unGlobalPosition :: GlobalPosition -> Int64
unGlobalPosition (GlobalPosition i) = i

encodedInsertEncoder :: E.Params EncodedInsert
encodedInsertEncoder =
    mconcat
        [ view #source >$< E.param (E.nonNullable E.text)
        , view #dedupeKey >$< E.param (E.nonNullable E.text)
        , view #messageId >$< E.param (E.nullable E.text)
        , view #sourceEventId >$< E.param (E.nullable E.uuid)
        , view #sourceGlobalPosition >$< E.param (E.nullable E.int8)
        , view #destination >$< E.param (E.nullable E.text)
        , view #eventType >$< E.param (E.nullable E.text)
        , view #schemaVersion >$< E.param (E.nullable E.int8)
        , view #contentType >$< E.param (E.nonNullable E.text)
        , view #schemaRegistry >$< E.param (E.nullable E.text)
        , view #schemaSubject >$< E.param (E.nullable E.text)
        , view #schemaVersionRef >$< E.param (E.nullable E.int8)
        , view #schemaId >$< E.param (E.nullable E.int8)
        , view #schemaFingerprint >$< E.param (E.nullable E.text)
        , view #causationId >$< E.param (E.nullable E.uuid)
        , view #correlationId >$< E.param (E.nullable E.uuid)
        , view #traceparent >$< E.param (E.nullable E.text)
        , view #tracestate >$< E.param (E.nullable E.text)
        , view #kafkaTopic >$< E.param (E.nullable E.text)
        , view #kafkaPartition >$< E.param (E.nullable E.int8)
        , view #kafkaOffset >$< E.param (E.nullable E.int8)
        , view #payloadBytes >$< E.param (E.nonNullable E.bytea)
        , view #attributes >$< E.param (E.nullable E.jsonb)
        , view #occurredAt >$< E.param (E.nullable E.timestamptz)
        , view #receivedAt >$< E.param (E.nonNullable E.timestamptz)
        ]

encodedFailedInsertEncoder :: E.Params EncodedFailedInsert
encodedFailedInsertEncoder =
    (view #insert >$< encodedInsertEncoder)
        <> (view #lastError >$< E.param (E.nonNullable E.text))

-- ---------------------------------------------------------------------------
-- Statements
-- ---------------------------------------------------------------------------

tryInsertStmt :: Statement EncodedInsert Bool
tryInsertStmt =
    preparable
        """
        INSERT INTO keiro.keiro_inbox
          ( source
          , dedupe_key
          , message_id
          , source_event_id
          , source_global_position
          , destination
          , event_type
          , schema_version
          , content_type
          , schema_registry
          , schema_subject
          , schema_version_ref
          , schema_id
          , schema_fingerprint
          , causation_id
          , correlation_id
          , traceparent
          , tracestate
          , kafka_topic
          , kafka_partition
          , kafka_offset
          , payload_bytes
          , attributes
          , occurred_at
          , received_at
          , status
          , completed_at
          )
        VALUES
          ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, 'completed', $25)
        ON CONFLICT (source, dedupe_key) DO NOTHING
        RETURNING TRUE
        """
        encodedInsertEncoder
        (fmap (fromMaybe False) (D.rowMaybe (D.column (D.nonNullable D.bool))))

markCompletedStmt :: Statement (Text, Text, UTCTime) ()
markCompletedStmt =
    preparable
        """
        UPDATE keiro.keiro_inbox
        SET status = 'completed',
            completed_at = $3,
            last_error = NULL
        WHERE source = $1 AND dedupe_key = $2
        """
        ( contrazip3
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.timestamptz))
        )
        D.noResult

markFailedStmt :: Statement (Text, Text, Text, UTCTime) ()
markFailedStmt =
    preparable
        """
        UPDATE keiro.keiro_inbox
        SET status = 'failed',
            failed_at = $4,
            last_error = $3
        WHERE source = $1 AND dedupe_key = $2
        """
        ( contrazip4
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.timestamptz))
        )
        D.noResult

recordFailedAttemptStmt :: Statement EncodedFailedInsert Int
recordFailedAttemptStmt =
    preparable
        """
        INSERT INTO keiro.keiro_inbox
          ( source
          , dedupe_key
          , message_id
          , source_event_id
          , source_global_position
          , destination
          , event_type
          , schema_version
          , content_type
          , schema_registry
          , schema_subject
          , schema_version_ref
          , schema_id
          , schema_fingerprint
          , causation_id
          , correlation_id
          , traceparent
          , tracestate
          , kafka_topic
          , kafka_partition
          , kafka_offset
          , payload_bytes
          , attributes
          , occurred_at
          , received_at
          , status
          , attempt_count
          , failed_at
          , last_error
          )
        VALUES
          ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, 'failed', 1, $25, $26)
        ON CONFLICT (source, dedupe_key) DO UPDATE
        SET status = 'failed',
            attempt_count = keiro_inbox.attempt_count + 1,
            last_error = EXCLUDED.last_error,
            failed_at = EXCLUDED.failed_at
        RETURNING attempt_count
        """
        encodedFailedInsertEncoder
        (fmap fromIntegral (D.singleRow (D.column (D.nonNullable D.int8))))

selectByKeyStmt :: Statement (Text, Text) (Maybe InboxRow)
selectByKeyStmt =
    preparable
        (selectAllSql <> " WHERE source = $1 AND dedupe_key = $2")
        ( contrazip2
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.text))
        )
        (D.rowMaybe inboxRowDecoder)

listBySourceStmt :: Statement Text [InboxRow]
listBySourceStmt =
    preparable
        (selectAllSql <> " WHERE source = $1 ORDER BY received_at, dedupe_key")
        (E.param (E.nonNullable E.text))
        (D.rowList inboxRowDecoder)

countInboxBacklogStmt :: Statement () Int
countInboxBacklogStmt =
    preparable
        "SELECT COUNT(*)::bigint FROM keiro.keiro_inbox WHERE status IN ('processing', 'failed')"
        E.noParams
        (fmap fromIntegral (D.singleRow (D.column (D.nonNullable D.int8))))

gcStmt :: Statement UTCTime Int64
gcStmt =
    preparable
        """
        WITH deleted AS (
          DELETE FROM keiro.keiro_inbox
          WHERE status = 'completed' AND completed_at < $1
          RETURNING 1
        )
        SELECT COALESCE(COUNT(*), 0)::bigint FROM deleted
        """
        (E.param (E.nonNullable E.timestamptz))
        (D.singleRow (D.column (D.nonNullable D.int8)))

selectAllSql :: Text
selectAllSql =
    """
    SELECT source, dedupe_key, message_id, source_event_id, source_global_position,
           destination, event_type, schema_version, content_type, schema_registry,
           schema_subject, schema_version_ref, schema_id, schema_fingerprint,
           causation_id, correlation_id, traceparent, tracestate, kafka_topic,
           kafka_partition, kafka_offset, payload_bytes, attributes, occurred_at,
           status, attempt_count, received_at, completed_at, failed_at, last_error
    FROM keiro.keiro_inbox
    """

inboxRowDecoder :: D.Row InboxRow
inboxRowDecoder = fmap assembleInboxRow rawDecoder

data RawInbox = RawInbox
    { source :: !Text
    , dedupeKey :: !Text
    , messageId :: !(Maybe Text)
    , sourceEventId :: !(Maybe EventId)
    , sourceGlobalPosition :: !(Maybe GlobalPosition)
    , destination :: !(Maybe Text)
    , eventType :: !(Maybe Text)
    , schemaVersion :: !(Maybe Int)
    , contentType :: !Text
    , schemaRegistry :: !(Maybe Text)
    , schemaSubject :: !(Maybe Text)
    , schemaVersionRef :: !(Maybe Int)
    , schemaId :: !(Maybe Int64)
    , schemaFingerprint :: !(Maybe Text)
    , causationId :: !(Maybe EventId)
    , correlationId :: !(Maybe EventId)
    , traceparent :: !(Maybe Text)
    , tracestate :: !(Maybe Text)
    , kafkaTopic :: !(Maybe Text)
    , kafkaPartition :: !(Maybe Int64)
    , kafkaOffset :: !(Maybe Int64)
    , payloadBytes :: !ByteString
    , attributes :: !(Maybe Value)
    , occurredAt :: !(Maybe UTCTime)
    , status :: !InboxStatus
    , attemptCount :: !Int
    , receivedAt :: !UTCTime
    , completedAt :: !(Maybe UTCTime)
    , failedAt :: !(Maybe UTCTime)
    , lastError :: !(Maybe Text)
    }
    deriving stock (Generic)

rawDecoder :: D.Row RawInbox
rawDecoder =
    RawInbox
        <$> D.column (D.nonNullable D.text)
        <*> D.column (D.nonNullable D.text)
        <*> D.column (D.nullable D.text)
        <*> (fmap EventId <$> D.column (D.nullable D.uuid))
        <*> (fmap GlobalPosition <$> D.column (D.nullable D.int8))
        <*> D.column (D.nullable D.text)
        <*> D.column (D.nullable D.text)
        <*> (fmap fromIntegral <$> D.column (D.nullable D.int8))
        <*> D.column (D.nonNullable D.text)
        <*> D.column (D.nullable D.text)
        <*> D.column (D.nullable D.text)
        <*> (fmap fromIntegral <$> D.column (D.nullable D.int8))
        <*> D.column (D.nullable D.int8)
        <*> D.column (D.nullable D.text)
        <*> (fmap EventId <$> D.column (D.nullable D.uuid))
        <*> (fmap EventId <$> D.column (D.nullable D.uuid))
        <*> D.column (D.nullable D.text)
        <*> D.column (D.nullable D.text)
        <*> D.column (D.nullable D.text)
        <*> D.column (D.nullable D.int8)
        <*> D.column (D.nullable D.int8)
        <*> D.column (D.nonNullable D.bytea)
        <*> D.column (D.nullable D.jsonb)
        <*> D.column (D.nullable D.timestamptz)
        <*> D.column (D.nonNullable (D.refine parseInboxStatus D.text))
        <*> (fromIntegral <$> D.column (D.nonNullable D.int8))
        <*> D.column (D.nonNullable D.timestamptz)
        <*> D.column (D.nullable D.timestamptz)
        <*> D.column (D.nullable D.timestamptz)
        <*> D.column (D.nullable D.text)

assembleInboxRow :: RawInbox -> InboxRow
assembleInboxRow raw =
    let traceContext = case raw ^. #traceparent of
            Nothing -> Nothing
            Just tp -> Just (TraceContext tp (raw ^. #tracestate))
        schemaReference =
            case ( raw ^. #schemaRegistry
                 , raw ^. #schemaSubject
                 , raw ^. #schemaVersionRef
                 , raw ^. #schemaId
                 , raw ^. #schemaFingerprint
                 ) of
                (Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing
                _ ->
                    Just
                        ( SchemaReference
                            (raw ^. #schemaRegistry)
                            (raw ^. #schemaSubject)
                            (raw ^. #schemaVersionRef)
                            (raw ^. #schemaId)
                            (raw ^. #schemaFingerprint)
                        )
        kafka =
            case (raw ^. #kafkaTopic, raw ^. #kafkaPartition, raw ^. #kafkaOffset) of
                (Just t, Just p, Just o) -> Just (KafkaDeliveryRef t p o)
                _ -> Nothing
        event =
            IntegrationEvent
                { messageId = fromMaybe mempty (raw ^. #messageId)
                , source = raw ^. #source
                , destination = fromMaybe mempty (raw ^. #destination)
                , key = Nothing
                , -- @key@ is not part of the inbox primary key and is not
                  -- carried separately on the row — the same partition info
                  -- lives in @kafka@. Receivers that need it can re-derive
                  -- from the payload.
                  eventType = fromMaybe mempty (raw ^. #eventType)
                , schemaVersion = fromMaybe 0 (raw ^. #schemaVersion)
                , contentType = parseContentType (raw ^. #contentType)
                , schemaReference
                , sourceEventId = raw ^. #sourceEventId
                , sourceGlobalPosition = raw ^. #sourceGlobalPosition
                , payloadBytes = raw ^. #payloadBytes
                , occurredAt = fromMaybe defaultEpoch (raw ^. #occurredAt)
                , causationId = raw ^. #causationId
                , correlationId = raw ^. #correlationId
                , traceContext
                , attributes = raw ^. #attributes
                }
     in InboxRow
            { source = raw ^. #source
            , dedupeKey = raw ^. #dedupeKey
            , event
            , kafka
            , status = raw ^. #status
            , attemptCount = raw ^. #attemptCount
            , receivedAt = raw ^. #receivedAt
            , completedAt = raw ^. #completedAt
            , failedAt = raw ^. #failedAt
            , lastError = raw ^. #lastError
            }

{- | Sentinel used when an inbox row was written without an @occurred_at@.

Inbox rows are observability-shaped; consumers that care should
re-read the timestamp from the envelope payload directly.
-}
defaultEpoch :: UTCTime
defaultEpoch = posixSecondsToUTCTime 0