Request.hs

{-# LANGUAGE ApplicativeDo #-}

-- | Internal module. Not part of the public API.
module Database.Bolty.Message.Request
  ( Request(..)
  , Hello(..)
  , Logon(..)
  , TelemetryApi(..)
  , RunExplicitTransaction(..)
  , RunAutoCommitTransaction(..), defaultRunAutoCommitTransaction, mkRunAutoCommit
  , Discard(..)
  , Begin(..)
  , Pull(..), defaultPull
  , Route(..)
  , RouteExtra(..)
  , RunExtra(..), defaultRunExtra
  ) where

import           Data.Int                      (Int64)
import           Data.Kind                     (Type)
import           Data.Maybe                    (fromMaybe)
import           GHC.Generics                  (Generic)
import qualified Data.HashMap.Lazy             as H
import qualified Data.Text                     as T
import qualified Data.Vector                   as V

import           Data.PackStream
import           Data.PackStream.Integer       (toPSInteger, fromPSInteger)
import           Database.Bolty.AccessMode     (AccessMode(..))
import           Database.Bolty.Connection.Instances ()
import           Database.Bolty.Connection.Type (UserAgent(..), Scheme(..), Routing(..))


-- | A BOLT client request message.
type Request :: Type
data Request
  = RHello Hello
  | RGoodbye
  | RReset
  | RRunExplicitTransaction RunExplicitTransaction
  | RRunAutoCommitTransaction RunAutoCommitTransaction
  | RDiscard Discard
  | RPull Pull
  | RBegin Begin
  | RCommit
  | RRollback
  | RRoute Route
  | RLogon Logon
  | RLogoff
  | RTelemetry TelemetryApi
  deriving stock (Generic)


instance PackStream Request where
  toPs (RHello hello) = structureSingleton 0x01 $ toPs hello
  toPs (RGoodbye) = PsStructure 0x02 V.empty
  toPs (RReset) = PsStructure 0x0F V.empty
  toPs (RRunExplicitTransaction x) = runExplicitTransactionToPs x
  toPs (RRunAutoCommitTransaction x) = runAutoCommitTransactiontoPs x
  toPs (RDiscard discard) = structureSingleton 0x2F $ toPs discard
  toPs (RPull pull) = structureSingleton 0x3F $ toPs pull
  toPs (RBegin begin) = structureSingleton 0x11 $ toPs begin
  toPs (RCommit) = PsStructure 0x12 V.empty
  toPs (RRollback) = PsStructure 0x13 V.empty
  toPs (RRoute route) = routeToPs route
  toPs (RLogon logon) = structureSingleton 0x6A $ logonToPs logon
  toPs (RLogoff) = PsStructure 0x6B V.empty
  toPs (RTelemetry api) = PsStructure 0x54 $ V.singleton $ PsInteger $ toPSInteger $ telemetryApiToInt api
  fromPs (PsStructure 0x01 fields) = fromOneField fields RHello "expected exactly one field for Hello"
  fromPs (PsStructure 0x02 fields) = if V.null fields then Success RGoodbye else Error "additional data on RGoodbye request is not allowed"
  fromPs (PsStructure 0x0F fields) = if V.null fields then Success RReset else Error "additional data on RReset request is not allowed"
  fromPs (PsStructure 0x10 fields) = fromPsToRunRequest fields
  fromPs (PsStructure 0x2F fields) = fromOneField fields RDiscard "expected exactly one field for Discard"
  fromPs (PsStructure 0x3F fields) = fromOneField fields RPull "expected exactly one field for Pull"
  fromPs (PsStructure 0x11 fields) = fromOneField fields RBegin "expected exactly one field for Begin"
  fromPs (PsStructure 0x12 fields) = if V.null fields then Success RCommit else Error "additional data on RCommit request is not allowed"
  fromPs (PsStructure 0x13 fields) = if V.null fields then Success RRollback else Error "additional data on RRollback request is not allowed"
  fromPs (PsStructure 0x66 fields) = fromPsToRoute fields
  fromPs (PsStructure 0x6A fields) = fromOneField fields RLogon "expected exactly one field for Logon"
  fromPs (PsStructure 0x6B fields) = if V.null fields then Success RLogoff else Error "additional data on RLogoff request is not allowed"
  fromPs (PsStructure 0x54 fields) = fromPsToTelemetry fields
  fromPs _                         = Error "Could not decode Request"


fromPsToRunRequest :: V.Vector Ps -> Result Request
fromPsToRunRequest fields =
    if V.length fields /= 1 then
      Error "expected exactly one field for Run"
    else do
      arg <- fromPs $ V.unsafeHead fields
      f arg
  where f :: H.HashMap T.Text Ps -> Result Request
        f map = if any (\f -> H.member f map) ["bookmarks", "tx_timeout", "tx_metadata", "mode", "db", "imp_user"] then do
                  query       :: T.Text              <- lookupWithError "query" map "\"query\" not found in Run Request"
                  parameters  :: H.HashMap T.Text Ps <- lookupWithError "parameters" map "\"parameters\" not found in Run Request"
                  extra       :: RunExtra            <- lookupWithError "extra" map "\"extra\" not found in Run Request"
                  bookmarks   :: V.Vector T.Text     <- lookupWithError "bookmarks" map "\"bookmarks\" not found in Run Request"
                  tx_timeout  :: Maybe Int64         <- lookupMaybe "tx_timeout" map
                  tx_metadata :: H.HashMap T.Text Ps <- lookupWithError "tx_metadata" map "\"tx_metadata\" not found in Run Request"
                  mode        :: AccessMode          <- lookupMode map
                  db          :: Maybe T.Text        <- lookupMaybe "db" map
                  imp_user    :: Maybe T.Text        <- lookupMaybe "imp_user" map
                  pure $ RRunAutoCommitTransaction RunAutoCommitTransaction{query, parameters, extra, bookmarks, tx_timeout, tx_metadata, mode, db, imp_user}
                else do
                  query       :: T.Text              <- lookupWithError "query" map "\"query\" not found in Run Request"
                  parameters  :: H.HashMap T.Text Ps <- lookupWithError "parameters" map "\"parameters\" not found in Run Request"
                  extra       :: RunExtra            <- lookupWithError "extra" map "\"extra\" not found in Run Request"
                  pure $ RRunExplicitTransaction RunExplicitTransaction{query, parameters, extra}



-- | HELLO message payload for initial handshake authentication.
type Hello :: Type
data Hello =  Hello
  { user_agent :: UserAgent
  , scheme     :: Scheme
  , routing    :: Routing
  , patchBolt  :: Bool
  -- ^ When True, include @"patch_bolt": ["utc"]@ in HELLO (BOLT 4.3\/4.4 only).
  }
  deriving stock (Generic)

instance PackStream Hello where
  toPs = helloToPs
  fromPs = helloFromPs


-- | RUN message payload within an explicit transaction.
type RunExplicitTransaction :: Type
data RunExplicitTransaction = RunExplicitTransaction
  { query      :: !T.Text
  , parameters :: !(H.HashMap T.Text Ps)
  , extra      :: RunExtra
  }
  deriving stock (Generic)

runExplicitTransactionToPs :: RunExplicitTransaction -> Ps
runExplicitTransactionToPs RunExplicitTransaction{query, parameters, extra} =
  PsStructure 0x10 $ V.fromList
    [ toPs query
    , toPs parameters
    , toPs extra
    ]

-- | Default 'RunAutoCommitTransaction' with empty query, parameters, and bookmarks.
defaultRunAutoCommitTransaction :: RunAutoCommitTransaction
defaultRunAutoCommitTransaction = RunAutoCommitTransaction
  { query       = ""
  , parameters  = H.empty
  , extra       = defaultRunExtra
  , bookmarks   = V.empty
  , tx_timeout  = Nothing
  , tx_metadata = H.empty
  , mode        = WriteAccess
  , db          = Nothing
  , imp_user    = Nothing
  }

-- | Build a 'RunAutoCommitTransaction' with the given query and parameters.
mkRunAutoCommit :: T.Text -> H.HashMap T.Text Ps -> RunAutoCommitTransaction
mkRunAutoCommit q p = RunAutoCommitTransaction
  { query       = q
  , parameters  = p
  , extra       = defaultRunExtra
  , bookmarks   = V.empty
  , tx_timeout  = Nothing
  , tx_metadata = H.empty
  , mode        = WriteAccess
  , db          = Nothing
  , imp_user    = Nothing
  }

-- | RUN message payload for an auto-commit transaction.
type RunAutoCommitTransaction :: Type
data RunAutoCommitTransaction = RunAutoCommitTransaction
  { query       :: !T.Text
  , parameters  :: !(H.HashMap T.Text Ps)
  , extra       :: RunExtra
  , bookmarks   :: V.Vector T.Text
  , tx_timeout  :: !(Maybe Int64)
  , tx_metadata :: !(H.HashMap T.Text Ps)
  , mode        :: !AccessMode
  , db          :: !(Maybe T.Text)
  , imp_user    :: !(Maybe T.Text)
  }
  deriving stock (Generic)

runAutoCommitTransactiontoPs :: RunAutoCommitTransaction -> Ps
runAutoCommitTransactiontoPs RunAutoCommitTransaction{query, parameters, bookmarks, tx_timeout, tx_metadata, mode, db, imp_user} =
  PsStructure 0x10 $ V.fromList
    [ toPs query
    , toPs parameters
    , PsDictionary $ H.fromList $
        let tx_timeout_list = case tx_timeout of
                                Nothing -> []
                                Just t  -> [("tx_timeout", toPs t)]
            imp_user_list = case imp_user of
                              Nothing -> []
                              Just i  -> [("imp_user", toPs i)]
        in  [ ("bookmarks", toPs bookmarks)
            , ("tx_metadata", toPs tx_metadata)
            , ("mode", toPs mode)
            , ("db", toPs db)
            ] <> tx_timeout_list <> imp_user_list
    ]



-- | DISCARD message payload with record count and query ID.
type Discard :: Type
data Discard = Discard
  { n   :: Int64
  , qid :: Int64
  }
  deriving stock (Generic)
instance PackStream Discard where
  toPs Discard{n, qid} = PsDictionary $ H.fromList
    [("n", toPs n), ("qid", toPs qid)]
  fromPs = withDictionary "Discard" $ \m -> do
    n' <- lookupWithError "n" m "\"n\" not found in Discard"
    qid' <- lookupWithError "qid" m "\"qid\" not found in Discard"
    pure $ Discard{n = n', qid = qid'}

-- | BEGIN message payload for starting an explicit transaction.
type Begin :: Type
data Begin = Begin
  { bookmarks   :: !(V.Vector T.Text)
  , tx_timeout  :: !(Maybe Int64)
  , tx_metadata :: !(H.HashMap T.Text Ps)
  , mode        :: !AccessMode
  , db          :: !(Maybe T.Text)
  , imp_user    :: !(Maybe T.Text)
  }
  deriving stock (Generic)
instance PackStream Begin where
  toPs Begin{bookmarks, tx_timeout, tx_metadata, mode, db, imp_user} = PsDictionary $ H.fromList
    [ ("bookmarks", toPs bookmarks)
    , ("tx_timeout", toPs tx_timeout)
    , ("tx_metadata", toPs tx_metadata)
    , ("mode", toPs mode)
    , ("db", toPs db)
    , ("imp_user", toPs imp_user)
    ]
  toBinary run_extra = putPs $ toPs run_extra
  fromPs = withDictionary "Begin" f
    where f map = do
            bookmarks :: V.Vector T.Text <- lookupWithError "bookmarks" map "\"bookmarks\" not found in Begin"
            tx_timeout :: Maybe Int64 <- lookupMaybe "tx_timeout" map
            tx_metadata :: H.HashMap T.Text Ps <- lookupWithError "tx_metadata" map "\"tx_metadata\" not found in Begin"
            mode :: AccessMode <- lookupMode map
            db :: (Maybe T.Text) <- lookupMaybe "db" map
            imp_user :: Maybe T.Text <- lookupMaybe "imp_user" map
            pure $ Begin{bookmarks, tx_timeout, tx_metadata, mode, db, imp_user}



-- | ROUTE message payload for fetching routing tables.
type Route :: Type
data Route = Route
  { routing   :: !(H.HashMap T.Text Ps) -- Only "address" specified
  , bookmarks :: !(V.Vector T.Text)
  , extra     :: RouteExtra
  }

routeToPs :: Route -> Ps
routeToPs Route{routing, bookmarks, extra} =
  PsStructure 0x66 $ V.fromList
    [ toPs routing
    , toPs bookmarks
    , toPs extra
    ]

fromPsToRoute :: V.Vector Ps -> Result Request
fromPsToRoute fields
  | V.length fields /= 3 = Error "expected exactly 3 fields for Route"
  | otherwise = do
      routing   <- fromPs (fields V.! 0)
      bookmarks <- fromPs (fields V.! 1)
      extra     <- fromPs (fields V.! 2)
      Success $ RRoute Route{routing, bookmarks, extra}

-- | Default 'Pull' requesting up to 1000 records per batch, with no explicit query ID.
-- Matches the default fetch size used by the official Neo4j drivers (Java, Python, JS, .NET).
-- Use @n = -1@ to request all records in a single batch (no back-pressure).
defaultPull :: Pull
defaultPull = Pull{n = 1000, qid = Nothing}

-- | PULL message payload with record count and optional query ID.
type Pull :: Type
data Pull = Pull
  { n   :: Int64
  , qid :: Maybe Int64
  }
  deriving stock (Generic)
instance PackStream Pull where
  toPs Pull{n, qid} = PsDictionary $ H.fromList $
    [("n", toPs n)] <>
    case qid of
      Nothing -> []
      Just q  -> [("qid", toPs q)]
  fromPs = withDictionary "Pull" $ \m -> do
    n' <- lookupWithError "n" m "\"n\" not found in Pull"
    qid' <- lookupMaybe "qid" m
    pure $ Pull{n = n', qid = qid'}

-- | Extra parameters for the ROUTE message (database and impersonated user).
type RouteExtra :: Type
data RouteExtra = RouteExtra
  { db       :: !(Maybe T.Text)
  , imp_user :: !(Maybe T.Text)
  }
instance PackStream RouteExtra where
  toPs RouteExtra{db, imp_user} = PsDictionary $ H.fromList $
    maybe [] (\d -> [("db", toPs d)]) db <>
    maybe [] (\i -> [("imp_user", toPs i)]) imp_user
  toBinary route_extra = putPs $ toPs route_extra
  fromPs = withDictionary "RouteExtra" f
    where f map = do
            db :: Maybe T.Text <- lookupMaybe "db" map
            imp_user :: Maybe T.Text <- lookupMaybe "imp_user" map
            pure $ RouteExtra{db, imp_user}


-- | Extra parameters for the RUN message (bookmarks, timeout, metadata, mode, db, impersonation).
type RunExtra :: Type
data RunExtra = RunExtra
  { bookmarks   :: !(V.Vector T.Text)
  , tx_timeout  :: !(Maybe Int64)
  , tx_metadata :: !(H.HashMap T.Text Ps)
  , mode        :: !AccessMode
  , db          :: !(Maybe T.Text)
  , imp_user    :: !(Maybe T.Text)
  }

-- | Default 'RunExtra' with empty bookmarks, no timeout, write mode, and no database.
defaultRunExtra :: RunExtra
defaultRunExtra = RunExtra
  { bookmarks   = V.empty
  , tx_timeout  = Nothing
  , tx_metadata = H.empty
  , mode        = WriteAccess
  , db          = Nothing
  , imp_user    = Nothing
  }

instance PackStream RunExtra where
  toPs RunExtra{bookmarks, tx_timeout, tx_metadata, mode, db, imp_user} = PsDictionary $ H.fromList
    [ ("bookmarks", toPs bookmarks)
    , ("tx_timeout", toPs tx_timeout)
    , ("tx_metadata", toPs tx_metadata)
    , ("mode", toPs mode)
    , ("db", toPs db)
    , ("imp_user", toPs imp_user)
    ]
  toBinary run_extra = putPs $ toPs run_extra
  fromPs = withDictionary "RunExtra" f
    where f map = do
            bookmarks :: V.Vector T.Text <- lookupWithError "bookmarks" map "\"bookmarks\" not found in RunExtra"
            tx_timeout :: Maybe Int64 <- lookupMaybe "tx_timeout" map
            tx_metadata :: H.HashMap T.Text Ps <- lookupWithError "tx_metadata" map "\"tx_metadata\" not found in RunExtra"
            mode :: AccessMode <- lookupMode map
            db :: Maybe T.Text <- lookupMaybe "db" map
            imp_user :: Maybe T.Text <- lookupMaybe "imp_user" map
            pure $ RunExtra{bookmarks, tx_timeout, tx_metadata, mode, db, imp_user}


helloToPs :: Hello -> Ps
helloToPs Hello{user_agent, scheme, routing, patchBolt} =
    PsDictionary $ H.fromList $
      [ ("user_agent", toPs user_agent)
      , ("scheme"    , toPs scheme')
      ] <>
        ( case scheme of
            None                        -> []
            Basic principal credentials -> [("principal", toPs principal), ("credentials", toPs credentials)]
            Kerberos                    -> []
            Bearer credentials          -> [("credentials", toPs credentials)]
        ) <>
        ( case routing of
            NoRouting                          -> []
            Routing                            -> [("routing", PsDictionary H.empty)]
            RoutingSpec address urlQueryParams -> [("routing", PsDictionary $ H.fromList $
                ("address", toPs address) : [(k, toPs v) | (k, v) <- H.toList urlQueryParams])]
        ) <>
        ( if patchBolt then [("patch_bolt", PsList $ V.singleton $ PsString "utc")] else [] )
  where
    scheme' :: T.Text
    scheme' = case scheme of
                None      -> "none"
                Basic _ _ -> "basic"
                Kerberos  -> "kerberos"
                Bearer _  -> "bearer"

helloFromPs :: Ps -> Result Hello
helloFromPs = withDictionary "Request:Hello" f
  where f :: H.HashMap T.Text Ps -> Result Hello
        f map = do
          user_agent :: UserAgent <- lookupWithError "user_agent" map "\"user_agent\" not found in hello request"
          scheme :: T.Text <- lookupWithError "scheme" map "\"scheme\" not found in hello request"
          routing <- case H.lookup "routing" map of
                       Nothing -> Success NoRouting
                       Just (PsDictionary m)
                         | H.null m  -> Success Routing
                         | otherwise -> case H.lookup "address" m of
                             Just addr -> do
                               a <- fromPs addr
                               let params = H.delete "address" m
                               ps <- traverse fromPs params
                               Success $ RoutingSpec a ps
                             Nothing -> Error "\"routing\" dictionary missing \"address\" key"
                       Just _ -> Error "\"routing\" field should be a dictionary"
          let credentials = H.lookup "credentials" map
          let principal   = H.lookup "principal" map
          let pb = case H.lookup "patch_bolt" map of
                     Just (PsList v) -> V.elem (PsString "utc") v
                     _               -> False
          case scheme of
            "none"     -> pure $ Hello {user_agent, scheme = None, routing, patchBolt = pb}
            "basic"    -> case (principal, credentials) of
              (Just p, Just c) -> do
                p2 <- fromPs p
                c2 <- fromPs c
                pure $ Hello {user_agent, scheme = Basic p2 c2, routing, patchBolt = pb}
              _ -> Error "\"principal\" and \"credentials\" required for scheme \"basic\""
            "kerberos" -> pure $ Hello {user_agent, scheme = Kerberos, routing, patchBolt = pb}
            "bearer"   -> case credentials of
              Just c  -> do
                c2 <- fromPs c
                pure $ Hello {user_agent, scheme = Bearer c2, routing, patchBolt = pb}
              Nothing -> Error "\"credentials\" required for scheme \"bearer\""
            _          -> Error $ "unknown scheme: " <> scheme


-- | LOGON message carries authentication credentials (Bolt 5.1+).
-- In 5.1+ HELLO no longer carries credentials; auth moves to LOGON.
type Logon :: Type
data Logon = Logon
  { scheme :: Scheme
  }
  deriving stock (Generic)

instance PackStream Logon where
  toPs = logonToPs
  fromPs = logonFromPs

logonToPs :: Logon -> Ps
logonToPs Logon{scheme} =
    PsDictionary $ H.fromList $
      [ ("scheme", toPs scheme')
      ] <>
        ( case scheme of
            None                        -> []
            Basic principal credentials -> [("principal", toPs principal), ("credentials", toPs credentials)]
            Kerberos                    -> []
            Bearer credentials          -> [("credentials", toPs credentials)]
        )
  where
    scheme' :: T.Text
    scheme' = case scheme of
                None      -> "none"
                Basic _ _ -> "basic"
                Kerberos  -> "kerberos"
                Bearer _  -> "bearer"

logonFromPs :: Ps -> Result Logon
logonFromPs = withDictionary "Request:Logon" f
  where f :: H.HashMap T.Text Ps -> Result Logon
        f map = do
          schemeTxt :: T.Text <- lookupWithError "scheme" map "\"scheme\" not found in logon request"
          let credentials = H.lookup "credentials" map
          let principal   = H.lookup "principal" map
          case schemeTxt of
            "none"     -> pure $ Logon{scheme = None}
            "basic"    -> case (principal, credentials) of
              (Just p, Just c) -> do
                p2 <- fromPs p
                c2 <- fromPs c
                pure $ Logon{scheme = Basic p2 c2}
              _ -> Error "\"principal\" and \"credentials\" required for scheme \"basic\""
            "kerberos" -> pure $ Logon{scheme = Kerberos}
            "bearer"   -> case credentials of
              Just c  -> do
                c2 <- fromPs c
                pure $ Logon{scheme = Bearer c2}
              Nothing -> Error "\"credentials\" required for scheme \"bearer\""
            _          -> Error $ "unknown scheme: " <> schemeTxt


-- | Telemetry API type indicators (Bolt 5.4+).
type TelemetryApi :: Type
data TelemetryApi
  = ManagedTransactions
  | ExplicitTransactions
  | ImplicitTransactions
  | ExecuteQuery
  deriving stock (Show, Eq)

telemetryApiToInt :: TelemetryApi -> Int64
telemetryApiToInt ManagedTransactions  = 0
telemetryApiToInt ExplicitTransactions = 1
telemetryApiToInt ImplicitTransactions = 2
telemetryApiToInt ExecuteQuery         = 3

telemetryApiFromInt :: Int64 -> Result TelemetryApi
telemetryApiFromInt 0 = Success ManagedTransactions
telemetryApiFromInt 1 = Success ExplicitTransactions
telemetryApiFromInt 2 = Success ImplicitTransactions
telemetryApiFromInt 3 = Success ExecuteQuery
telemetryApiFromInt n = Error $ "unknown telemetry api value: " <> T.pack (show n)

fromPsToTelemetry :: V.Vector Ps -> Result Request
fromPsToTelemetry fields
  | V.length fields /= 1 = Error "expected exactly one field for Telemetry"
  | otherwise = case V.unsafeHead fields of
      PsInteger n -> case fromPSInteger n of
        Just i  -> fmap RTelemetry $ telemetryApiFromInt i
        Nothing -> Error "telemetry api integer out of Int64 range"
      _           -> Error "expected integer field for Telemetry"


-- | Look up the BEGIN\/RUN @mode@ field, defaulting to 'WriteAccess' when
-- absent. Per the Bolt protocol, @mode@ is an optional field with default
-- @"w"@; spec-compliant peers may omit it on write transactions to save
-- bytes. A bare @lookupWithError@ would reject those messages.
lookupMode :: H.HashMap T.Text Ps -> Result AccessMode
lookupMode m = fromMaybe WriteAccess <$> lookupMaybe "mode" m