-- | Postgres Execute Prepare
--
-- Deals with translating (session) variables to SQL expressions. Uses a state
-- monad to keep track of things like variables and generating fresh variable
-- names.
--
-- See 'Hasura.Backends.Postgres.Instances.Execute'.
module Hasura.Backends.Postgres.Execute.Prepare
  ( PlanVariables,
    PrepArgMap,
    PlanningSt (..),
    ExecutionPlan,
    ExecutionStep (..),
    initPlanningSt,
    prepareWithPlan,
    prepareWithoutPlan,
    withUserVars,
  )
where

import Data.Aeson qualified as J
import Data.HashMap.Strict qualified as Map
import Data.IntMap qualified as IntMap
import Data.Text.Extended
import Database.PG.Query qualified as Q
import Hasura.Backends.Postgres.Connection.MonadTx
import Hasura.Backends.Postgres.SQL.DML qualified as S
import Hasura.Backends.Postgres.SQL.Value
import Hasura.Backends.Postgres.Translate.Column
import Hasura.Backends.Postgres.Types.Column
import Hasura.Base.Error
  ( Code (NotFound),
    QErr,
    throw400,
  )
import Hasura.GraphQL.Execute.Backend
import Hasura.GraphQL.Parser.Names
import Hasura.Prelude
import Hasura.RQL.IR.Value
import Hasura.RQL.Types.Column
import Hasura.SQL.Backend
import Hasura.Session
  ( SessionVariables,
    UserInfo (_uiSession),
    getSessionVariableValue,
    sessionVariableToText,
  )
import Language.GraphQL.Draft.Syntax qualified as G

type PlanVariables = Map.HashMap G.Name Int

-- | The value is (Q.PrepArg, PGScalarValue) because we want to log the human-readable value of the
-- prepared argument and not the binary encoding in PG format
type PrepArgMap = IntMap.IntMap (Q.PrepArg, PGScalarValue)

data PlanningSt = PlanningSt
  { PlanningSt -> Int
_psArgNumber :: Int,
    PlanningSt -> PlanVariables
_psVariables :: PlanVariables,
    PlanningSt -> PrepArgMap
_psPrepped :: PrepArgMap
  }
  deriving stock (PlanningSt -> PlanningSt -> Bool
(PlanningSt -> PlanningSt -> Bool)
-> (PlanningSt -> PlanningSt -> Bool) -> Eq PlanningSt
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PlanningSt -> PlanningSt -> Bool
$c/= :: PlanningSt -> PlanningSt -> Bool
== :: PlanningSt -> PlanningSt -> Bool
$c== :: PlanningSt -> PlanningSt -> Bool
Eq, Int -> PlanningSt -> ShowS
[PlanningSt] -> ShowS
PlanningSt -> String
(Int -> PlanningSt -> ShowS)
-> (PlanningSt -> String)
-> ([PlanningSt] -> ShowS)
-> Show PlanningSt
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PlanningSt] -> ShowS
$cshowList :: [PlanningSt] -> ShowS
show :: PlanningSt -> String
$cshow :: PlanningSt -> String
showsPrec :: Int -> PlanningSt -> ShowS
$cshowsPrec :: Int -> PlanningSt -> ShowS
Show)

initPlanningSt :: PlanningSt
initPlanningSt :: PlanningSt
initPlanningSt = Int -> PlanVariables -> PrepArgMap -> PlanningSt
PlanningSt Int
2 PlanVariables
forall k v. HashMap k v
Map.empty PrepArgMap
forall a. IntMap a
IntMap.empty

-- | If we're preparing a value with planning state, we favour referring to
-- values by their prepared argument index. If the value refers to a session
-- value, we look for it in prepared value (1) and access the particular keys
-- using the JSONB @->>@ accessor.
prepareWithPlan ::
  ( MonadState PlanningSt m,
    MonadError QErr m
  ) =>
  UserInfo ->
  UnpreparedValue ('Postgres pgKind) ->
  m S.SQLExp
prepareWithPlan :: UserInfo -> UnpreparedValue ('Postgres pgKind) -> m SQLExp
prepareWithPlan UserInfo
userInfo = \case
  UVParameter Maybe VariableInfo
varInfoM ColumnValue {ScalarValue ('Postgres pgKind)
ColumnType ('Postgres pgKind)
cvValue :: forall (b :: BackendType). ColumnValue b -> ScalarValue b
cvType :: forall (b :: BackendType). ColumnValue b -> ColumnType b
cvValue :: ScalarValue ('Postgres pgKind)
cvType :: ColumnType ('Postgres pgKind)
..} -> do
    Int
argNum <- m Int -> (VariableInfo -> m Int) -> Maybe VariableInfo -> m Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m Int
forall (m :: * -> *). MonadState PlanningSt m => m Int
getNextArgNum (Name -> m Int
forall (m :: * -> *). MonadState PlanningSt m => Name -> m Int
getVarArgNum (Name -> m Int) -> (VariableInfo -> Name) -> VariableInfo -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VariableInfo -> Name
forall a. HasName a => a -> Name
getName) Maybe VariableInfo
varInfoM
    Int -> (PrepArg, PGScalarValue) -> m ()
forall (m :: * -> *).
MonadState PlanningSt m =>
Int -> (PrepArg, PGScalarValue) -> m ()
addPrepArg Int
argNum (PGScalarValue -> PrepArg
binEncoder ScalarValue ('Postgres pgKind)
PGScalarValue
cvValue, ScalarValue ('Postgres pgKind)
PGScalarValue
cvValue)
    SQLExp -> m SQLExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SQLExp -> m SQLExp) -> SQLExp -> m SQLExp
forall a b. (a -> b) -> a -> b
$ Int -> PGScalarType -> SQLExp
toPrepParam Int
argNum (ColumnType ('Postgres pgKind) -> PGScalarType
forall (pgKind :: PostgresKind).
ColumnType ('Postgres pgKind) -> PGScalarType
unsafePGColumnToBackend ColumnType ('Postgres pgKind)
cvType)
  UVSessionVar SessionVarType ('Postgres pgKind)
ty SessionVariable
sessVar -> do
    -- For queries, we need to make sure the session variables are passed. However,
    -- we want to keep them as variables in the resulting SQL in order to keep
    -- hitting query caching for similar queries.
    SessionVariableValue
_ <-
      SessionVariable -> SessionVariables -> Maybe SessionVariableValue
getSessionVariableValue SessionVariable
sessVar (UserInfo -> SessionVariables
_uiSession UserInfo
userInfo)
        Maybe SessionVariableValue
-> m SessionVariableValue -> m SessionVariableValue
forall (m :: * -> *) a. Applicative m => Maybe a -> m a -> m a
`onNothing` Code -> SessionVariableValue -> m SessionVariableValue
forall (m :: * -> *) a.
QErrM m =>
Code -> SessionVariableValue -> m a
throw400
          Code
NotFound
          (SessionVariableValue
"missing session variable: " SessionVariableValue
-> SessionVariableValue -> SessionVariableValue
forall t.
ToTxt t =>
SessionVariableValue -> t -> SessionVariableValue
<>> SessionVariable -> SessionVariableValue
sessionVariableToText SessionVariable
sessVar)
    let sessVarVal :: SQLExp
sessVarVal =
          SQLOp -> [SQLExp] -> SQLExp
S.SEOpApp
            (SessionVariableValue -> SQLOp
S.SQLOp SessionVariableValue
"->>")
            [SQLExp
currentSessionExp, SessionVariableValue -> SQLExp
S.SELit (SessionVariableValue -> SQLExp) -> SessionVariableValue -> SQLExp
forall a b. (a -> b) -> a -> b
$ SessionVariable -> SessionVariableValue
sessionVariableToText SessionVariable
sessVar]
    SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SQLExp -> m SQLExp) -> SQLExp -> m SQLExp
forall a b. (a -> b) -> a -> b
$ CollectableType PGScalarType -> SQLExp -> SQLExp
withTypeAnn SessionVarType ('Postgres pgKind)
CollectableType PGScalarType
ty SQLExp
sessVarVal
  UVLiteral SQLExpression ('Postgres pgKind)
sqlExp -> SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure SQLExpression ('Postgres pgKind)
SQLExp
sqlExp
  UnpreparedValue ('Postgres pgKind)
UVSession -> SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure SQLExp
currentSessionExp
  where
    currentSessionExp :: SQLExp
currentSessionExp = Int -> SQLExp
S.SEPrep Int
1

-- | If we're /not/ using a prepared statement, substitution is pretty naïve:
-- we resolve session variable names, ignore parameter names, and substitute
-- into the 'S.SQLExp'.
prepareWithoutPlan ::
  (MonadError QErr m) =>
  UserInfo ->
  UnpreparedValue ('Postgres pgKind) ->
  m S.SQLExp
prepareWithoutPlan :: UserInfo -> UnpreparedValue ('Postgres pgKind) -> m SQLExp
prepareWithoutPlan UserInfo
userInfo = \case
  UVParameter Maybe VariableInfo
_ ColumnValue ('Postgres pgKind)
cv -> SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SQLExp -> m SQLExp) -> SQLExp -> m SQLExp
forall a b. (a -> b) -> a -> b
$ ColumnValue ('Postgres pgKind) -> SQLExp
forall (pgKind :: PostgresKind).
ColumnValue ('Postgres pgKind) -> SQLExp
toTxtValue ColumnValue ('Postgres pgKind)
cv
  UVLiteral SQLExpression ('Postgres pgKind)
sqlExp -> SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure SQLExpression ('Postgres pgKind)
SQLExp
sqlExp
  UnpreparedValue ('Postgres pgKind)
UVSession -> SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SQLExp -> m SQLExp) -> SQLExp -> m SQLExp
forall a b. (a -> b) -> a -> b
$ SessionVariables -> SQLExp
sessionInfoJsonExp (SessionVariables -> SQLExp) -> SessionVariables -> SQLExp
forall a b. (a -> b) -> a -> b
$ UserInfo -> SessionVariables
_uiSession UserInfo
userInfo
  UVSessionVar SessionVarType ('Postgres pgKind)
ty SessionVariable
sessVar -> do
    let maybeSessionVariableValue :: Maybe SessionVariableValue
maybeSessionVariableValue =
          SessionVariable -> SessionVariables -> Maybe SessionVariableValue
getSessionVariableValue SessionVariable
sessVar (UserInfo -> SessionVariables
_uiSession UserInfo
userInfo)
    SQLExp
sessionVariableValue <-
      (SessionVariableValue -> SQLExp)
-> m SessionVariableValue -> m SQLExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SessionVariableValue -> SQLExp
S.SELit
        (m SessionVariableValue -> m SQLExp)
-> (m SessionVariableValue -> m SessionVariableValue)
-> m SessionVariableValue
-> m SQLExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SessionVariableValue
-> m SessionVariableValue -> m SessionVariableValue
forall (m :: * -> *) a. Applicative m => Maybe a -> m a -> m a
onNothing Maybe SessionVariableValue
maybeSessionVariableValue
        (m SessionVariableValue -> m SQLExp)
-> m SessionVariableValue -> m SQLExp
forall a b. (a -> b) -> a -> b
$ Code -> SessionVariableValue -> m SessionVariableValue
forall (m :: * -> *) a.
QErrM m =>
Code -> SessionVariableValue -> m a
throw400 Code
NotFound (SessionVariableValue -> m SessionVariableValue)
-> SessionVariableValue -> m SessionVariableValue
forall a b. (a -> b) -> a -> b
$
          SessionVariableValue
"missing session variable: " SessionVariableValue
-> SessionVariableValue -> SessionVariableValue
forall t.
ToTxt t =>
SessionVariableValue -> t -> SessionVariableValue
<>> SessionVariable -> SessionVariableValue
sessionVariableToText SessionVariable
sessVar
    SQLExp -> m SQLExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SQLExp -> m SQLExp) -> SQLExp -> m SQLExp
forall a b. (a -> b) -> a -> b
$ CollectableType PGScalarType -> SQLExp -> SQLExp
withTypeAnn SessionVarType ('Postgres pgKind)
CollectableType PGScalarType
ty SQLExp
sessionVariableValue

-- | The map of user session variables is always given the number (1) as its
-- variable argument number (see 'getVarArgNum'). If we want to refer to a
-- particular variable in this map, we use JSONB functions to interrogate
-- variable (1).
withUserVars :: SessionVariables -> PrepArgMap -> PrepArgMap
withUserVars :: SessionVariables -> PrepArgMap -> PrepArgMap
withUserVars SessionVariables
usrVars PrepArgMap
list =
  let usrVarsAsPgScalar :: PGScalarValue
usrVarsAsPgScalar = JSON -> PGScalarValue
PGValJSON (JSON -> PGScalarValue) -> JSON -> PGScalarValue
forall a b. (a -> b) -> a -> b
$ Value -> JSON
Q.JSON (Value -> JSON) -> Value -> JSON
forall a b. (a -> b) -> a -> b
$ SessionVariables -> Value
forall a. ToJSON a => a -> Value
J.toJSON SessionVariables
usrVars
      prepArg :: PrepArg
prepArg = AltJ SessionVariables -> PrepArg
forall a. ToPrepArg a => a -> PrepArg
Q.toPrepVal (SessionVariables -> AltJ SessionVariables
forall a. a -> AltJ a
Q.AltJ SessionVariables
usrVars)
   in Int -> (PrepArg, PGScalarValue) -> PrepArgMap -> PrepArgMap
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
1 (PrepArg
prepArg, PGScalarValue
usrVarsAsPgScalar) PrepArgMap
list

-- | In prepared statements, we refer to variables by a number, not their name.
-- If the statement already refers to a variable, then we'll already have a
-- number for it, and so we just return that. Otherwise, we produce a new
-- number, and that will refer to the variable from now on.
getVarArgNum :: (MonadState PlanningSt m) => G.Name -> m Int
getVarArgNum :: Name -> m Int
getVarArgNum Name
var = do
  PlanningSt Int
curArgNum PlanVariables
vars PrepArgMap
prepped <- m PlanningSt
forall s (m :: * -> *). MonadState s m => m s
get
  Name -> PlanVariables -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup Name
var PlanVariables
vars Maybe Int -> m Int -> m Int
forall (m :: * -> *) a. Applicative m => Maybe a -> m a -> m a
`onNothing` do
    PlanningSt -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (PlanningSt -> m ()) -> PlanningSt -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> PlanVariables -> PrepArgMap -> PlanningSt
PlanningSt (Int
curArgNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Name -> Int -> PlanVariables -> PlanVariables
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert Name
var Int
curArgNum PlanVariables
vars) PrepArgMap
prepped
    Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
curArgNum

-- | Add a prepared argument to the prepared argument map. These are keyed by
-- the variable argument numbers, which can be computed using 'getVarArgNum'.
addPrepArg ::
  (MonadState PlanningSt m) =>
  Int ->
  (Q.PrepArg, PGScalarValue) ->
  m ()
addPrepArg :: Int -> (PrepArg, PGScalarValue) -> m ()
addPrepArg Int
argNum (PrepArg, PGScalarValue)
arg = (PlanningSt -> PlanningSt) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify \PlanningSt
s ->
  PlanningSt
s {_psPrepped :: PrepArgMap
_psPrepped = Int -> (PrepArg, PGScalarValue) -> PrepArgMap -> PrepArgMap
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
argNum (PrepArg, PGScalarValue)
arg (PlanningSt -> PrepArgMap
_psPrepped PlanningSt
s)}

-- | Get '_psArgNumber' from inside the 'PlanningSt' and increment it for the
-- next operation. Think of this as a pure analogue to 'Data.Unique.newUnique'.
getNextArgNum :: (MonadState PlanningSt m) => m Int
getNextArgNum :: m Int
getNextArgNum = (PlanningSt -> (Int, PlanningSt)) -> m Int
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state \PlanningSt
s ->
  ( PlanningSt -> Int
_psArgNumber PlanningSt
s,
    PlanningSt
s {_psArgNumber :: Int
_psArgNumber = PlanningSt -> Int
_psArgNumber PlanningSt
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1}
  )