Response.hs

{-# LANGUAGE ApplicativeDo #-}

-- | Internal module. Not part of the public API.
module Database.Bolty.Message.Response
  ( Response(..)
  , Failure(..)
  , SuccessHello(..)
  , SuccessRun(..), successFields, successTFirst
  , SuccessRunAutoCommitTransaction(..)
  , SuccessRunExplicitTransaction(..)
  , makeResponseRunAutoCommitTransaction
  , SuccessDiscard(..), SuccessDiscardLast(..)
  , SuccessPull(..), QueryMeta(..)
  , makeSuccessPull, mapToSuccessPull
  , SuccessCommit(..)
  , RoutingTable(..)
  , parseRoutingTable
  , extractBookmark
  ) where

import           Data.Kind                     (Type)
import           Data.Int                      (Int64)
import qualified Data.HashMap.Lazy             as H
import qualified Data.Text                     as T
import qualified Data.Vector                   as V

import qualified Data.PackStream               as PS
import           Data.PackStream.Ps            as PS
import           Data.PackStream.Integer       (fromPSInteger)
import           Database.Bolty.Notification   (Notification, parseNotifications)
import           Database.Bolty.Plan           (PlanNode, ProfileNode, parsePlan, parseProfile)
import           Database.Bolty.Stats          (QueryStats, parseStats)
import           Database.Bolty.Record         (Record)

-- | A BOLT server response message.
--
-- the server always responds with one summary message if the connection is still open
-- the server always responds with zero or more detail messages before sending a summary message
type Response :: Type
data Response
  = RSuccess !(H.HashMap T.Text Ps)
  -- ^ https://neo4j.com/docs/bolt/current/bolt/message/#messages-success
  | RIgnored
  -- ^ https://neo4j.com/docs/bolt/current/bolt/message/#messages-ignored
  | RFailure Failure
  -- ^ https://neo4j.com/docs/bolt/current/bolt/message/#messages-failure
  | RRecord !Record
  -- ^ https://neo4j.com/docs/bolt/current/bolt/message/#messages-record
instance PS.PackStream Response where
  toPs (RSuccess success) = structureSingleton 0x70 $ toPs success
  toPs (RIgnored) = PsStructure 0x7E V.empty
  toPs (RFailure failure) = structureSingleton 0x7F $ toPs failure
  toPs (RRecord records) = structureSingleton 0x71 $ toPs records
  fromPs (PsStructure 0x70 fields) = PS.fromOneField fields RSuccess "expected exactly one field for Success"
  fromPs (PsStructure 0x7E fields) = if V.null fields then PS.Success RIgnored else PS.Error "additional data on Ignored request is not allowed"
  fromPs (PsStructure 0x7F fields) = PS.fromOneField fields RFailure "expected exactly one field for Failure"
  fromPs (PsStructure 0x71 fields) = PS.fromOneField fields RRecord "expected exactly one field for Record"
  fromPs _                         = PS.Error "Could not decode Response"



-- | Parse a RUN SUCCESS response for an auto-commit transaction.
makeResponseRunAutoCommitTransaction :: H.HashMap T.Text Ps -> Either T.Text SuccessRun
makeResponseRunAutoCommitTransaction map = do
  fields <- PS.lookupMaybeError "fields" map
  t_first <- PS.lookupMaybeError "t_first" map
  pure $ AutoCommitTransaction $ SuccessRunAutoCommitTransaction fields t_first



-- | Build a 'SuccessPull' from collected records and PULL SUCCESS metadata (Either version).
makeSuccessPull :: V.Vector Record -> H.HashMap T.Text Ps -> Either T.Text SuccessPull
makeSuccessPull records map = do
  let bookmark = case PS.lookupMaybeError "bookmark" map of
                   Right b -> Just b
                   Left _  -> Nothing
  t_last <- PS.lookupMaybeError "t_last" map
  type_ <- PS.lookupMaybeError "type" map
  db <- PS.lookupMaybeError "db" map
  let stats = H.lookup "stats" map
  let parsedStats = parseStats stats
  let plan = H.lookup "plan" map
  let profile = H.lookup "profile" map
  let notifications = H.lookup "notifications" map
  let parsedNotifications = parseNotifications notifications
  let parsedPlan = parsePlan plan
  let parsedProfile = parseProfile profile
  let infos = QueryMeta{bookmark, t_last, type_, stats, parsedStats, plan, profile, notifications, parsedNotifications, parsedPlan, parsedProfile, db}
  pure $ SuccessPull{records, infos}





-- | A server-reported failure with a Neo4j error code and human-readable message.
type Failure :: Type
data Failure = Failure
  { code :: !T.Text
  , message :: !T.Text
  }
instance PS.PackStream Failure where
  toPs (Failure code message) = PsDictionary $ H.fromList [("code", toPs code), ("message", toPs message)]
  fromPs (PsDictionary map) = do
    code <- PS.lookupWithError "code" map "expected \"code\" in Failure response"
    message <- PS.lookupWithError "message" map "expected \"message\" in Failure response"
    PS.Success $ Failure code message
  fromPs _ = PS.Error "expected dictionary for Failure"




-- | Parsed HELLO SUCCESS metadata.
type SuccessHello :: Type
data SuccessHello = SuccessHello
  { server :: !T.Text
  -- ^ server agent string, example "Neo4j/4.1.0"
  , connection_id :: !T.Text
  -- ^ unique identifier of the bolt connection used on the server side, example: "bolt-61"
  , hints :: !(H.HashMap T.Text Ps)
  -- ^ set of optional configuration hints to be considered by the driver
  }

-- | Parsed RUN SUCCESS metadata, either auto-commit or explicit transaction.
type SuccessRun :: Type
data SuccessRun
  = AutoCommitTransaction SuccessRunAutoCommitTransaction
  | ExplicitTransaction SuccessRunExplicitTransaction

-- | Extract the result field names from a 'SuccessRun'.
successFields :: SuccessRun -> V.Vector T.Text
successFields (AutoCommitTransaction (SuccessRunAutoCommitTransaction fields _)) = fields
successFields (ExplicitTransaction (SuccessRunExplicitTransaction fields _ _)) = fields

-- | Extract the @t_first@ timing value from a 'SuccessRun'.
successTFirst :: SuccessRun -> Int64
successTFirst (AutoCommitTransaction (SuccessRunAutoCommitTransaction _ tf)) = tf
successTFirst (ExplicitTransaction (SuccessRunExplicitTransaction _ tf _)) = tf

-- | RUN SUCCESS metadata for an auto-commit transaction.
type SuccessRunAutoCommitTransaction :: Type
data SuccessRunAutoCommitTransaction = SuccessRunAutoCommitTransaction
  { fields :: !(V.Vector T.Text)
  -- ^ the fields of the return result. e.g. [“name”, “age”, …]
  , t_first :: !Int64
  -- ^ the time, specified in ms, which the first record in the result stream is available after
  }

-- | RUN SUCCESS metadata for an explicit transaction.
type SuccessRunExplicitTransaction :: Type
data SuccessRunExplicitTransaction = SuccessRunExplicitTransaction
  { fields :: !(V.Vector T.Text)
  -- ^ the fields of the return result. e.g. [“name”, “age”, …]
  , t_first :: !Int64
  -- ^ the time, specified in ms, which the first record in the result stream is available after
  , qid :: !Int64
  -- ^ specifies the server assigned statement ID to reference the server side result-set with commencing BEGIN+RUN+PULL and BEGIN+RUN+DISCARD messages.
  }

-- | Parsed DISCARD SUCCESS response.
type SuccessDiscard :: Type
data SuccessDiscard
  = DiscardMore
  | DiscardLast SuccessDiscardLast

-- | Final DISCARD SUCCESS metadata with bookmark and database.
type SuccessDiscardLast :: Type
data SuccessDiscardLast = SuccessDiscardLast
  { bookmark :: !T.Text
  -- ^ the bookmark after committing this transaction (Auto-commit Transaction only)
  , db :: !T.Text
  -- ^ the database name where the query is executed.
  }

-- | Build a 'SuccessPull' from collected records and PULL SUCCESS metadata (Result version).
mapToSuccessPull :: V.Vector Record -> H.HashMap T.Text Ps -> PS.Result SuccessPull
mapToSuccessPull records map = do
  bookmark :: Maybe T.Text <- PS.lookupMaybe "bookmark" map
  t_last :: Int64    <- PS.lookupWithError "t_last" map "could not find \"t_last\" in Pull response"
  type_ :: T.Text    <- PS.lookupWithError "type" map "could not find \"type\" in Pull response"
  db :: T.Text       <- PS.lookupWithError "db" map "could not find \"db\" in Pull response"
  let stats         = H.lookup "stats" map
  let parsedStats   = parseStats stats
  let plan          = H.lookup "plan" map
  let profile       = H.lookup "profile" map
  let notifications = H.lookup "notifications" map
  let parsedNotifications = parseNotifications notifications
  let parsedPlan    = parsePlan plan
  let parsedProfile = parseProfile profile
  let infos = QueryMeta{bookmark, t_last, type_, stats, parsedStats, plan, profile, notifications, parsedNotifications, parsedPlan, parsedProfile, db}
  pure $ SuccessPull{records, infos}




-- | Full query result: records plus server metadata.
type SuccessPull :: Type
data SuccessPull = SuccessPull
  { records :: !(V.Vector Record)
  , infos   :: !QueryMeta
  }

-- | Server metadata from a PULL SUCCESS response.
type QueryMeta :: Type
data QueryMeta = QueryMeta
  { bookmark :: !(Maybe T.Text)
  -- ^ the bookmark after committing this transaction (Autocommit Transaction only, absent in explicit transactions).
  , t_last :: !Int64
  -- ^ the time, specified in ms, which the last record in the result stream is consumed after.
  , type_ :: !T.Text
  -- ^ the type of the statement, e.g. "r" for read-only statement, "w" for write-only statement, "rw" for read-and-write, and "s" for schema only.
  , stats :: !(Maybe Ps)
  -- ^ counter information, such as db-hits etc. May be omitted.
  , parsedStats :: !(Maybe QueryStats)
  -- ^ parsed stats from the raw field above.
  , plan :: !(Maybe Ps)
  -- ^ plan result. May be omitted.
  , profile :: !(Maybe Ps)
  -- ^ profile result. May be omitted.
  , notifications :: !(Maybe Ps)
  -- ^ a list of all notifications generated during execution of this statement. May be omitted.
  , parsedNotifications :: !(V.Vector Notification)
  -- ^ parsed notifications from the raw field above.
  , parsedPlan :: !(Maybe PlanNode)
  -- ^ parsed plan tree from the raw @plan@ field above.
  , parsedProfile :: !(Maybe ProfileNode)
  -- ^ parsed profile tree from the raw @profile@ field above.
  , db :: !T.Text
  -- ^ the database name where the query was executed (v4.0+).
  }
  deriving stock (Show)

-- | Parsed COMMIT SUCCESS metadata.
type SuccessCommit :: Type
data SuccessCommit = SuccessCommit
  { bookmark :: !T.Text
  -- ^ the bookmark after committing this transaction.
  }


-- | A parsed routing table returned by the ROUTE message.
type RoutingTable :: Type
data RoutingTable = RoutingTable
  { ttl     :: !Int64
  -- ^ Seconds to cache this routing table.
  , db      :: !T.Text
  -- ^ Database name this routing table is for.
  , routers :: !(V.Vector T.Text)
  -- ^ "host:port" addresses for ROUTE role.
  , readers :: !(V.Vector T.Text)
  -- ^ "host:port" addresses for READ role.
  , writers :: !(V.Vector T.Text)
  -- ^ "host:port" addresses for WRITE role.
  }
  deriving stock (Show, Eq)


-- | Parse a routing table from the SUCCESS response metadata of a ROUTE message.
-- The expected wire format is:
-- @{"rt": {"ttl": N, "db": "name", "servers": [{"addresses": [...], "role": "WRITE"}, ...]}}@
parseRoutingTable :: H.HashMap T.Text Ps -> Either T.Text RoutingTable
parseRoutingTable meta = do
  rtPs <- case H.lookup "rt" meta of
    Just (PsDictionary m) -> Right m
    Just _  -> Left "\"rt\" field should be a dictionary"
    Nothing -> Left "\"rt\" key not found in ROUTE response"
  ttl <- case H.lookup "ttl" rtPs of
    Just (PsInteger n) -> case fromPSInteger n of
      Just i  -> Right i
      Nothing -> Left "\"ttl\" integer out of Int64 range"
    Just _  -> Left "\"ttl\" should be an integer"
    Nothing -> Left "\"ttl\" not found in routing table"
  db <- case H.lookup "db" rtPs of
    Just (PsString s) -> Right s
    Just _  -> Left "\"db\" should be a string"
    Nothing -> Left "\"db\" not found in routing table"
  serversPs <- case H.lookup "servers" rtPs of
    Just (PsList v) -> Right v
    Just _  -> Left "\"servers\" should be a list"
    Nothing -> Left "\"servers\" not found in routing table"
  let (routers, readers, writers) = V.foldl' classifyServer (V.empty, V.empty, V.empty) serversPs
  Right RoutingTable{ttl, db, routers, readers, writers}


-- | Classify a single server entry from the routing table by its role.
classifyServer :: (V.Vector T.Text, V.Vector T.Text, V.Vector T.Text)
               -> Ps
               -> (V.Vector T.Text, V.Vector T.Text, V.Vector T.Text)
classifyServer (rout, rd, wr) (PsDictionary m) =
  case (H.lookup "role" m, H.lookup "addresses" m) of
    (Just (PsString role), Just (PsList addrs)) ->
      let addrTexts = V.mapMaybe extractText addrs
      in case role of
           "ROUTE" -> (rout <> addrTexts, rd, wr)
           "READ"  -> (rout, rd <> addrTexts, wr)
           "WRITE" -> (rout, rd, wr <> addrTexts)
           _       -> (rout, rd, wr)
    _ -> (rout, rd, wr)
classifyServer acc _ = acc


extractText :: Ps -> Maybe T.Text
extractText (PsString t) = Just t
extractText _            = Nothing


-- | Extract the bookmark from a SUCCESS response metadata (e.g. from COMMIT).
-- Returns 'Nothing' if the @bookmark@ key is absent or not a text value.
extractBookmark :: H.HashMap T.Text Ps -> Maybe T.Text
extractBookmark meta = case H.lookup "bookmark" meta of
  Just (PsString t) -> Just t
  _                 -> Nothing