ConsumerGroupSql.hs

{-# LANGUAGE OverloadedStrings #-}

-- | SQL-level tests for consumer-group partition routing and per-member
-- checkpoints (ExecPlan 28 / EP-1). These exercise the new prepared statements in
-- "Kiroku.Store.SQL" directly through the connection pool, with no subscription
-- runtime, on a fresh ephemeral PostgreSQL per test.
--
-- Terms: a /consumer group/ of /size/ N has members 0..N-1; each source stream is
-- assigned to exactly one member by 'Kiroku.Store.SQL'-encoded hash routing. The
-- properties proven here — disjointness, completeness, per-stream affinity,
-- determinism, and size-1 equivalence — are the contract EP-2's runtime depends on.
module Test.ConsumerGroupSql (spec) where

import Contravariant.Extras (contrazip2)
import Control.Lens ((^.))
import Data.Aeson qualified as Aeson
import Data.Generics.Labels ()
import Data.Int (Int32, Int64)
import Data.List (sort)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import Data.UUID (UUID)
import Data.Vector (Vector)
import Data.Vector qualified as V
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import Hasql.Statement (Statement, preparable)
import Kiroku.Store
import Kiroku.Store.SQL qualified as SQL
import Test.Helpers (makeEvent, withTestStore)
import Test.Hspec

-- | A generous limit that returns every seeded event in one read.
bigLimit :: Int32
bigLimit = 100000

-- | Run a statement against the store's pool, failing the test on a usage error.
runStmt :: KirokuStore -> Session.Session a -> IO a
runStmt store session = do
  result <- Pool.use (store ^. #pool) session
  case result of
    Left err -> error ("ConsumerGroupSql statement failed: " <> show err)
    Right a -> pure a

-- | Append one event to each of the named streams, seeding the store. Each
-- stream gets a single event; per-stream affinity is what we test, so one event
-- per stream is enough, but the helper appends 'n' events to vary versions.
seedStreams :: KirokuStore -> [Text] -> Int -> IO ()
seedStreams store names n =
  mapM_
    ( \name -> do
        let events = map (\i -> makeEvent ("E" <> T.pack (show i)) (Aeson.object [])) [1 .. n]
        r <- runStoreIO store $ appendToStream (StreamName name) NoStream events
        case r of
          Left err -> error ("seed append failed for " <> T.unpack name <> ": " <> show err)
          Right _ -> pure ()
    )
    names

-- | Collect the event ids from a read result, in order.
eventIds :: Vector RecordedEvent -> [UUID]
eventIds = map (\e -> case e ^. #eventId of EventId u -> u) . V.toList

-- | Collect (originalStreamId, globalPosition) pairs from a read result.
streamPositions :: Vector RecordedEvent -> [(Int64, Int64)]
streamPositions =
  map
    ( \e ->
        ( case e ^. #originalStreamId of StreamId s -> s,
          case e ^. #globalPosition of GlobalPosition p -> p
        )
    )
    . V.toList

-- | Direct call to the partition rule for one stream id and size, returning the
-- member index PostgreSQL computes. Mirrors IP-1 exactly so the test pins the
-- formula, not just its consequences.
runMemberOf :: KirokuStore -> Int64 -> Int32 -> IO Int32
runMemberOf store streamId size = runStmt store (Session.statement (streamId, size) stmt)
  where
    stmt :: Statement (Int64, Int32) Int32
    stmt =
      preparable
        "SELECT (((hashtextextended($1::text, 0) % $2) + $2) % $2)::int4"
        ( contrazip2
            (E.param (E.nonNullable E.int8))
            (E.param (E.nonNullable E.int4))
        )
        (D.singleRow (D.column (D.nonNullable D.int4)))

-- | All distinct (originalStreamId) values present in the unpartitioned read.
distinctStreamIds :: Vector RecordedEvent -> [Int64]
distinctStreamIds = Set.toList . Set.fromList . map fst . streamPositions

spec :: Spec
spec = do
  describe "ConsumerGroupSql" $ do
    categorySpec
    checkpointSpec
    allSpec
    memberOfSpec

-- ---------------------------------------------------------------------------
-- M1: category partitioning
-- ---------------------------------------------------------------------------

categorySpec :: Spec
categorySpec = around withTestStore $ do
  describe "category consumer-group partitioning (size 4)" $ do
    let cat = "acct"
        names = map (\i -> "acct-" <> T.pack (show i)) [1 .. 50 :: Int]
        size = 4 :: Int32

    let readMember store m =
          runStmt store $
            Session.statement (0 :: Int64, cat, m, size, bigLimit) SQL.readCategoryForwardConsumerGroupStmt
        readFull store =
          runStmt store $
            Session.statement (0 :: Int64, cat, bigLimit) SQL.readCategoryForwardStmt

    it "splits a category into 4 pairwise-disjoint member slices" $ \store -> do
      seedStreams store names 2
      slices <- mapM (readMember store) [0 .. size - 1]
      let idSets = map (Set.fromList . eventIds) slices
      -- pairwise disjoint: union of sizes equals size of union
      let totalIds = sum (map Set.size idSets)
          unionIds = Set.size (Set.unions idSets)
      totalIds `shouldBe` unionIds

    it "union of all member slices equals the unpartitioned category read" $ \store -> do
      seedStreams store names 2
      slices <- mapM (readMember store) [0 .. size - 1]
      full <- readFull store
      let unionIds = Set.unions (map (Set.fromList . eventIds) slices)
          fullIds = Set.fromList (eventIds full)
      unionIds `shouldBe` fullIds

    it "every stream's events go to exactly one member, in ascending global position" $ \store -> do
      seedStreams store names 3
      slices <- mapM (readMember store) [0 .. size - 1]
      -- For each member slice, every stream present must have all its events
      -- in ascending global position, and no stream may appear in two slices.
      let perMemberStreams = map (Set.fromList . map fst . streamPositions) slices
      -- per-stream affinity: stream sets are pairwise disjoint
      let totalStreams = sum (map Set.size perMemberStreams)
          unionStreams = Set.size (Set.unions perMemberStreams)
      totalStreams `shouldBe` unionStreams
      -- ascending global position within each slice
      mapM_
        ( \slice -> do
            let ps = map snd (streamPositions slice)
            ps `shouldBe` sort ps
        )
        slices

    it "member assignment is deterministic across repeated reads" $ \store -> do
      seedStreams store names 1
      firstReads <- mapM (\m -> eventIds <$> readMember store m) [0 .. size - 1]
      secondReads <- mapM (\m -> eventIds <$> readMember store m) [0 .. size - 1]
      firstReads `shouldBe` secondReads

    it "size 1 is equivalent to an unpartitioned category read" $ \store -> do
      seedStreams store names 2
      one <-
        runStmt store $
          Session.statement (0 :: Int64, cat, 0 :: Int32, 1 :: Int32, bigLimit) SQL.readCategoryForwardConsumerGroupStmt
      full <- readFull store
      eventIds one `shouldBe` eventIds full

-- ---------------------------------------------------------------------------
-- M2: per-member checkpoints
-- ---------------------------------------------------------------------------

checkpointSpec :: Spec
checkpointSpec = around withTestStore $ do
  describe "per-member checkpoints" $ do
    let subName = "proj-acct" :: Text

    it "stores and reads independent checkpoints per member" $ \store -> do
      runStmt store $ Session.statement (subName, 0 :: Int32, 7 :: Int64, 2 :: Int32, "unbound", Nothing) SQL.saveCheckpointMemberStmt
      runStmt store $ Session.statement (subName, 1 :: Int32, 13 :: Int64, 2 :: Int32, "unbound", Nothing) SQL.saveCheckpointMemberStmt
      m0 <- runStmt store $ Session.statement (subName, 0 :: Int32) SQL.getCheckpointMemberStmt
      m1 <- runStmt store $ Session.statement (subName, 1 :: Int32) SQL.getCheckpointMemberStmt
      m0 `shouldBe` Just 7
      m1 `shouldBe` Just 13

    it "never moves a member checkpoint backward (GREATEST monotonicity)" $ \store -> do
      runStmt store $ Session.statement (subName, 0 :: Int32, 20 :: Int64, 2 :: Int32, "unbound", Nothing) SQL.saveCheckpointMemberStmt
      runStmt store $ Session.statement (subName, 0 :: Int32, 5 :: Int64, 2 :: Int32, "unbound", Nothing) SQL.saveCheckpointMemberStmt
      m0 <- runStmt store $ Session.statement (subName, 0 :: Int32) SQL.getCheckpointMemberStmt
      m0 `shouldBe` Just 20

    it "missing member checkpoint reads as Nothing" $ \store -> do
      m9 <- runStmt store $ Session.statement (subName, 9 :: Int32) SQL.getCheckpointMemberStmt
      m9 `shouldBe` Nothing

    it "the existing name-keyed checkpoint statements still round-trip (as member 0)" $ \store -> do
      runStmt store $ Session.statement (subName, 42 :: Int64) SQL.saveCheckpointStmt
      -- name-keyed read returns the same row (member 0)
      byName <- runStmt store $ Session.statement subName SQL.getCheckpointStmt
      byMember0 <- runStmt store $ Session.statement (subName, 0 :: Int32) SQL.getCheckpointMemberStmt
      byName `shouldBe` Just 42
      byMember0 `shouldBe` Just 42

-- ---------------------------------------------------------------------------
-- M3: $all partitioning
-- ---------------------------------------------------------------------------

allSpec :: Spec
allSpec = around withTestStore $ do
  describe "$all consumer-group partitioning (size 4)" $ do
    -- \$all spans several categories; partitioning is by originating stream.
    let names =
          concatMap
            (\c -> map (\i -> c <> "-" <> T.pack (show i)) [1 .. 20 :: Int])
            ["acct", "user", "order"]
        size = 4 :: Int32

    let readMember store m =
          runStmt store $
            Session.statement (0 :: Int64, m, size, bigLimit) SQL.readAllForwardConsumerGroupStmt
        readFull store =
          runStmt store $
            Session.statement (0 :: Int64, bigLimit) SQL.readAllForwardStmt

    it "splits $all into 4 pairwise-disjoint member slices" $ \store -> do
      seedStreams store names 2
      slices <- mapM (readMember store) [0 .. size - 1]
      let idSets = map (Set.fromList . eventIds) slices
          totalIds = sum (map Set.size idSets)
          unionIds = Set.size (Set.unions idSets)
      totalIds `shouldBe` unionIds

    it "union of all member slices equals the unpartitioned $all read" $ \store -> do
      seedStreams store names 2
      slices <- mapM (readMember store) [0 .. size - 1]
      full <- readFull store
      let unionIds = Set.unions (map (Set.fromList . eventIds) slices)
          fullIds = Set.fromList (eventIds full)
      unionIds `shouldBe` fullIds

    it "every stream's events go to exactly one member, in ascending global position" $ \store -> do
      seedStreams store names 3
      slices <- mapM (readMember store) [0 .. size - 1]
      let perMemberStreams = map (Set.fromList . map fst . streamPositions) slices
          totalStreams = sum (map Set.size perMemberStreams)
          unionStreams = Set.size (Set.unions perMemberStreams)
      totalStreams `shouldBe` unionStreams
      mapM_
        ( \slice -> do
            let ps = map snd (streamPositions slice)
            ps `shouldBe` sort ps
        )
        slices

    it "size 1 is equivalent to an unpartitioned $all read" $ \store -> do
      seedStreams store names 2
      one <-
        runStmt store $
          Session.statement (0 :: Int64, 0 :: Int32, 1 :: Int32, bigLimit) SQL.readAllForwardConsumerGroupStmt
      full <- readFull store
      eventIds one `shouldBe` eventIds full

-- ---------------------------------------------------------------------------
-- The partition rule, pinned directly.
-- ---------------------------------------------------------------------------

memberOfSpec :: Spec
memberOfSpec = around withTestStore $ do
  describe "member_of assignment rule" $ do
    it "returns a member index in [0, size) for every stream" $ \store -> do
      let names = map (\i -> "rule-" <> T.pack (show i)) [1 .. 30 :: Int]
      seedStreams store names 1
      full <- runStmt store $ Session.statement (0 :: Int64, bigLimit) SQL.readAllForwardStmt
      let sids = distinctStreamIds full
      mapM_
        ( \sid -> do
            m <- runMemberOf store sid 4
            m `shouldSatisfy` (\x -> x >= 0 && x < 4)
        )
        sids