Plan.hs

-- | Typed plan and profile trees from EXPLAIN\/PROFILE queries.
module Database.Bolty.Plan
  ( PlanNode(..)
  , ProfileNode(..)
  , parsePlan
  , parseProfile
  , renderPlan
  ) where

import           Prelude

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

import           Data.PackStream.Ps            (Ps(..))
import           Data.PackStream.Integer       (fromPSInteger)


-- | A node in the query execution plan tree (from EXPLAIN).
type PlanNode :: Type
data PlanNode = PlanNode
  { pnOperatorType :: !T.Text
  , pnArguments    :: !(H.HashMap T.Text Ps)
  , pnIdentifiers  :: !(V.Vector T.Text)
  , pnEstimatedRows :: !Double
  , pnChildren     :: !(V.Vector PlanNode)
  } deriving stock (Show, Eq)


-- | A node in the query profile tree (from PROFILE).
-- Extends 'PlanNode' with actual execution statistics.
type ProfileNode :: Type
data ProfileNode = ProfileNode
  { prOperatorType    :: !T.Text
  , prArguments       :: !(H.HashMap T.Text Ps)
  , prIdentifiers     :: !(V.Vector T.Text)
  , prEstimatedRows   :: !Double
  , prDbHits          :: !Int64
  , prRows            :: !Int64
  , prPageCacheHits   :: !Int64
  , prPageCacheMisses :: !Int64
  , prTime            :: !Int64              -- ^ microseconds
  , prChildren        :: !(V.Vector ProfileNode)
  } deriving stock (Show, Eq)


-- | Parse the raw @plan@ field from PULL SUCCESS metadata into a typed 'PlanNode' tree.
parsePlan :: Maybe Ps -> Maybe PlanNode
parsePlan Nothing = Nothing
parsePlan (Just (PsDictionary m)) = parsePlanDict m
parsePlan (Just _) = Nothing


-- | Parse the raw @profile@ field from PULL SUCCESS metadata into a typed 'ProfileNode' tree.
parseProfile :: Maybe Ps -> Maybe ProfileNode
parseProfile Nothing = Nothing
parseProfile (Just (PsDictionary m)) = parseProfileDict m
parseProfile (Just _) = Nothing


parsePlanDict :: H.HashMap T.Text Ps -> Maybe PlanNode
parsePlanDict m = do
  opType <- lookupText "operatorType" m
  let args = case H.lookup "args" m of
               Just (PsDictionary a) -> a
               _                     -> H.empty
  let idents = case H.lookup "identifiers" m of
                 Just (PsList v) -> V.mapMaybe extractText v
                 _               -> V.empty
  let estRows = lookupDouble "estimatedRows" m
  let children = case H.lookup "children" m of
                   Just (PsList v) -> V.mapMaybe parsePlanPs v
                   _               -> V.empty
  Just PlanNode
    { pnOperatorType  = opType
    , pnArguments     = args
    , pnIdentifiers   = idents
    , pnEstimatedRows = estRows
    , pnChildren      = children
    }


parsePlanPs :: Ps -> Maybe PlanNode
parsePlanPs (PsDictionary m) = parsePlanDict m
parsePlanPs _ = Nothing


parseProfileDict :: H.HashMap T.Text Ps -> Maybe ProfileNode
parseProfileDict m = do
  opType <- lookupText "operatorType" m
  let args = case H.lookup "args" m of
               Just (PsDictionary a) -> a
               _                     -> H.empty
  let idents = case H.lookup "identifiers" m of
                 Just (PsList v) -> V.mapMaybe extractText v
                 _               -> V.empty
  let estRows = lookupDouble "estimatedRows" m
  let dbHits = lookupInt64OrZero "dbHits" m
  let rows = lookupInt64OrZero "rows" m
  let pcHits = lookupInt64OrZero "pageCacheHits" m
  let pcMisses = lookupInt64OrZero "pageCacheMisses" m
  let time = lookupInt64OrZero "time" m
  let children = case H.lookup "children" m of
                   Just (PsList v) -> V.mapMaybe parseProfilePs v
                   _               -> V.empty
  Just ProfileNode
    { prOperatorType    = opType
    , prArguments       = args
    , prIdentifiers     = idents
    , prEstimatedRows   = estRows
    , prDbHits          = dbHits
    , prRows            = rows
    , prPageCacheHits   = pcHits
    , prPageCacheMisses = pcMisses
    , prTime            = time
    , prChildren        = children
    }


parseProfilePs :: Ps -> Maybe ProfileNode
parseProfilePs (PsDictionary m) = parseProfileDict m
parseProfilePs _ = Nothing


-- Helpers

lookupText :: T.Text -> H.HashMap T.Text Ps -> Maybe T.Text
lookupText key m = case H.lookup key m of
  Just (PsString t) -> Just t
  _                 -> Nothing


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


lookupDouble :: T.Text -> H.HashMap T.Text Ps -> Double
lookupDouble key m = case H.lookup key m of
  Just (PsFloat d) -> d
  Just (PsInteger n) -> case fromPSInteger n of
    Just (i :: Int64) -> fromIntegral i
    Nothing           -> 0.0
  _ -> 0.0


lookupInt64OrZero :: T.Text -> H.HashMap T.Text Ps -> Int64
lookupInt64OrZero key m = case H.lookup key m of
  Just (PsInteger n) -> case fromPSInteger n of
    Just i  -> i
    Nothing -> 0
  _ -> 0


-- | Render a 'PlanNode' tree as an indented text block. Each child is
-- indented two spaces deeper than its parent. Identifier vectors (the
-- @vars:@ line) are shown when non-empty so a reader can match
-- variables to operators.
--
-- Operator arguments ('pnArguments') are intentionally dropped for
-- terseness -- the operator name + estimated rows + identifiers are
-- enough to read most plans, and the args dictionary is verbose.
--
-- Pass @0@ for the initial @depth@. Suitable for printing the result
-- of 'queryExplain'.
--
-- Example output:
--
-- > + ProduceResults  (~1.0 rows)
-- >     vars: u, name
-- >   + Projection  (~1.0 rows)
-- >       vars: u, name
-- >     + NodeUniqueIndexSeek  (~1.0 rows)
-- >         vars: u
renderPlan :: Int -> PlanNode -> T.Text
renderPlan depth (PlanNode op _args idents est children) =
  let indent = T.replicate depth "  "
      header = indent <> "+ " <> op
             <> "  (~" <> T.pack (show est) <> " rows)"
      idLine = if V.null idents
                 then ""
                 else indent <> "    vars: " <> T.intercalate ", " (V.toList idents) <> "\n"
      kids   = T.concat $ V.toList $ V.map (renderPlan (depth + 1)) children
  in header <> "\n" <> idLine <> kids