Session.hs

-- | Sessions: managed transactions with automatic bookmark tracking for causal consistency.
module Database.Bolty.Session
  ( -- * Bookmark management
    BookmarkManager
  , newBookmarkManager
  , getBookmarks
  , updateBookmark
    -- * Session
  , Session(..)
  , SessionPool(..)
  , SessionConfig(..)
  , defaultSessionConfig
  , createSession
  , createRoutingSession
    -- * Managed transactions
  , readTransaction
  , writeTransaction
    -- * Querying
  , getLastBookmarks
  ) where

import           Control.Concurrent             (threadDelay)
import           Control.Exception              (SomeException, fromException, onException,
                                                 throwIO, try)
import           Control.Monad                  (when)
import qualified Data.HashMap.Lazy              as H
import           Data.IORef                     (IORef, newIORef, readIORef, writeIORef)
import           Data.Kind                      (Type)
import           Data.Text                      (Text)
import qualified Data.Vector                    as V
import           GHC.Stack                      (HasCallStack)

import           Database.Bolty.AccessMode      (AccessMode(..))
import           Database.Bolty.Connection.Type
import qualified Database.Bolty.Connection.Pipe as P
import           Database.Bolty.Message.Request (Begin(Begin), TelemetryApi(..))
import           Database.Bolty.Pool
import           Database.Bolty.Routing (RoutingPool(..), withRoutingConnection,
                                        invalidateRoutingTable)


-- ---------------------------------------------------------------------------
-- Bookmark management
-- ---------------------------------------------------------------------------

-- | Mutable bookmark holder for causal consistency across transactions.
-- Within a session, each committed transaction produces a bookmark that
-- supersedes all previous bookmarks. The manager tracks the latest
-- bookmark(s) so they can be passed to the next transaction's BEGIN.
type BookmarkManager :: Type
newtype BookmarkManager = BookmarkManager (IORef [Text])

-- | Create a bookmark manager with optional initial bookmarks.
newBookmarkManager :: [Text] -> IO BookmarkManager
newBookmarkManager initial = BookmarkManager <$> newIORef initial

-- | Get the current bookmarks.
getBookmarks :: BookmarkManager -> IO [Text]
getBookmarks (BookmarkManager ref) = readIORef ref

-- | Replace the current bookmarks with the new one from a COMMIT response.
updateBookmark :: BookmarkManager -> Text -> IO ()
updateBookmark (BookmarkManager ref) bm = writeIORef ref [bm]


-- ---------------------------------------------------------------------------
-- Session
-- ---------------------------------------------------------------------------

-- | Session configuration.
type SessionConfig :: Type
data SessionConfig = SessionConfig
  { database         :: !(Maybe Text)
  -- ^ Database to use (Nothing = default database).
  , accessMode       :: !AccessMode
  -- ^ Default access mode for auto-commit queries.
  , sessionBookmarks :: ![Text]
  -- ^ Initial bookmarks for causal consistency (e.g. from a previous session).
  }

-- | Default session configuration: default database, write access, no initial bookmarks.
defaultSessionConfig :: SessionConfig
defaultSessionConfig = SessionConfig
  { database         = Nothing
  , accessMode       = WriteAccess
  , sessionBookmarks = []
  }


-- | Which pool type the session uses (internal).
type SessionPool :: Type
data SessionPool
  = DirectPool !BoltPool
  | RoutedPool !RoutingPool

-- | A session bundles a connection pool with bookmark management and configuration.
-- Sessions track bookmarks automatically across managed transactions
-- ('readTransaction', 'writeTransaction') to ensure causal consistency.
type Session :: Type
data Session = Session
  { sPool            :: !SessionPool
  , sBookmarks       :: !BookmarkManager
  , sConfig          :: !SessionConfig
  , sRetry           :: !RetryConfig
  , sTelemetrySent   :: !(IORef Bool)
  }


-- | Create a session using a direct (non-routing) connection pool.
createSession :: BoltPool -> SessionConfig -> IO Session
createSession pool cfg = do
  bm <- newBookmarkManager (sessionBookmarks cfg)
  ts <- newIORef False
  pure Session
    { sPool            = DirectPool pool
    , sBookmarks       = bm
    , sConfig          = cfg
    , sRetry           = bpRetryConfig pool
    , sTelemetrySent   = ts
    }


-- | Create a session using a routing-aware connection pool.
createRoutingSession :: RoutingPool -> SessionConfig -> IO Session
createRoutingSession rp cfg = do
  bm <- newBookmarkManager (sessionBookmarks cfg)
  ts <- newIORef False
  pure Session
    { sPool            = RoutedPool rp
    , sBookmarks       = bm
    , sConfig          = cfg
    , sRetry           = retryConfig (rpPoolConfig rp)
    , sTelemetrySent   = ts
    }


-- | Get the last bookmarks from the session. Pass these to a new session's
-- 'SessionConfig' to ensure the new session sees all committed writes.
getLastBookmarks :: Session -> IO [Text]
getLastBookmarks Session{sBookmarks} = getBookmarks sBookmarks


-- ---------------------------------------------------------------------------
-- Managed transactions
-- ---------------------------------------------------------------------------

-- | Run a read transaction. Uses 'ReadAccess' for routing (directs to
-- read replicas in a cluster). Automatically handles BEGIN, COMMIT,
-- ROLLBACK, bookmark propagation, and retries on transient failures.
readTransaction :: HasCallStack => Session -> (Connection -> IO a) -> IO a
readTransaction session action = managedTransaction session ReadAccess action


-- | Run a write transaction. Uses 'WriteAccess' for routing (directs to
-- the leader in a cluster). Automatically handles BEGIN, COMMIT,
-- ROLLBACK, bookmark propagation, and retries on transient failures.
writeTransaction :: HasCallStack => Session -> (Connection -> IO a) -> IO a
writeTransaction session action = managedTransaction session WriteAccess action


-- | Internal: run a managed transaction with the given access mode.
-- Re-acquires a connection on each retry so that routing errors (NotALeader)
-- result in fresh routing table lookups.
managedTransaction :: HasCallStack => Session -> AccessMode -> (Connection -> IO a) -> IO a
managedTransaction Session{sPool, sBookmarks, sConfig, sRetry, sTelemetrySent} mode action =
  let maxR = maxRetries sRetry
      initD = initialDelay sRetry
      maxD = maxDelay sRetry
  in go maxR initD maxD
  where
    go 0 _ _    = attempt
    go n delay maxD = do
      result <- try attempt
      case result of
        Right x -> pure x
        Left (e :: SomeException) -> case fromException e :: Maybe Error of
          Just err
            | isTransient err -> do
                threadDelay delay
                go (n - 1) (min maxD (delay * 2)) maxD
            | isRoutingError err -> do
                invalidatePool sPool
                threadDelay delay
                go (n - 1) (min maxD (delay * 2)) maxD
          _ -> throwIO e

    attempt =
      withSessionConnection sPool mode $ \conn -> do
        -- Get current bookmarks and begin the transaction
        bms <- getBookmarks sBookmarks
        P.beginTx conn $ Begin (V.fromList bms) Nothing H.empty mode (database sConfig) Nothing
        -- Run the user's action, rollback on error
        result <- action conn `onException` P.tryRollback conn
        -- Commit and extract bookmark
        mbBookmark <- P.commitTx conn
        case mbBookmark of
          Just bm -> updateBookmark sBookmarks bm
          Nothing -> pure ()
        -- Send telemetry after first successful managed transaction
        sent <- readIORef sTelemetrySent
        when (not sent) $ do
          writeIORef sTelemetrySent True
          P.sendTelemetry conn ManagedTransactions
        pure result

    invalidatePool (RoutedPool rp) = invalidateRoutingTable rp
    invalidatePool (DirectPool _)  = pure ()


-- | Acquire a connection based on pool type and access mode.
withSessionConnection :: SessionPool -> AccessMode -> (Connection -> IO a) -> IO a
withSessionConnection (DirectPool pool) _mode action =
  withConnection pool action
withSessionConnection (RoutedPool rp) mode action =
  withRoutingConnection rp mode action