Schema.hs

-- | The @keiro_snapshots@ table: persistence for aggregate snapshots.
--
-- One row per stream holds the latest snapshot of its folded state as JSONB,
-- tagged with the 'stateCodecVersion', 'regfileShapeHash', and 'stateShapeHash'
-- that produced it. 'lookupSnapshot' fetches the newest row matching all three
-- discriminators (so incompatible snapshots are simply not found). Within one
-- discriminator tuple, 'writeSnapshotRow' keeps only the highest stream version,
-- so a late or out-of-order write cannot regress the snapshot.
--
-- A write with a /different/ discriminator deliberately replaces the row even
-- at a lower stream version. This lets a rolled-back deployment
-- reclaim the single snapshot slot instead of being locked out by a newer codec
-- forever. During a mixed-version deployment, however, writers with incompatible
-- codecs can thrash that row and each side will miss the other's snapshot. The
-- cost is repeated full replay, never incorrect state: the event log remains the
-- source of truth.
--
-- This module is the storage layer beneath "Keiro.Snapshot"; callers normally
-- go through 'Keiro.Snapshot.hydrateWithSnapshot' and
-- 'Keiro.Snapshot.writeSnapshot' rather than these statements directly.
module Keiro.Snapshot.Schema
  ( -- * Rows
    SnapshotRow (..),
    SnapshotWrite (..),

    -- * Storage
    lookupSnapshot,
    writeSnapshotRow,
  )
where

import Contravariant.Extras (contrazip4, contrazip6)
import Effectful (Eff, (:>))
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Keiro.Prelude
import Kiroku.Store.Effect (Store)
import Kiroku.Store.Transaction (runTransaction)
import Kiroku.Store.Types (StreamId (..), StreamVersion (..))
import "hasql-transaction" Hasql.Transaction qualified as Tx
import Prelude qualified

-- | A snapshot row as read back from @keiro_snapshots@: the stored 'state'
-- JSON, the 'streamVersion' it captures, the 'stateCodecVersion' and
-- 'regfileShapeHash' and 'stateShapeHash' that gate compatibility, and the create/update
-- timestamps.
data SnapshotRow = SnapshotRow
  { streamId :: !StreamId,
    streamVersion :: !StreamVersion,
    state :: !Value,
    stateCodecVersion :: !Int,
    regfileShapeHash :: !Text,
    stateShapeHash :: !Text,
    createdAt :: !UTCTime,
    updatedAt :: !UTCTime
  }
  deriving stock (Generic, Eq, Show)

-- | The fields needed to write a snapshot — 'SnapshotRow' minus the
-- database-managed timestamps.
data SnapshotWrite = SnapshotWrite
  { streamId :: !StreamId,
    streamVersion :: !StreamVersion,
    state :: !Value,
    stateCodecVersion :: !Int,
    regfileShapeHash :: !Text,
    stateShapeHash :: !Text
  }
  deriving stock (Generic, Eq, Show)

-- | Fetch the latest snapshot for a stream that matches all three codec
-- discriminators. Returns 'Nothing' when no compatible snapshot exists, so an
-- incompatible one is treated as absent.
lookupSnapshot ::
  (Store :> es) =>
  StreamId ->
  Int ->
  Text ->
  Text ->
  Eff es (Maybe SnapshotRow)
lookupSnapshot streamId version shapeHash stateShapeHash =
  runTransaction
    $ Tx.statement
      (streamIdToInt streamId, Prelude.fromIntegral version, shapeHash, stateShapeHash)
      lookupSnapshotStmt

-- | Upsert a snapshot row for its stream. For the same discriminator tuple,
-- the write only takes effect when its 'streamVersion' is at least the stored
-- one. Any incompatible discriminator replaces the row even at a lower version
-- so codec rollback can make progress; see the module header for the
-- mixed-deployment performance caveat.
writeSnapshotRow ::
  (Store :> es) =>
  SnapshotWrite ->
  Eff es ()
writeSnapshotRow snapshot =
  runTransaction
    $ Tx.statement (snapshotWriteParams snapshot) writeSnapshotStmt

lookupSnapshotStmt :: Statement (Int64, Int64, Text, Text) (Maybe SnapshotRow)
lookupSnapshotStmt =
  preparable
    """
    SELECT stream_id, stream_version, state, state_codec_version, regfile_shape_hash, state_shape_hash, created_at, updated_at
    FROM keiro.keiro_snapshots
    WHERE stream_id = $1
      AND state_codec_version = $2
      AND regfile_shape_hash = $3
      AND state_shape_hash = $4
    ORDER BY stream_version DESC
    LIMIT 1
    """
    ( contrazip4
        (E.param (E.nonNullable E.int8))
        (E.param (E.nonNullable E.int8))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
    )
    (D.rowMaybe snapshotRowDecoder)

writeSnapshotStmt :: Statement (Int64, Int64, Value, Int64, Text, Text) ()
writeSnapshotStmt =
  preparable
    """
    INSERT INTO keiro.keiro_snapshots
      (stream_id, stream_version, state, state_codec_version, regfile_shape_hash, state_shape_hash)
    VALUES
      ($1, $2, $3, $4, $5, $6)
    ON CONFLICT (stream_id) DO UPDATE
      SET stream_version = EXCLUDED.stream_version,
          state = EXCLUDED.state,
          state_codec_version = EXCLUDED.state_codec_version,
          regfile_shape_hash = EXCLUDED.regfile_shape_hash,
          state_shape_hash = EXCLUDED.state_shape_hash,
          updated_at = now()
      WHERE keiro_snapshots.stream_version <= EXCLUDED.stream_version
         OR keiro_snapshots.state_codec_version <> EXCLUDED.state_codec_version
         OR keiro_snapshots.regfile_shape_hash <> EXCLUDED.regfile_shape_hash
         OR keiro_snapshots.state_shape_hash <> EXCLUDED.state_shape_hash
    """
    ( contrazip6
        (E.param (E.nonNullable E.int8))
        (E.param (E.nonNullable E.int8))
        (E.param (E.nonNullable E.jsonb))
        (E.param (E.nonNullable E.int8))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
    )
    D.noResult

snapshotRowDecoder :: D.Row SnapshotRow
snapshotRowDecoder =
  SnapshotRow
    <$> (StreamId <$> D.column (D.nonNullable D.int8))
    <*> (StreamVersion <$> D.column (D.nonNullable D.int8))
    <*> D.column (D.nonNullable D.jsonb)
    <*> (Prelude.fromIntegral <$> D.column (D.nonNullable D.int8))
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nonNullable D.timestamptz)
    <*> D.column (D.nonNullable D.timestamptz)

streamIdToInt :: StreamId -> Int64
streamIdToInt (StreamId value) = value

streamVersionToInt :: StreamVersion -> Int64
streamVersionToInt (StreamVersion value) = value

snapshotWriteParams :: SnapshotWrite -> (Int64, Int64, Value, Int64, Text, Text)
snapshotWriteParams snapshot =
  ( streamIdToInt (snapshot ^. #streamId),
    streamVersionToInt (snapshot ^. #streamVersion),
    snapshot ^. #state,
    Prelude.fromIntegral (snapshot ^. #stateCodecVersion),
    snapshot ^. #regfileShapeHash,
    snapshot ^. #stateShapeHash
  )