{-# LANGUAGE DeriveAnyClass #-}

-- | 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 HashMap
import Data.IntMap qualified as IntMap
import Data.Text.Extended
import Database.PG.Query qualified as PG
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.BackendType
import Hasura.RQL.Types.Column
import Hasura.Session
  ( SessionVariables,
    UserInfo (_uiSession),
    getSessionVariableValue,
    sessionVariableToText,
  )
import Language.GraphQL.Draft.Syntax qualified as G

type PlanVariables = HashMap.HashMap PlanVariable Int

-- | A variable used within the 'PlanVariables' map. We make the distinction
-- between internal (i.e. generated within `graphql-engine`) and external (i.e.
-- generated by users or as part of a query) variables, which avoids accidental
-- name collisions between the two.
data PlanVariable
  = -- | An external GQL API name
    External G.Name
  | -- | An internal HGE name
    Internal Text
  deriving stock (PlanVariable -> PlanVariable -> Bool
(PlanVariable -> PlanVariable -> Bool)
-> (PlanVariable -> PlanVariable -> Bool) -> Eq PlanVariable
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PlanVariable -> PlanVariable -> Bool
== :: PlanVariable -> PlanVariable -> Bool
$c/= :: PlanVariable -> PlanVariable -> Bool
/= :: PlanVariable -> PlanVariable -> Bool
Eq, (forall x. PlanVariable -> Rep PlanVariable x)
-> (forall x. Rep PlanVariable x -> PlanVariable)
-> Generic PlanVariable
forall x. Rep PlanVariable x -> PlanVariable
forall x. PlanVariable -> Rep PlanVariable x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. PlanVariable -> Rep PlanVariable x
from :: forall x. PlanVariable -> Rep PlanVariable x
$cto :: forall x. Rep PlanVariable x -> PlanVariable
to :: forall x. Rep PlanVariable x -> PlanVariable
Generic, Int -> PlanVariable -> ShowS
[PlanVariable] -> ShowS
PlanVariable -> String
(Int -> PlanVariable -> ShowS)
-> (PlanVariable -> String)
-> ([PlanVariable] -> ShowS)
-> Show PlanVariable
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PlanVariable -> ShowS
showsPrec :: Int -> PlanVariable -> ShowS
$cshow :: PlanVariable -> String
show :: PlanVariable -> String
$cshowList :: [PlanVariable] -> ShowS
showList :: [PlanVariable] -> ShowS
Show)
  deriving anyclass (Eq PlanVariable
Eq PlanVariable
-> (Int -> PlanVariable -> Int)
-> (PlanVariable -> Int)
-> Hashable PlanVariable
Int -> PlanVariable -> Int
PlanVariable -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
$chashWithSalt :: Int -> PlanVariable -> Int
hashWithSalt :: Int -> PlanVariable -> Int
$chash :: PlanVariable -> Int
hash :: PlanVariable -> Int
Hashable)

-- | The value is (PG.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 (PG.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
$c== :: PlanningSt -> PlanningSt -> Bool
== :: PlanningSt -> PlanningSt -> Bool
$c/= :: PlanningSt -> PlanningSt -> Bool
/= :: 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
$cshowsPrec :: Int -> PlanningSt -> ShowS
showsPrec :: Int -> PlanningSt -> ShowS
$cshow :: PlanningSt -> String
show :: PlanningSt -> String
$cshowList :: [PlanningSt] -> ShowS
showList :: [PlanningSt] -> ShowS
Show)

initPlanningSt :: PlanningSt
initPlanningSt :: PlanningSt
initPlanningSt = Int -> PlanVariables -> PrepArgMap -> PlanningSt
PlanningSt Int
2 PlanVariables
forall k v. HashMap k v
HashMap.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 :: forall (m :: * -> *) (pgKind :: PostgresKind).
(MonadState PlanningSt m, MonadError QErr m) =>
UserInfo -> UnpreparedValue ('Postgres pgKind) -> m SQLExp
prepareWithPlan UserInfo
userInfo = \case
  UVParameter Provenance
provenance ColumnValue {ScalarValue ('Postgres pgKind)
ColumnType ('Postgres pgKind)
cvType :: ColumnType ('Postgres pgKind)
cvValue :: ScalarValue ('Postgres pgKind)
cvType :: forall (b :: BackendType). ColumnValue b -> ColumnType b
cvValue :: forall (b :: BackendType). ColumnValue b -> ScalarValue b
..} -> do
    Int
argNum <- case Provenance
provenance of
      FromInternal Text
name -> PlanVariable -> m Int
forall (m :: * -> *).
MonadState PlanningSt m =>
PlanVariable -> m Int
getVarArgNum (Text -> PlanVariable
Internal Text
name)
      FromGraphQL VariableInfo
varInfo -> PlanVariable -> m Int
forall (m :: * -> *).
MonadState PlanningSt m =>
PlanVariable -> m Int
getVarArgNum (Name -> PlanVariable
External (VariableInfo -> Name
forall a. HasName a => a -> Name
getName VariableInfo
varInfo))
      Provenance
FreshVar -> m Int
forall (m :: * -> *). MonadState PlanningSt m => m Int
getNextArgNum
    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 a. a -> m a
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.
    Text
_ <-
      SessionVariable -> SessionVariables -> Maybe Text
getSessionVariableValue SessionVariable
sessVar (UserInfo -> SessionVariables
_uiSession UserInfo
userInfo)
        Maybe Text -> m Text -> m Text
forall (m :: * -> *) a. Applicative m => Maybe a -> m a -> m a
`onNothing` Code -> Text -> m Text
forall (m :: * -> *) a. QErrM m => Code -> Text -> m a
throw400
          Code
NotFound
          (Text
"missing session variable: " Text -> Text -> Text
forall t. ToTxt t => Text -> t -> Text
<>> SessionVariable -> Text
sessionVariableToText SessionVariable
sessVar)
    let sessVarVal :: SQLExp
sessVarVal =
          SQLOp -> [SQLExp] -> SQLExp
S.SEOpApp
            (Text -> SQLOp
S.SQLOp Text
"->>")
            [SQLExp
currentSessionExp, Text -> SQLExp
S.SELit (Text -> SQLExp) -> Text -> SQLExp
forall a b. (a -> b) -> a -> b
$ SessionVariable -> Text
sessionVariableToText SessionVariable
sessVar]
    SQLExp -> m SQLExp
forall a. a -> m a
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SQLExpression ('Postgres pgKind)
SQLExp
sqlExp
  UnpreparedValue ('Postgres pgKind)
UVSession -> SQLExp -> m SQLExp
forall a. a -> m a
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 :: forall (m :: * -> *) (pgKind :: PostgresKind).
MonadError QErr m =>
UserInfo -> UnpreparedValue ('Postgres pgKind) -> m SQLExp
prepareWithoutPlan UserInfo
userInfo = \case
  UVParameter Provenance
_ ColumnValue ('Postgres pgKind)
cv -> SQLExp -> m SQLExp
forall a. a -> m a
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SQLExpression ('Postgres pgKind)
SQLExp
sqlExp
  UnpreparedValue ('Postgres pgKind)
UVSession -> SQLExp -> m SQLExp
forall a. a -> m a
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 Text
maybeSessionVariableValue =
          SessionVariable -> SessionVariables -> Maybe Text
getSessionVariableValue SessionVariable
sessVar (UserInfo -> SessionVariables
_uiSession UserInfo
userInfo)
    SQLExp
sessionVariableValue <-
      (Text -> SQLExp) -> m Text -> m SQLExp
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Text -> SQLExp
S.SELit
        (m Text -> m SQLExp) -> (m Text -> m Text) -> m Text -> m SQLExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text -> m Text -> m Text
forall (m :: * -> *) a. Applicative m => Maybe a -> m a -> m a
onNothing Maybe Text
maybeSessionVariableValue
        (m Text -> m SQLExp) -> m Text -> m SQLExp
forall a b. (a -> b) -> a -> b
$ Code -> Text -> m Text
forall (m :: * -> *) a. QErrM m => Code -> Text -> m a
throw400 Code
NotFound
        (Text -> m Text) -> Text -> m Text
forall a b. (a -> b) -> a -> b
$ Text
"missing session variable: "
        Text -> Text -> Text
forall t. ToTxt t => Text -> t -> Text
<>> SessionVariable -> Text
sessionVariableToText SessionVariable
sessVar
    SQLExp -> m SQLExp
forall a. a -> m a
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
PG.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 = ViaJSON SessionVariables -> PrepArg
forall a. ToPrepArg a => a -> PrepArg
PG.toPrepVal (SessionVariables -> ViaJSON SessionVariables
forall a. a -> ViaJSON a
PG.ViaJSON 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) => PlanVariable -> m Int
getVarArgNum :: forall (m :: * -> *).
MonadState PlanningSt m =>
PlanVariable -> m Int
getVarArgNum PlanVariable
var = do
  PlanningSt Int
curArgNum PlanVariables
vars PrepArgMap
prepped <- m PlanningSt
forall s (m :: * -> *). MonadState s m => m s
get
  PlanVariable -> PlanVariables -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup PlanVariable
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) (PlanVariable -> Int -> PlanVariables -> PlanVariables
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert PlanVariable
var Int
curArgNum PlanVariables
vars) PrepArgMap
prepped
    Int -> m Int
forall a. a -> m a
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 ->
  (PG.PrepArg, PGScalarValue) ->
  m ()
addPrepArg :: forall (m :: * -> *).
MonadState PlanningSt m =>
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 :: forall (m :: * -> *). MonadState PlanningSt m => m Int
getNextArgNum = (PlanningSt -> (Int, PlanningSt)) -> m Int
forall a. (PlanningSt -> (a, PlanningSt)) -> m a
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}
  )