-- | Typed plan and profile trees from EXPLAIN\/PROFILE queries.
module Database.Bolty.Plan
  ( PlanNode(..)
  , ProfileNode(..)
  , parsePlan
  , parseProfile
  ) 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
  { PlanNode -> Text
pnOperatorType :: !T.Text
  , PlanNode -> HashMap Text Ps
pnArguments    :: !(H.HashMap T.Text Ps)
  , PlanNode -> Vector Text
pnIdentifiers  :: !(V.Vector T.Text)
  , PlanNode -> Double
pnEstimatedRows :: !Double
  , PlanNode -> Vector PlanNode
pnChildren     :: !(V.Vector PlanNode)
  } deriving stock (Int -> PlanNode -> ShowS
[PlanNode] -> ShowS
PlanNode -> String
(Int -> PlanNode -> ShowS)
-> (PlanNode -> String) -> ([PlanNode] -> ShowS) -> Show PlanNode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PlanNode -> ShowS
showsPrec :: Int -> PlanNode -> ShowS
$cshow :: PlanNode -> String
show :: PlanNode -> String
$cshowList :: [PlanNode] -> ShowS
showList :: [PlanNode] -> ShowS
Show, PlanNode -> PlanNode -> Bool
(PlanNode -> PlanNode -> Bool)
-> (PlanNode -> PlanNode -> Bool) -> Eq PlanNode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PlanNode -> PlanNode -> Bool
== :: PlanNode -> PlanNode -> Bool
$c/= :: PlanNode -> PlanNode -> Bool
/= :: PlanNode -> PlanNode -> Bool
Eq)


-- | A node in the query profile tree (from PROFILE).
-- Extends 'PlanNode' with actual execution statistics.
type ProfileNode :: Type
data ProfileNode = ProfileNode
  { ProfileNode -> Text
prOperatorType    :: !T.Text
  , ProfileNode -> HashMap Text Ps
prArguments       :: !(H.HashMap T.Text Ps)
  , ProfileNode -> Vector Text
prIdentifiers     :: !(V.Vector T.Text)
  , ProfileNode -> Double
prEstimatedRows   :: !Double
  , ProfileNode -> Int64
prDbHits          :: !Int64
  , ProfileNode -> Int64
prRows            :: !Int64
  , ProfileNode -> Int64
prPageCacheHits   :: !Int64
  , ProfileNode -> Int64
prPageCacheMisses :: !Int64
  , ProfileNode -> Int64
prTime            :: !Int64              -- ^ microseconds
  , ProfileNode -> Vector ProfileNode
prChildren        :: !(V.Vector ProfileNode)
  } deriving stock (Int -> ProfileNode -> ShowS
[ProfileNode] -> ShowS
ProfileNode -> String
(Int -> ProfileNode -> ShowS)
-> (ProfileNode -> String)
-> ([ProfileNode] -> ShowS)
-> Show ProfileNode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ProfileNode -> ShowS
showsPrec :: Int -> ProfileNode -> ShowS
$cshow :: ProfileNode -> String
show :: ProfileNode -> String
$cshowList :: [ProfileNode] -> ShowS
showList :: [ProfileNode] -> ShowS
Show, ProfileNode -> ProfileNode -> Bool
(ProfileNode -> ProfileNode -> Bool)
-> (ProfileNode -> ProfileNode -> Bool) -> Eq ProfileNode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ProfileNode -> ProfileNode -> Bool
== :: ProfileNode -> ProfileNode -> Bool
$c/= :: ProfileNode -> ProfileNode -> Bool
/= :: ProfileNode -> ProfileNode -> Bool
Eq)


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


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


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


parsePlanPs :: Ps -> Maybe PlanNode
parsePlanPs :: Ps -> Maybe PlanNode
parsePlanPs (PsDictionary HashMap Text Ps
m) = HashMap Text Ps -> Maybe PlanNode
parsePlanDict HashMap Text Ps
m
parsePlanPs Ps
_ = Maybe PlanNode
forall a. Maybe a
Nothing


parseProfileDict :: H.HashMap T.Text Ps -> Maybe ProfileNode
parseProfileDict :: HashMap Text Ps -> Maybe ProfileNode
parseProfileDict HashMap Text Ps
m = do
  opType <- Text -> HashMap Text Ps -> Maybe Text
lookupText Text
"operatorType" HashMap Text Ps
m
  let args = case Text -> HashMap Text Ps -> Maybe Ps
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
H.lookup Text
"args" HashMap Text Ps
m of
               Just (PsDictionary HashMap Text Ps
a) -> HashMap Text Ps
a
               Maybe Ps
_                     -> HashMap Text Ps
forall k v. HashMap k v
H.empty
  let idents = case Text -> HashMap Text Ps -> Maybe Ps
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
H.lookup Text
"identifiers" HashMap Text Ps
m of
                 Just (PsList Vector Ps
v) -> (Ps -> Maybe Text) -> Vector Ps -> Vector Text
forall a b. (a -> Maybe b) -> Vector a -> Vector b
V.mapMaybe Ps -> Maybe Text
extractText Vector Ps
v
                 Maybe Ps
_               -> Vector Text
forall a. Vector a
V.empty
  let estRows = Text -> HashMap Text Ps -> Double
lookupDouble Text
"estimatedRows" HashMap Text Ps
m
  let dbHits = Text -> HashMap Text Ps -> Int64
lookupInt64OrZero Text
"dbHits" HashMap Text Ps
m
  let rows = Text -> HashMap Text Ps -> Int64
lookupInt64OrZero Text
"rows" HashMap Text Ps
m
  let pcHits = Text -> HashMap Text Ps -> Int64
lookupInt64OrZero Text
"pageCacheHits" HashMap Text Ps
m
  let pcMisses = Text -> HashMap Text Ps -> Int64
lookupInt64OrZero Text
"pageCacheMisses" HashMap Text Ps
m
  let time = Text -> HashMap Text Ps -> Int64
lookupInt64OrZero Text
"time" HashMap Text Ps
m
  let children = case Text -> HashMap Text Ps -> Maybe Ps
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
H.lookup Text
"children" HashMap Text Ps
m of
                   Just (PsList Vector Ps
v) -> (Ps -> Maybe ProfileNode) -> Vector Ps -> Vector ProfileNode
forall a b. (a -> Maybe b) -> Vector a -> Vector b
V.mapMaybe Ps -> Maybe ProfileNode
parseProfilePs Vector Ps
v
                   Maybe Ps
_               -> Vector ProfileNode
forall a. Vector a
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 :: Ps -> Maybe ProfileNode
parseProfilePs (PsDictionary HashMap Text Ps
m) = HashMap Text Ps -> Maybe ProfileNode
parseProfileDict HashMap Text Ps
m
parseProfilePs Ps
_ = Maybe ProfileNode
forall a. Maybe a
Nothing


-- Helpers

lookupText :: T.Text -> H.HashMap T.Text Ps -> Maybe T.Text
lookupText :: Text -> HashMap Text Ps -> Maybe Text
lookupText Text
key HashMap Text Ps
m = case Text -> HashMap Text Ps -> Maybe Ps
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
H.lookup Text
key HashMap Text Ps
m of
  Just (PsString Text
t) -> Text -> Maybe Text
forall a. a -> Maybe a
Just Text
t
  Maybe Ps
_                 -> Maybe Text
forall a. Maybe a
Nothing


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


lookupDouble :: T.Text -> H.HashMap T.Text Ps -> Double
lookupDouble :: Text -> HashMap Text Ps -> Double
lookupDouble Text
key HashMap Text Ps
m = case Text -> HashMap Text Ps -> Maybe Ps
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
H.lookup Text
key HashMap Text Ps
m of
  Just (PsFloat Double
d) -> Double
d
  Just (PsInteger PSInteger
n) -> case PSInteger -> Maybe Int64
forall a. FromPSInteger a => PSInteger -> Maybe a
fromPSInteger PSInteger
n of
    Just (Int64
i :: Int64) -> Int64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
i
    Maybe Int64
Nothing           -> Double
0.0
  Maybe Ps
_ -> Double
0.0


lookupInt64OrZero :: T.Text -> H.HashMap T.Text Ps -> Int64
lookupInt64OrZero :: Text -> HashMap Text Ps -> Int64
lookupInt64OrZero Text
key HashMap Text Ps
m = case Text -> HashMap Text Ps -> Maybe Ps
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
H.lookup Text
key HashMap Text Ps
m of
  Just (PsInteger PSInteger
n) -> case PSInteger -> Maybe Int64
forall a. FromPSInteger a => PSInteger -> Maybe a
fromPSInteger PSInteger
n of
    Just Int64
i  -> Int64
i
    Maybe Int64
Nothing -> Int64
0
  Maybe Ps
_ -> Int64
0