Connection.hs

-- | Internal module. Not part of the public API.
module Database.Bolty.Connection
  ( queryIO
  , queryPIO
  , queryIO'
  , queryPIO'
  , queryWithFieldsIO
  , queryPWithFieldsIO
  , queryPMetaIO
  , requestResponseRunIO
  , requestResponsePullIO
  , withTxMode
  ) where

import Prelude

import           Database.Bolty.AccessMode        (AccessMode)
import qualified Database.Bolty.Connection.Pipe as P
import           Database.Bolty.Connection.Type
import           Database.Bolty.Logging
import           Database.Bolty.Message.Request   (Begin(Begin))
import           Database.Bolty.Record

import           Control.Exception             (SomeException, onException, throwIO, try)
import           Data.IORef                    (IORef, newIORef, readIORef, writeIORef)
import           Data.Kind                     (Type)
import           Data.Text                     (Text)
import           Data.Word                     (Word64)
import           Debug.Trace                   (traceEventIO)
import           GHC.Clock                     (getMonotonicTimeNSec)
import           GHC.Stack                     (HasCallStack)

import           Database.Bolty.Message.Response
import           Database.Bolty.Message.Request
import qualified Data.HashMap.Lazy as H
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import           Data.PackStream.Ps (Ps, fromPs)
import           Data.PackStream.Result (Result(..))


-- ---------------------------------------------------------------------------
-- Growable vector (O(n) amortized appends via mutable IOVector + doubling)
-- ---------------------------------------------------------------------------

type GrowVec :: Type -> Type
type role GrowVec representational
data GrowVec a = GrowVec
  { gvBuf  :: !(IORef (MV.IOVector a))
  , gvLen  :: !(IORef Int)
  }

newGrowVec :: IO (GrowVec a)
newGrowVec = do
  buf <- MV.new 64
  GrowVec <$> newIORef buf <*> newIORef 0

pushGrowVec :: GrowVec a -> a -> IO ()
pushGrowVec GrowVec{gvBuf, gvLen} x = do
  i <- readIORef gvLen
  buf <- readIORef gvBuf
  buf' <- if i >= MV.length buf
          then MV.unsafeGrow buf (MV.length buf)
          else pure buf
  MV.write buf' i x
  writeIORef gvBuf buf'
  writeIORef gvLen (i + 1)

freezeGrowVec :: GrowVec a -> IO (V.Vector a)
freezeGrowVec GrowVec{gvBuf, gvLen} = do
  n <- readIORef gvLen
  buf <- readIORef gvBuf
  V.unsafeFreeze (MV.take n buf)


-- | Run a query, return only records.
queryIO :: HasCallStack => Connection -> Text -> IO (V.Vector Record)
queryIO conn cypher = records <$> queryPIO' conn cypher H.empty

-- | Run a parameterized query, return only records.
queryPIO :: HasCallStack => Connection -> Text -> H.HashMap Text Ps -> IO (V.Vector Record)
queryPIO conn cypher params = records <$> queryPIO' conn cypher params

-- | Run a query, return full response.
queryIO' :: HasCallStack => Connection -> Text -> IO SuccessPull
queryIO' conn cypher = queryPIO' conn cypher H.empty

-- | Run a parameterized query, return full response.
queryPIO' :: HasCallStack => Connection -> Text -> H.HashMap Text Ps -> IO SuccessPull
queryPIO' conn cypher params = do
  (runResp, pullResp, clientNs) <- timedQuery conn cypher params
  fireLogger conn cypher params runResp pullResp clientNs
  fireNotifications conn pullResp
  pure pullResp

-- | Send RUN and receive the run response.
requestResponseRunIO :: HasCallStack => Connection -> Text -> H.HashMap Text Ps -> IO SuccessRun
requestResponseRunIO conn cypher params = do
  st <- P.getState conn
  P.requireStateIO conn [Ready, TXready] "RUN"
  let request = case st of
        TXready -> RRunExplicitTransaction $ RunExplicitTransaction cypher params defaultRunExtra
        _       -> RRunAutoCommitTransaction $ mkRunAutoCommit cypher params
  P.flushIO conn request
  response <- P.fetchIO conn
  case response of
    RSuccess meta -> do
      P.setState conn $ case st of
        TXready -> TXstreaming
        _       -> Streaming
      case makeResponseRunAutoCommitTransaction meta of
        Left err -> throwIO $ WrongMessageFormat err
        Right v  -> pure v
    RIgnored -> do
      P.reset conn
      throwIO ResponseErrorIgnored
    RFailure Failure{code, message} -> do
      P.setState conn Failed
      P.reset conn
      throwIO $ ResponseErrorFailure code message
    RRecord _ -> do
      P.reset conn
      throwIO ResponseErrorRecords

-- | Pull all records from an in-progress query.
requestResponsePullIO :: HasCallStack => Connection -> IO SuccessPull
requestResponsePullIO conn = do
  P.requireStateIO conn [Streaming, TXstreaming] "PULL"
  let pull = Pull { n = conn.fetchSize, qid = Nothing }
  P.flushIO conn $ RPull pull
  gv <- newGrowVec
  go gv
  where
    go gv = do
      response <- P.fetchIO conn
      case response of
        RSuccess meta -> do
          let has_more = case H.lookup "has_more" meta of
                           Just hm -> case fromPs hm of
                             Success True -> True
                             _            -> False
                           Nothing -> False
          if has_more then do
            let pull = Pull { n = conn.fetchSize, qid = Nothing }
            P.flushIO conn $ RPull pull
            go gv
          else do
            recs <- freezeGrowVec gv
            st <- P.getState conn
            P.setState conn $ case st of
              TXstreaming -> TXready
              _           -> Ready
            case makeSuccessPull recs meta of
              Left err -> throwIO $ WrongMessageFormat err
              Right v  -> pure v
        RIgnored -> do
          P.reset conn
          throwIO ResponseErrorIgnored
        RFailure Failure{code, message} -> do
          P.setState conn Failed
          P.reset conn
          throwIO $ ResponseErrorFailure code message
        RRecord record -> do
          pushGrowVec gv record
          go gv

-- | Run a parameterized query, returning column names alongside records.
queryPWithFieldsIO :: HasCallStack => Connection -> T.Text -> H.HashMap T.Text Ps -> IO (V.Vector T.Text, V.Vector Record)
queryPWithFieldsIO conn cypher params = do
  (runResp, pullResp, clientNs) <- timedQuery conn cypher params
  fireLogger conn cypher params runResp pullResp clientNs
  fireNotifications conn pullResp
  pure (successFields runResp, records pullResp)

-- | Run a query, returning column names alongside records.
queryWithFieldsIO :: HasCallStack => Connection -> T.Text -> IO (V.Vector T.Text, V.Vector Record)
queryWithFieldsIO conn cypher = queryPWithFieldsIO conn cypher H.empty


-- | Run a parameterized query, returning field names, records, and metadata.
queryPMetaIO :: HasCallStack => Connection -> T.Text -> H.HashMap T.Text Ps -> IO (V.Vector T.Text, V.Vector Record, QueryMeta)
queryPMetaIO conn cypher params = do
  (runResp, pullResp, clientNs) <- timedQuery conn cypher params
  fireLogger conn cypher params runResp pullResp clientNs
  fireNotifications conn pullResp
  pure (successFields runResp, records pullResp, infos pullResp)


-- ---------------------------------------------------------------------------
-- Internal: timing + logging
-- ---------------------------------------------------------------------------

-- | Run a query with monotonic clock timing, returning the responses and elapsed nanoseconds.
-- Uses Bolt pipelining: sends RUN + PULL before reading either response,
-- saving one network round-trip compared to the sequential approach.
timedQuery :: HasCallStack => Connection -> T.Text -> H.HashMap T.Text Ps -> IO (SuccessRun, SuccessPull, Word64)
timedQuery conn cypher params = do
  traceEventIO "START boltWireIO"
  t0 <- getMonotonicTimeNSec

  -- Build RUN request
  st <- P.getState conn
  P.requireStateIO conn [Ready, TXready] "RUN"
  let runRequest = case st of
        TXready -> RRunExplicitTransaction $ RunExplicitTransaction cypher params defaultRunExtra
        _       -> RRunAutoCommitTransaction $ mkRunAutoCommit cypher params

  -- Pipeline: send RUN + PULL before receiving either response
  let pull = Pull { n = conn.fetchSize, qid = Nothing }
  P.flushIO conn runRequest
  P.flushIO conn (RPull pull)

  -- Receive RUN response
  runResult <- try @SomeException $ do
    runResponse <- P.fetchIO conn
    case runResponse of
      RSuccess meta -> do
        P.setState conn $ case st of
          TXready -> TXstreaming
          _       -> Streaming
        case makeResponseRunAutoCommitTransaction meta of
          Left err -> throwIO $ WrongMessageFormat err
          Right v  -> pure v
      RFailure Failure{code, message} -> do
        P.setState conn Failed
        throwIO $ ResponseErrorFailure code message
      RIgnored ->
        throwIO ResponseErrorIgnored
      RRecord _ ->
        throwIO ResponseErrorRecords

  case runResult of
    Left err -> do
      -- RUN failed: consume the IGNORED response for the pipelined PULL
      _ <- P.fetchIO conn
      P.reset conn
      throwIO err
    Right runResp -> do
      -- Receive PULL responses
      gv <- newGrowVec
      pullResp <- pullRecords conn gv
      t1 <- getMonotonicTimeNSec
      traceEventIO "STOP boltWireIO"
      pure (runResp, pullResp, t1 - t0)

-- | Pull all records from an in-progress pipelined query.
-- The initial PULL has already been sent; this receives responses and handles has_more pagination.
pullRecords :: HasCallStack => Connection -> GrowVec Record -> IO SuccessPull
pullRecords conn gv = do
  response <- P.fetchIO conn
  case response of
    RSuccess meta -> do
      let has_more = case H.lookup "has_more" meta of
                       Just hm -> case fromPs hm of
                         Success True -> True
                         _            -> False
                       Nothing -> False
      if has_more then do
        let pull = Pull { n = conn.fetchSize, qid = Nothing }
        P.flushIO conn $ RPull pull
        pullRecords conn gv
      else do
        recs <- freezeGrowVec gv
        st <- P.getState conn
        P.setState conn $ case st of
          TXstreaming -> TXready
          _           -> Ready
        case makeSuccessPull recs meta of
          Left err -> throwIO $ WrongMessageFormat err
          Right v  -> pure v
    RIgnored -> do
      P.reset conn
      throwIO ResponseErrorIgnored
    RFailure Failure{code, message} -> do
      P.setState conn Failed
      P.reset conn
      throwIO $ ResponseErrorFailure code message
    RRecord record -> do
      pushGrowVec gv record
      pullRecords conn gv

-- | Fire the query logger callback if configured.
fireLogger :: Connection -> T.Text -> H.HashMap T.Text Ps -> SuccessRun -> SuccessPull -> Word64 -> IO ()
fireLogger Connection{queryLogger} cypher params runResp SuccessPull{records, infos} clientNs =
  case queryLogger of
    Nothing -> pure ()
    Just logger -> logger QueryLog
      { qlCypher      = cypher
      , qlParameters  = params
      , qlRowCount    = V.length records
      , qlServerFirst = successTFirst runResp
      , qlServerLast  = t_last infos
      , qlClientTime  = fromIntegral clientNs / 1_000_000
      } infos

-- | Fire the notification handler for each parsed notification, if configured.
fireNotifications :: Connection -> SuccessPull -> IO ()
fireNotifications Connection{notificationHandler} SuccessPull{infos = QueryMeta{parsedNotifications}} =
  case notificationHandler of
    Nothing      -> pure ()
    Just handler -> V.mapM_ handler parsedNotifications


-- | Run an action inside an explicit transaction opened with the given
-- 'AccessMode'. Automatically commits on success and attempts a rollback
-- on failure.
--
-- The mode controls cluster routing: 'ReadAccess' routes to a follower
-- and engages the server's read-only enforcement; 'WriteAccess' routes
-- to the leader. Higher-level helpers in "Database.Bolty"
-- ('Database.Bolty.withTransaction' / 'Database.Bolty.withReadTransaction')
-- are thin wrappers over this primitive.
withTxMode :: HasCallStack => AccessMode -> Connection -> (Connection -> IO a) -> IO a
withTxMode mode conn action = do
  P.beginTx conn (Begin V.empty Nothing H.empty mode Nothing Nothing)
  result <- action conn `onException` P.tryRollback conn
  _ <- P.commitTx conn
  pure result