SQL.hs

{-# LANGUAGE MultilineStrings #-}

module Kiroku.Store.HistoryRetention.SQL
  ( acquireLeaseStmt,
    lockCoordinatorStmt,
    readLeaseForUpdateStmt,
    renewLeaseStmt,
    releaseLeaseStmt,
    leaseInventoryStmt,
    pruneLeasesStmt,
    activeConflictStmt,
    lockStreamHistoryStmt,
    lockAffectedStreamsForHardDeleteStmt,
  )
where

import Contravariant.Extras (contrazip2, contrazip3)
import Data.Int (Int32, Int64)
import Data.Text (Text)
import Data.Time.Clock (DiffTime, UTCTime)
import Data.UUID (UUID)
import Data.Vector (Vector)
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Kiroku.Store.HistoryRetention.Types
import Kiroku.Store.Types (GlobalPosition (..), StreamId (..), StreamInfo (..), StreamName (..), StreamVersion (..))

acquireLeaseStmt :: Statement (Text, Text, DiffTime) HistoryRetentionLease
acquireLeaseStmt =
  preparable
    """
    WITH coordinator AS MATERIALIZED (
      SELECT singleton
      FROM history_retention_coordinator
      WHERE singleton
      FOR UPDATE
    ), snapshot AS MATERIALIZED (
      SELECT streams.stream_version AS protected_through,
             clock_timestamp() AS database_now
      FROM streams
      CROSS JOIN coordinator
      WHERE streams.stream_id = 0
    ), inserted AS (
      INSERT INTO history_retention_leases
        (owner, reason, protected_through, created_at, renewed_at, expires_at)
      SELECT $1, $2, protected_through, database_now, database_now, database_now + $3
      FROM snapshot
      RETURNING lease_id, owner, reason, protected_through,
                created_at, renewed_at, expires_at, released_at
    )
    SELECT lease_id, owner, reason, protected_through,
           created_at, renewed_at, expires_at, released_at,
           'active'::text
    FROM inserted
    """
    (contrazip3 textParam textParam intervalParam)
    (D.singleRow leaseRow)

lockCoordinatorStmt :: Statement () Bool
lockCoordinatorStmt =
  preparable
    """
    SELECT singleton
    FROM history_retention_coordinator
    WHERE singleton
    FOR UPDATE
    """
    E.noParams
    (D.singleRow (column D.bool))

readLeaseForUpdateStmt :: Statement UUID (Maybe HistoryRetentionLease)
readLeaseForUpdateStmt =
  preparable
    """
    SELECT lease_id, owner, reason, protected_through,
           created_at, renewed_at, expires_at, released_at,
           CASE
             WHEN released_at IS NOT NULL THEN 'released'
             WHEN expires_at <= clock_timestamp() THEN 'expired'
             ELSE 'active'
           END::text
    FROM history_retention_leases
    WHERE lease_id = $1
    FOR UPDATE
    """
    uuidParam
    (D.rowMaybe leaseRow)

renewLeaseStmt :: Statement (UUID, DiffTime) HistoryRetentionLease
renewLeaseStmt =
  preparable
    """
    WITH renewal_time AS MATERIALIZED (
      SELECT clock_timestamp() AS database_now
    ), renewed AS (
      UPDATE history_retention_leases
      SET renewed_at = renewal_time.database_now,
          expires_at = GREATEST(
            history_retention_leases.expires_at,
            renewal_time.database_now + $2
          )
      FROM renewal_time
      WHERE lease_id = $1
      RETURNING lease_id, owner, reason, protected_through,
                created_at, renewed_at, expires_at, released_at
    )
    SELECT lease_id, owner, reason, protected_through,
           created_at, renewed_at, expires_at, released_at,
           'active'::text
    FROM renewed
    """
    (contrazip2 uuidParam intervalParam)
    (D.singleRow leaseRow)

releaseLeaseStmt :: Statement UUID HistoryRetentionLease
releaseLeaseStmt =
  preparable
    """
    WITH released AS (
      UPDATE history_retention_leases
      SET released_at = clock_timestamp()
      WHERE lease_id = $1
      RETURNING lease_id, owner, reason, protected_through,
                created_at, renewed_at, expires_at, released_at
    )
    SELECT lease_id, owner, reason, protected_through,
           created_at, renewed_at, expires_at, released_at,
           'released'::text
    FROM released
    """
    uuidParam
    (D.singleRow leaseRow)

leaseInventoryStmt :: Statement Int32 (Vector HistoryRetentionLease)
leaseInventoryStmt =
  preparable
    """
    WITH inventory_time AS MATERIALIZED (
      SELECT clock_timestamp() AS database_now
    )
    SELECT lease_id, owner, reason, protected_through,
           created_at, renewed_at, expires_at, released_at,
           CASE
             WHEN released_at IS NOT NULL THEN 'released'
             WHEN expires_at <= inventory_time.database_now THEN 'expired'
             ELSE 'active'
           END::text
    FROM history_retention_leases
    CROSS JOIN inventory_time
    ORDER BY created_at, lease_id
    LIMIT $1
    """
    int4Param
    (D.rowVector leaseRow)

pruneLeasesStmt :: Statement UTCTime HistoryRetentionPruneResult
pruneLeasesStmt =
  preparable
    """
    WITH released AS (
      DELETE FROM history_retention_leases
      WHERE released_at IS NOT NULL
        AND released_at < $1
      RETURNING 1
    ), expired AS (
      DELETE FROM history_retention_leases
      WHERE released_at IS NULL
        AND expires_at < $1
      RETURNING 1
    )
    SELECT (SELECT count(*) FROM expired),
           (SELECT count(*) FROM released)
    """
    timestamptzParam
    ( D.singleRow
        ( HistoryRetentionPruneResult
            <$> column D.int8
            <*> column D.int8
        )
    )

activeConflictStmt :: Statement () (Maybe HistoryRetentionConflict)
activeConflictStmt =
  preparable
    """
    SELECT count(*), min(expires_at)
    FROM history_retention_leases
    WHERE released_at IS NULL
      AND expires_at > clock_timestamp()
    HAVING count(*) > 0
    """
    E.noParams
    ( D.rowMaybe
        ( HistoryRetentionConflict
            <$> column D.int8
            <*> column D.timestamptz
        )
    )

lockStreamHistoryStmt :: Statement Text (Maybe StreamInfo)
lockStreamHistoryStmt =
  preparable
    """
    SELECT stream_id, stream_name, stream_version,
           created_at, deleted_at, truncate_before
    FROM streams
    WHERE stream_name = $1
    FOR SHARE
    """
    textParam
    (D.rowMaybe streamInfoRow)

lockAffectedStreamsForHardDeleteStmt :: Statement Int64 (Vector Int64)
lockAffectedStreamsForHardDeleteStmt =
  preparable
    """
    WITH affected_streams AS (
      SELECT $1::bigint AS stream_id
      UNION
      SELECT stream_events.stream_id
      FROM stream_events
      WHERE stream_events.original_stream_id = $1
        AND stream_events.stream_id <> 0
    )
    SELECT streams.stream_id
    FROM streams
    JOIN affected_streams USING (stream_id)
    ORDER BY streams.stream_id
    FOR UPDATE OF streams
    """
    int8Param
    (D.rowVector (column D.int8))

leaseRow :: D.Row HistoryRetentionLease
leaseRow =
  makeLease
    <$> column D.uuid
    <*> column D.text
    <*> column D.text
    <*> column D.int8
    <*> column D.timestamptz
    <*> column D.timestamptz
    <*> column D.timestamptz
    <*> D.column (D.nullable D.timestamptz)
    <*> column D.text

makeLease :: UUID -> Text -> Text -> Int64 -> UTCTime -> UTCTime -> UTCTime -> Maybe UTCTime -> Text -> HistoryRetentionLease
makeLease leaseUuid ownerText reasonText frontier created renewed expires released stateText =
  HistoryRetentionLease
    { leaseId = HistoryRetentionLeaseId leaseUuid,
      owner = requireValidated (mkHistoryRetentionLeaseOwner ownerText),
      reason = requireValidated (mkHistoryRetentionLeaseReason reasonText),
      protectedThrough = GlobalPosition frontier,
      createdAt = created,
      renewedAt = renewed,
      expiresAt = expires,
      releasedAt = released,
      state = case stateText of
        "active" -> HistoryRetentionLeaseActive
        "expired" -> HistoryRetentionLeaseExpired
        "released" -> HistoryRetentionLeaseReleased
        unexpected -> error ("unknown history retention lease state from database: " <> show unexpected)
    }

streamInfoRow :: D.Row StreamInfo
streamInfoRow =
  StreamInfo
    <$> (StreamId <$> column D.int8)
    <*> (StreamName <$> column D.text)
    <*> (StreamVersion <$> column D.int8)
    <*> column D.timestamptz
    <*> D.column (D.nullable D.timestamptz)
    <*> (StreamVersion <$> column D.int8)

requireValidated :: Either error value -> value
requireValidated = either (const (error "database returned an invalid history retention lease")) (\value -> value)

column :: D.Value value -> D.Row value
column = D.column . D.nonNullable

textParam :: E.Params Text
textParam = E.param (E.nonNullable E.text)

uuidParam :: E.Params UUID
uuidParam = E.param (E.nonNullable E.uuid)

int4Param :: E.Params Int32
int4Param = E.param (E.nonNullable E.int4)

int8Param :: E.Params Int64
int8Param = E.param (E.nonNullable E.int8)

intervalParam :: E.Params DiffTime
intervalParam = E.param (E.nonNullable E.interval)

timestamptzParam :: E.Params UTCTime
timestamptzParam = E.param (E.nonNullable E.timestamptz)