feat: Allow getting the EXPLAIN plan of a request

2022-07-27 19:33:34 -05:00
parent d2df289696
commit 8911afd079
26 changed files with 582 additions and 218 deletions
@@ -74,6 +74,7 @@ import PostgREST.GucHeader               (GucHeader,
                                          addHeadersIfNotIncluded,
                                          unwrapGucHeader)
 import PostgREST.MediaType               (MediaType (..))
+import PostgREST.Query.Statements        (ResultSet (..))
 import PostgREST.Request.ApiRequest      (Action (..),
                                          ApiRequest (..),
                                          InvokeMethod (..),
@@ -259,9 +260,9 @@ handleRead headersOnly identifier context@RequestContext{..} = do
    AppConfig{..} = ctxConfig
    countQuery = QueryBuilder.readRequestToCountQuery req

-  (tableTotal, queryTotal, _ , body, gucHeaders, gucStatus) <-
-    lift . SQL.statement mempty $
-      Statements.createReadStatement
+  resultSet <-
+     lift . SQL.statement mempty $
+      Statements.prepareRead
        (QueryBuilder.readRequestToQuery req)
        (if iPreferCount == Just EstimatedCount then
           -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
@@ -274,23 +275,28 @@ handleRead headersOnly identifier context@RequestContext{..} = do
        bField
        configDbPreparedStatements

-  total <- readTotal ctxConfig ctxApiRequest tableTotal countQuery
-  response <- liftEither $ gucResponse <$> gucStatus <*> gucHeaders
+  case resultSet of
+    RSStandard{..} -> do
+      total <- readTotal ctxConfig ctxApiRequest rsTableTotal countQuery
+      response <- liftEither $ gucResponse <$> rsGucStatus <*> rsGucHeaders

-  let
-    (status, contentRange) = RangeQuery.rangeStatusHeader iTopLevelRange queryTotal total
-    headers =
-      [ contentRange
-      , ( "Content-Location"
-        , "/"
-            <> toUtf8 (qiName identifier)
-            <> if BS.null (qsCanonical iQueryParams) then mempty else "?" <> qsCanonical iQueryParams
-        )
-      ]
-      ++ contentTypeHeaders context
+      let
+        (status, contentRange) = RangeQuery.rangeStatusHeader iTopLevelRange rsQueryTotal total
+        headers =
+          [ contentRange
+          , ( "Content-Location"
+            , "/"
+                <> toUtf8 (qiName identifier)
+                <> if BS.null (qsCanonical iQueryParams) then mempty else "?" <> qsCanonical iQueryParams
+            )
+          ]
+          ++ contentTypeHeaders context

-  failNotSingular iAcceptMediaType queryTotal . response status headers $
-    if headersOnly then mempty else LBS.fromStrict body
+      failNotSingular iAcceptMediaType rsQueryTotal . response status headers $
+        if headersOnly then mempty else LBS.fromStrict rsBody
+
+    RSPlan plan ->
+      pure $ Wai.responseLBS HTTP.status200 (contentTypeHeaders context) $ LBS.fromStrict plan

 readTotal :: AppConfig -> ApiRequest -> Maybe Int64 -> SQL.Snippet -> DbHandler (Maybe Int64)
 readTotal AppConfig{..} ApiRequest{..} tableTotal countQuery =
@@ -306,7 +312,7 @@ readTotal AppConfig{..} ApiRequest{..} tableTotal countQuery =
      return tableTotal
  where
    explain =
-      lift . SQL.statement mempty . Statements.createExplainStatement countQuery $
+      lift . SQL.statement mempty . Statements.preparePlanRows countQuery $
        configDbPreparedStatements

 handleCreate :: QualifiedIdentifier -> RequestContext -> DbHandler Wai.Response
@@ -317,34 +323,41 @@ handleCreate identifier@QualifiedIdentifier{..} context@RequestContext{..} = do
      then maybe mempty tablePKCols $ HM.lookup identifier $ dbTables ctxDbStructure
      else mempty

-  WriteQueryResult{..} <- writeQuery MutationCreate identifier True pkCols context
+  resultSet <- writeQuery MutationCreate identifier True pkCols context

-  let
-    response = gucResponse resGucStatus resGucHeaders
-    headers =
-      catMaybes
-        [ if null resFields then
-            Nothing
-          else
-            Just
-              ( HTTP.hLocation
-              , "/"
-                  <> toUtf8 qiName
-                  <> HTTP.renderSimpleQuery True (splitKeyValue <$> resFields)
-              )
-        , Just . RangeQuery.contentRangeH 1 0 $
-            if shouldCount iPreferCount then Just resQueryTotal else Nothing
-        , if null pkCols && isNothing (qsOnConflict iQueryParams) then
-            Nothing
-          else
-            toAppliedHeader <$> iPreferResolution
-        ]
+  case resultSet of
+    RSStandard{..} -> do

-  failNotSingular iAcceptMediaType resQueryTotal $
-    if iPreferRepresentation == Full then
-      response HTTP.status201 (headers ++ contentTypeHeaders context) (LBS.fromStrict resBody)
-    else
-      response HTTP.status201 headers mempty
+      response <- liftEither $ gucResponse <$> rsGucStatus <*> rsGucHeaders
+
+      let
+        headers =
+          catMaybes
+            [ if null rsLocation then
+                Nothing
+              else
+                Just
+                  ( HTTP.hLocation
+                  , "/"
+                      <> toUtf8 qiName
+                      <> HTTP.renderSimpleQuery True rsLocation
+                  )
+            , Just . RangeQuery.contentRangeH 1 0 $
+                if shouldCount iPreferCount then Just rsQueryTotal else Nothing
+            , if null pkCols && isNothing (qsOnConflict iQueryParams) then
+                Nothing
+              else
+                toAppliedHeader <$> iPreferResolution
+            ]
+
+      failNotSingular iAcceptMediaType rsQueryTotal $
+        if iPreferRepresentation == Full then
+          response HTTP.status201 (headers ++ contentTypeHeaders context) (LBS.fromStrict rsBody)
+        else
+          response HTTP.status201 headers mempty
+
+    RSPlan plan ->
+      pure $ Wai.responseLBS HTTP.status200 (contentTypeHeaders context) $ LBS.fromStrict plan

 handleUpdate :: QualifiedIdentifier -> RequestContext -> DbHandler Wai.Response
 handleUpdate identifier context@RequestContext{..} = do
@@ -352,68 +365,87 @@ handleUpdate identifier context@RequestContext{..} = do
    ApiRequest{..} = ctxApiRequest
    pkCols = maybe mempty tablePKCols $ HM.lookup identifier $ dbTables ctxDbStructure

-  WriteQueryResult{..} <- writeQuery MutationUpdate identifier False pkCols context
+  resultSet <- writeQuery MutationUpdate identifier False pkCols context

-  let
-    response = gucResponse resGucStatus resGucHeaders
-    fullRepr = iPreferRepresentation == Full
-    updateIsNoOp = S.null iColumns
-    status
-      | resQueryTotal == 0 && not updateIsNoOp = HTTP.status404
-      | fullRepr = HTTP.status200
-      | otherwise = HTTP.status204
-    contentRangeHeader =
-      RangeQuery.contentRangeH 0 (resQueryTotal - 1) $
-        if shouldCount iPreferCount then Just resQueryTotal else Nothing
+  case resultSet of
+    RSStandard{..} -> do
+      response <- liftEither $ gucResponse <$> rsGucStatus <*> rsGucHeaders

-  failChangesOffLimits (RangeQuery.rangeLimit iTopLevelRange) resQueryTotal =<<
-    failNotSingular iAcceptMediaType resQueryTotal (
-      if fullRepr then
-        response status (contentTypeHeaders context ++ [contentRangeHeader]) (LBS.fromStrict resBody)
-      else
-        response status [contentRangeHeader] mempty)
+      let
+        fullRepr = iPreferRepresentation == Full
+        updateIsNoOp = S.null iColumns
+        status
+          | rsQueryTotal == 0 && not updateIsNoOp = HTTP.status404
+          | fullRepr = HTTP.status200
+          | otherwise = HTTP.status204
+        contentRangeHeader =
+          RangeQuery.contentRangeH 0 (rsQueryTotal - 1) $
+            if shouldCount iPreferCount then Just rsQueryTotal else Nothing
+
+      failChangesOffLimits (RangeQuery.rangeLimit iTopLevelRange) rsQueryTotal =<<
+        failNotSingular iAcceptMediaType rsQueryTotal (
+          if fullRepr then
+            response status (contentTypeHeaders context ++ [contentRangeHeader]) (LBS.fromStrict rsBody)
+          else
+            response status [contentRangeHeader] mempty)
+
+    RSPlan plan ->
+      pure $ Wai.responseLBS HTTP.status200 (contentTypeHeaders context) $ LBS.fromStrict plan

 handleSingleUpsert :: QualifiedIdentifier -> RequestContext-> DbHandler Wai.Response
 handleSingleUpsert identifier context@(RequestContext _ ctxDbStructure ApiRequest{..} _) = do
  let pkCols = maybe mempty tablePKCols $ HM.lookup identifier $ dbTables ctxDbStructure

-  WriteQueryResult{..} <- writeQuery MutationSingleUpsert identifier False pkCols context
+  resultSet <- writeQuery MutationSingleUpsert identifier False pkCols context

-  let response = gucResponse resGucStatus resGucHeaders
+  case resultSet of
+    RSStandard {..} -> do

-  -- Makes sure the querystring pk matches the payload pk
-  -- e.g. PUT /items?id=eq.1 { "id" : 1, .. } is accepted,
-  -- PUT /items?id=eq.14 { "id" : 2, .. } is rejected.
-  -- If this condition is not satisfied then nothing is inserted,
-  -- check the WHERE for INSERT in QueryBuilder.hs to see how it's done
-  when (resQueryTotal /= 1) $ do
-    lift SQL.condemn
-    throwError Error.PutMatchingPkError
+      response <- liftEither $ gucResponse <$> rsGucStatus <*> rsGucHeaders

-  return $
-    if iPreferRepresentation == Full then
-      response HTTP.status200 (contentTypeHeaders context) (LBS.fromStrict resBody)
-    else
-      response HTTP.status204 [] mempty
+      -- Makes sure the querystring pk matches the payload pk
+      -- e.g. PUT /items?id=eq.1 { "id" : 1, .. } is accepted,
+      -- PUT /items?id=eq.14 { "id" : 2, .. } is rejected.
+      -- If this condition is not satisfied then nothing is inserted,
+      -- check the WHERE for INSERT in QueryBuilder.hs to see how it's done
+      when (rsQueryTotal /= 1) $ do
+        lift SQL.condemn
+        throwError Error.PutMatchingPkError
+
+      return $
+        if iPreferRepresentation == Full then
+          response HTTP.status200 (contentTypeHeaders context) (LBS.fromStrict rsBody)
+        else
+          response HTTP.status204 [] mempty
+
+    RSPlan plan ->
+      pure $ Wai.responseLBS HTTP.status200 (contentTypeHeaders context) $ LBS.fromStrict plan

 handleDelete :: QualifiedIdentifier -> RequestContext -> DbHandler Wai.Response
 handleDelete identifier context@(RequestContext _ _ ApiRequest{..} _) = do
-  WriteQueryResult{..} <- writeQuery MutationDelete identifier False mempty context
+  resultSet <- writeQuery MutationDelete identifier False mempty context

-  let
-    response = gucResponse resGucStatus resGucHeaders
-    contentRangeHeader =
-      RangeQuery.contentRangeH 1 0 $
-        if shouldCount iPreferCount then Just resQueryTotal else Nothing
+  case resultSet of
+    RSStandard {..} -> do

-  failChangesOffLimits (RangeQuery.rangeLimit iTopLevelRange) resQueryTotal =<<
-    failNotSingular iAcceptMediaType resQueryTotal (
-      if iPreferRepresentation == Full then
-        response HTTP.status200
-          (contentTypeHeaders context ++ [contentRangeHeader])
-          (LBS.fromStrict resBody)
-      else
-        response HTTP.status204 [contentRangeHeader] mempty)
+      response <- liftEither $ gucResponse <$> rsGucStatus <*> rsGucHeaders
+
+      let
+        contentRangeHeader =
+          RangeQuery.contentRangeH 1 0 $
+            if shouldCount iPreferCount then Just rsQueryTotal else Nothing
+
+      failChangesOffLimits (RangeQuery.rangeLimit iTopLevelRange) rsQueryTotal =<<
+        failNotSingular iAcceptMediaType rsQueryTotal (
+          if iPreferRepresentation == Full then
+            response HTTP.status200
+              (contentTypeHeaders context ++ [contentRangeHeader])
+              (LBS.fromStrict rsBody)
+          else
+            response HTTP.status204 [contentRangeHeader] mempty)
+
+    RSPlan plan ->
+      pure $ Wai.responseLBS HTTP.status200 (contentTypeHeaders context) $ LBS.fromStrict plan

 handleInfo :: Monad m => Target -> RequestContext -> Handler m Wai.Response
 handleInfo target RequestContext{..} =
@@ -453,9 +485,9 @@ handleInvoke invMethod proc context@RequestContext{..} = do

  let callReq = ReqBuilder.callRequest proc ctxApiRequest req

-  (tableTotal, queryTotal, body, gucHeaders, gucStatus) <-
+  resultSet <-
    lift . SQL.statement mempty $
-      Statements.callProcStatement
+      Statements.prepareCall
        (Proc.procReturnsScalar proc)
        (Proc.procReturnsSingle proc)
        (QueryBuilder.requestToCallProcQuery callReq)
@@ -467,19 +499,23 @@ handleInvoke invMethod proc context@RequestContext{..} = do
        bField
        (configDbPreparedStatements ctxConfig)

-  response <- liftEither $ gucResponse <$> gucStatus <*> gucHeaders
+  case resultSet of
+    RSStandard {..} -> do
+      response <- liftEither $ gucResponse <$> rsGucStatus <*> rsGucHeaders
+      let
+        (status, contentRange) =
+          RangeQuery.rangeStatusHeader iTopLevelRange rsQueryTotal rsTableTotal

-  let
-    (status, contentRange) =
-      RangeQuery.rangeStatusHeader iTopLevelRange queryTotal tableTotal
+      failNotSingular iAcceptMediaType rsQueryTotal $
+        if Proc.procReturnsVoid proc then
+          response HTTP.status204 [contentRange] mempty
+        else
+          response status
+            (contentTypeHeaders context ++ [contentRange])
+            (if invMethod == InvHead then mempty else LBS.fromStrict rsBody)

-  failNotSingular iAcceptMediaType queryTotal $
-    if Proc.procReturnsVoid proc then
-      response HTTP.status204 [contentRange] mempty
-    else
-      response status
-        (contentTypeHeaders context ++ [contentRange])
-        (if invMethod == InvHead then mempty else LBS.fromStrict body)
+    RSPlan plan ->
+      pure $ Wai.responseLBS HTTP.status200 (contentTypeHeaders context) $ LBS.fromStrict plan

 handleOpenApi :: Bool -> Schema -> RequestContext -> DbHandler Wai.Response
 handleOpenApi headersOnly tSchema (RequestContext conf@AppConfig{..} dbStructure apiRequest ctxPgVersion) = do
@@ -523,16 +559,7 @@ txMode ApiRequest{..} =
    _ ->
      SQL.Write

-- | Result from executing a write query on the database
-data WriteQueryResult = WriteQueryResult
-  { resQueryTotal :: Int64
-  , resFields     :: [ByteString]
-  , resBody       :: ByteString
-  , resGucStatus  :: Maybe HTTP.Status
-  , resGucHeaders :: [GucHeader]
-  }
-
-writeQuery :: Mutation -> QualifiedIdentifier -> Bool -> [Text] -> RequestContext -> DbHandler WriteQueryResult
+writeQuery :: Mutation -> QualifiedIdentifier -> Bool -> [Text] -> RequestContext -> DbHandler ResultSet
 writeQuery mutation identifier@QualifiedIdentifier{..} isInsert pkCols context@RequestContext{..} = do
  readReq <- readRequest identifier context

@@ -542,18 +569,15 @@ writeQuery mutation identifier@QualifiedIdentifier{..} isInsert pkCols context@R
        pkCols
        readReq

-  (_, queryTotal, fields, body, gucHeaders, gucStatus) <-
-    lift . SQL.statement mempty $
-      Statements.createWriteStatement
-        (QueryBuilder.readRequestToQuery readReq)
-        (QueryBuilder.mutateRequestToQuery mutateReq)
-        isInsert
-        (iAcceptMediaType ctxApiRequest)
-        (iPreferRepresentation ctxApiRequest)
-        pkCols
-        (configDbPreparedStatements ctxConfig)
-
-  liftEither $ WriteQueryResult queryTotal fields body <$> gucStatus <*> gucHeaders
+  lift . SQL.statement mempty $
+    Statements.prepareWrite
+      (QueryBuilder.readRequestToQuery readReq)
+      (QueryBuilder.mutateRequestToQuery mutateReq)
+      isInsert
+      (iAcceptMediaType ctxApiRequest)
+      (iPreferRepresentation ctxApiRequest)
+      pkCols
+      (configDbPreparedStatements ctxConfig)

 -- | Response with headers and status overridden from GUCs.
 gucResponse
@@ -632,9 +656,3 @@ binaryField RequestContext{..} readReq
 profileHeader :: ApiRequest -> Maybe HTTP.Header
 profileHeader ApiRequest{..} =
  (,) "Content-Profile" <$> (toUtf8 <$> iProfile)
-
-splitKeyValue :: ByteString -> (ByteString, ByteString)
-splitKeyValue kv =
-  (k, BS.tail v)
-  where
-    (k, v) = BS.break (== '=') kv
@@ -144,6 +144,9 @@ exampleConfigFile =
      |## Limit rows in response
      |# db-max-rows = 1000
      |
+      |## Allow getting the EXPLAIN plan through the `Accept: application/vnd.pgrst.plan` header
+      |# db-plan-enabled = false
+      |
      |## Number of open connections in the pool
      |db-pool = 10
      |
@@ -69,6 +69,7 @@ data AppConfig = AppConfig
  , configDbChannelEnabled      :: Bool
  , configDbExtraSearchPath     :: [Text]
  , configDbMaxRows             :: Maybe Integer
+  , configDbPlanEnabled         :: Bool
  , configDbPoolSize            :: Int
  , configDbPoolTimeout         :: NominalDiffTime
  , configDbPreRequest          :: Maybe QualifiedIdentifier
@@ -128,6 +129,7 @@ toText conf =
      ,("db-channel-enabled",            T.toLower . show . configDbChannelEnabled)
      ,("db-extra-search-path",      q . T.intercalate "," . configDbExtraSearchPath)
      ,("db-max-rows",                   maybe "\"\"" show . configDbMaxRows)
+      ,("db-plan-enabled",               T.toLower . show . configDbPlanEnabled)
      ,("db-pool",                       show . configDbPoolSize)
      ,("db-pool-timeout",               show . floor . configDbPoolTimeout)
      ,("db-pre-request",            q . maybe mempty dumpQi . configDbPreRequest)
@@ -216,6 +218,7 @@ parser optPath env dbSettings =
    <*> (maybe ["public"] splitOnCommas <$> optValue "db-extra-search-path")
    <*> optWithAlias (optInt "db-max-rows")
                     (optInt "max-rows")
+    <*> (fromMaybe False <$> optBool "db-plan-enabled")
    <*> (fromMaybe 10 <$> optInt "db-pool")
    <*> (fromIntegral . fromMaybe 3600 <$> optInt "db-pool-timeout")
    <*> (fmap toQi <$> optWithAlias (optString "db-pre-request")
@@ -9,6 +9,7 @@ module PostgREST.Config.PgVersion
  , pgVersion110
  , pgVersion112
  , pgVersion114
+  , pgVersion120
  , pgVersion121
  , pgVersion130
  , pgVersion140
@@ -50,6 +51,9 @@ pgVersion112 = PgVersion 110002 "11.2"
 pgVersion114 :: PgVersion
 pgVersion114 = PgVersion 110004 "11.4"

+pgVersion120 :: PgVersion
+pgVersion120 = PgVersion 120000 "12.0"
+
 pgVersion121 :: PgVersion
 pgVersion121 = PgVersion 120001 "12.1"

@@ -2,6 +2,9 @@

 module PostgREST.MediaType
  ( MediaType(..)
+  , MTPlanOption (..)
+  , MTPlanFormat (..)
+  , MTPlanAttrs(..)
  , toContentType
  , toMime
  , decodeMediaType
@@ -27,7 +30,18 @@ data MediaType
  | MTOctetStream
  | MTAny
  | MTOther ByteString
-  deriving (Eq)
+  | MTPlan MTPlanAttrs
+  deriving Eq
+
+data MTPlanAttrs = MTPlanAttrs MTPlanFormat [MTPlanOption]
+instance Eq MTPlanAttrs where
+  MTPlanAttrs {} == MTPlanAttrs {} = True -- we don't care about the attributes when comparing two MTPlan media types
+
+data MTPlanOption
+  = PlanAnalyze | PlanVerbose | PlanSettings | PlanBuffers | PlanWAL
+
+data MTPlanFormat
+  = PlanJSON | PlanText

 -- | Convert MediaType to a Content-Type HTTP Header
 toContentType :: MediaType -> Header
@@ -51,20 +65,47 @@ toMime MTUrlEncoded      = "application/x-www-form-urlencoded"
 toMime MTOctetStream     = "application/octet-stream"
 toMime MTAny             = "*/*"
 toMime (MTOther ct)      = ct
+toMime (MTPlan (MTPlanAttrs fmt opts)) = "application/vnd.pgrst.plan+" <> toMimePlanFormat fmt <>
+                            if null opts then mempty else "; options=" <> BS.intercalate "|" (toMimePlanOption <$> opts)
+
+toMimePlanOption :: MTPlanOption -> ByteString
+toMimePlanOption PlanAnalyze  = "analyze"
+toMimePlanOption PlanVerbose  = "verbose"
+toMimePlanOption PlanSettings = "settings"
+toMimePlanOption PlanBuffers  = "buffers"
+toMimePlanOption PlanWAL      = "wal"
+
+toMimePlanFormat :: MTPlanFormat -> ByteString
+toMimePlanFormat PlanJSON = "json"
+toMimePlanFormat PlanText = "text"

 -- | Convert from ByteString to MediaType. Warning: discards MIME parameters
 decodeMediaType :: BS.ByteString -> MediaType
-decodeMediaType ct =
-  case BS.takeWhile (/= BS.c2w ';') ct of
-    "application/json"                  -> MTApplicationJSON
-    "application/geo+json"              -> MTGeoJSON
-    "text/csv"                          -> MTTextCSV
-    "text/plain"                        -> MTTextPlain
-    "text/xml"                          -> MTTextXML
-    "application/openapi+json"          -> MTOpenAPI
-    "application/vnd.pgrst.object+json" -> MTSingularJSON
-    "application/vnd.pgrst.object"      -> MTSingularJSON
-    "application/x-www-form-urlencoded" -> MTUrlEncoded
-    "application/octet-stream"          -> MTOctetStream
-    "*/*"                               -> MTAny
-    ct'                                 -> MTOther ct'
+decodeMediaType mt =
+  case BS.split (BS.c2w ';') mt of
+    "application/json":_                   -> MTApplicationJSON
+    "application/geo+json":_               -> MTGeoJSON
+    "text/csv":_                           -> MTTextCSV
+    "text/plain":_                         -> MTTextPlain
+    "text/xml":_                           -> MTTextXML
+    "application/openapi+json":_           -> MTOpenAPI
+    "application/vnd.pgrst.object+json":_  -> MTSingularJSON
+    "application/vnd.pgrst.object":_       -> MTSingularJSON
+    "application/x-www-form-urlencoded":_  -> MTUrlEncoded
+    "application/octet-stream":_           -> MTOctetStream
+    "application/vnd.pgrst.plan":rest      -> getPlan PlanJSON rest
+    "application/vnd.pgrst.plan+json":rest -> getPlan PlanJSON rest
+    "application/vnd.pgrst.plan+text":rest -> getPlan PlanText rest
+    "*/*":_                                -> MTAny
+    other:_                                -> MTOther other
+    _                                      -> MTAny
+  where
+    getPlan fmt rest =
+     let opts = BS.split (BS.c2w '|') $ fromMaybe mempty (BS.stripPrefix "options=" =<< find (BS.isPrefixOf "options=") rest)
+         inOpts str = str `elem` opts in
+     MTPlan $ MTPlanAttrs fmt $
+      [PlanAnalyze  | inOpts "analyze" ] ++
+      [PlanVerbose  | inOpts "verbose" ] ++
+      [PlanSettings | inOpts "settings"] ++
+      [PlanBuffers  | inOpts "buffers" ] ++
+      [PlanWAL      | inOpts "wal"     ]
@@ -37,6 +37,7 @@ module PostgREST.Query.SqlFragment
  , sourceCTEName
  , unknownEncoder
  , intercalateSnippet
+  , explainF
  ) where

 import qualified Data.ByteString.Char8           as BS
@@ -50,6 +51,8 @@ import Text.InterpolatedString.Perl6 (qc)

 import PostgREST.DbStructure.Identifiers (FieldName,
                                          QualifiedIdentifier (..))
+import PostgREST.MediaType               (MTPlanFormat (..),
+                                          MTPlanOption (..))
 import PostgREST.RangeQuery              (NonnegRange, allRange,
                                          rangeLimit, rangeOffset)
 import PostgREST.Request.ReadQuery       (SelectItem)
@@ -367,3 +370,19 @@ unknownLiteral = unknownEncoder . encodeUtf8
 intercalateSnippet :: ByteString -> [SQL.Snippet] -> SQL.Snippet
 intercalateSnippet _ [] = mempty
 intercalateSnippet frag snippets = foldr1 (\a b -> a <> SQL.sql frag <> b) snippets
+
+explainF :: MTPlanFormat -> [MTPlanOption] -> SQL.Snippet -> SQL.Snippet
+explainF fmt opts snip =
+  "EXPLAIN (" <>
+    SQL.sql (BS.intercalate ", " (fmtPlanFmt fmt : (fmtPlanOpt <$> opts))) <>
+  ") " <> snip
+  where
+    fmtPlanOpt :: MTPlanOption -> BS.ByteString
+    fmtPlanOpt PlanAnalyze  = "ANALYZE"
+    fmtPlanOpt PlanVerbose  = "VERBOSE"
+    fmtPlanOpt PlanSettings = "SETTINGS"
+    fmtPlanOpt PlanBuffers  = "BUFFERS"
+    fmtPlanOpt PlanWAL      = "WAL"
+
+    fmtPlanFmt PlanJSON = "FORMAT JSON"
+    fmtPlanFmt PlanText = "FORMAT TEXT"
@@ -8,10 +8,11 @@ This module constructs single SQL statements that can be parametrized and prepar
 - It generates the body format and some headers of the final HTTP response.
 -}
 module PostgREST.Query.Statements
-  ( createWriteStatement
-  , createReadStatement
-  , callProcStatement
-  , createExplainStatement
+  ( prepareWrite
+  , prepareRead
+  , prepareCall
+  , preparePlanRows
+  , ResultSet (..)
  ) where

 import qualified Data.Aeson                        as JSON
@@ -32,23 +33,38 @@ import PostgREST.Error     (Error (..))
 import PostgREST.GucHeader (GucHeader)

 import PostgREST.DbStructure.Identifiers (FieldName)
-import PostgREST.MediaType               (MediaType (..))
+import PostgREST.MediaType               (MTPlanAttrs (..),
+                                          MTPlanFormat (..),
+                                          MediaType (..))
 import PostgREST.Query.SqlFragment
 import PostgREST.Request.Preferences

 import Protolude

-{-| The generic query result format used by API responses. The location header
-    is represented as a list of strings containing variable bindings like
-    @"k1=eq.42"@, or the empty list if there is no location header.
-}
-type ResultsWithCount = (Maybe Int64, Int64, [BS.ByteString], BS.ByteString, Either Error [GucHeader], Either Error (Maybe Status))
+-- | Standard result set format used for all queries
+data ResultSet
+  = RSStandard
+  { rsTableTotal :: Maybe Int64
+  -- ^ count of all the table rows
+  , rsQueryTotal :: Int64
+  -- ^ count of the query rows
+  , rsLocation   :: [(BS.ByteString, BS.ByteString)]
+  -- ^ The Location header(only used for inserts) is represented as a list of strings containing
+  -- variable bindings like @"k1=eq.42"@, or the empty list if there is no location header.
+  , rsBody       :: BS.ByteString
+  -- ^ the aggregated body of the query
+  , rsGucHeaders :: Either Error [GucHeader]
+  -- ^ the HTTP headers to be added to the response
+  , rsGucStatus  :: Either Error (Maybe Status)
+  -- ^ the HTTP status to be added to the response
+  }
+  | RSPlan BS.ByteString -- ^ the plan of the query

-createWriteStatement :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType ->
-                        PreferRepresentation -> [Text] -> Bool ->
-                        SQL.Statement () ResultsWithCount
-createWriteStatement selectQuery mutateQuery isInsert mediaType rep pKeys =
-  SQL.dynamicallyParameterized snippet decodeStandard
+
+prepareWrite :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType ->
+                PreferRepresentation -> [Text] -> Bool -> SQL.Statement () ResultSet
+prepareWrite selectQuery mutateQuery isInsert mediaType rep pKeys =
+  SQL.dynamicallyParameterized (mtSnippet mediaType snippet) decodeIt
 where
  snippet =
    "WITH " <> SQL.sql sourceCTEName <> " AS (" <> mutateQuery <> ") " <>
@@ -84,14 +100,14 @@ createWriteStatement selectQuery mutateQuery isInsert mediaType rep pKeys =
    | rep /= Full = SQL.sql ("SELECT * FROM " <> sourceCTEName)
    | otherwise   = selectQuery

-  decodeStandard :: HD.Result ResultsWithCount
-  decodeStandard =
-   fromMaybe (Nothing, 0, [], mempty, Right [], Right Nothing) <$> HD.rowMaybe standardRow
+  decodeIt :: HD.Result ResultSet
+  decodeIt = case mediaType of
+    MTPlan{} -> planRow
+    _        -> fromMaybe (RSStandard Nothing 0 mempty mempty (Right []) (Right Nothing)) <$> HD.rowMaybe (standardRow False)

-createReadStatement :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> Maybe FieldName -> Bool ->
-                       SQL.Statement () ResultsWithCount
-createReadStatement selectQuery countQuery countTotal mediaType binaryField =
-  SQL.dynamicallyParameterized snippet decodeStandard
+prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> Maybe FieldName -> Bool -> SQL.Statement () ResultSet
+prepareRead selectQuery countQuery countTotal mediaType binaryField =
+  SQL.dynamicallyParameterized (mtSnippet mediaType snippet) decodeIt
 where
  snippet =
    "WITH " <>
@@ -100,7 +116,6 @@ createReadStatement selectQuery countQuery countTotal mediaType binaryField =
    SQL.sql ("SELECT " <>
      countResultF <> " AS total_result_set, " <>
      "pg_catalog.count(_postgrest_t) AS page_total, " <>
-      noLocationF <> " AS header, " <>
      bodyF <> " AS body, " <>
      responseHeadersF <> " AS response_headers, " <>
      responseStatusF <> " AS response_status " <>
@@ -116,27 +131,16 @@ createReadStatement selectQuery countQuery countTotal mediaType binaryField =
    | isJust binaryField                           = asBinaryF $ fromJust binaryField
    | otherwise                                    = asJsonF False

-  decodeStandard :: HD.Result ResultsWithCount
-  decodeStandard =
-    HD.singleRow standardRow
+  decodeIt :: HD.Result ResultSet
+  decodeIt = case mediaType of
+    MTPlan{} -> planRow
+    _        -> HD.singleRow $ standardRow True

-{-| Read and Write api requests use a similar response format which includes
-    various record counts and possible location header. This is the decoder
-    for that common type of query.
-}
-standardRow :: HD.Row ResultsWithCount
-standardRow = (,,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
-                      <*> arrayColumn HD.bytea <*> column HD.bytea
-                      <*> (fromMaybe (Right []) <$> nullableColumn decodeGucHeaders)
-                      <*> (fromMaybe (Right Nothing) <$> nullableColumn decodeGucStatus)
-
-type ProcResults = (Maybe Int64, Int64, ByteString, Either Error [GucHeader], Either Error (Maybe Status))
-
-callProcStatement :: Bool -> Bool -> SQL.Snippet -> SQL.Snippet -> SQL.Snippet -> Bool ->
-                     MediaType -> Bool -> Maybe FieldName -> Bool ->
-                     SQL.Statement () ProcResults
-callProcStatement returnsScalar returnsSingle callProcQuery selectQuery countQuery countTotal mediaType multObjects binaryField =
-  SQL.dynamicallyParameterized snippet decodeProc
+prepareCall :: Bool -> Bool -> SQL.Snippet -> SQL.Snippet -> SQL.Snippet -> Bool ->
+               MediaType -> Bool -> Maybe FieldName -> Bool ->
+               SQL.Statement () ResultSet
+prepareCall returnsScalar returnsSingle callProcQuery selectQuery countQuery countTotal mediaType multObjects binaryField =
+  SQL.dynamicallyParameterized (mtSnippet mediaType snippet) decodeIt
  where
    snippet =
      "WITH " <> SQL.sql sourceCTEName <> " AS (" <> callProcQuery <> ") " <>
@@ -161,35 +165,42 @@ callProcStatement returnsScalar returnsSingle callProcQuery selectQuery countQue
     | returnsSingle && not multObjects             = asJsonSingleF returnsScalar
     | otherwise                                    = asJsonF returnsScalar

-    decodeProc :: HD.Result ProcResults
-    decodeProc =
-      fromMaybe (Just 0, 0, mempty, defGucHeaders, defGucStatus) <$> HD.rowMaybe procRow
-      where
-        defGucHeaders = Right []
-        defGucStatus  = Right Nothing
-        procRow = (,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
-                         <*> column HD.bytea
-                         <*> (fromMaybe defGucHeaders <$> nullableColumn decodeGucHeaders)
-                         <*> (fromMaybe defGucStatus <$> nullableColumn decodeGucStatus)
+    decodeIt :: HD.Result ResultSet
+    decodeIt = case mediaType of
+      MTPlan{} -> planRow
+      _        -> fromMaybe (RSStandard (Just 0) 0 mempty mempty (Right []) (Right Nothing)) <$> HD.rowMaybe (standardRow True)

-createExplainStatement :: SQL.Snippet -> Bool -> SQL.Statement () (Maybe Int64)
-createExplainStatement countQuery =
-  SQL.dynamicallyParameterized snippet decodeExplain
+preparePlanRows :: SQL.Snippet -> Bool -> SQL.Statement () (Maybe Int64)
+preparePlanRows countQuery =
+  SQL.dynamicallyParameterized snippet decodeIt
  where
-    snippet = "EXPLAIN (FORMAT JSON) " <> countQuery
-    -- |
-    -- An `EXPLAIN (FORMAT JSON) select * from items;` output looks like this:
-    -- [{
-    --   "Plan": {
-    --     "Node Type": "Seq Scan", "Parallel Aware": false, "Relation Name": "items",
-    --     "Alias": "items", "Startup Cost": 0.00, "Total Cost": 32.60,
-    --     "Plan Rows": 2260,"Plan Width": 8} }]
-    -- We only obtain the Plan Rows here.
-    decodeExplain :: HD.Result (Maybe Int64)
-    decodeExplain =
+    snippet = explainF PlanJSON mempty countQuery
+    decodeIt :: HD.Result (Maybe Int64)
+    decodeIt =
      let row = HD.singleRow $ column HD.bytea in
      (^? L.nth 0 . L.key "Plan" .  L.key "Plan Rows" . L._Integral) <$> row

+standardRow :: Bool -> HD.Row ResultSet
+standardRow noLocation =
+  RSStandard <$> nullableColumn HD.int8 <*> column HD.int8
+             <*> (if noLocation then pure mempty else fmap splitKeyValue <$> arrayColumn HD.bytea) <*> column HD.bytea
+             <*> (fromMaybe (Right []) <$> nullableColumn decodeGucHeaders)
+             <*> (fromMaybe (Right Nothing) <$> nullableColumn decodeGucStatus)
+  where
+    splitKeyValue :: ByteString -> (ByteString, ByteString)
+    splitKeyValue kv =
+      let (k, v) = BS.break (== '=') kv in
+      (k, BS.tail v)
+
+mtSnippet :: MediaType -> SQL.Snippet -> SQL.Snippet
+mtSnippet mediaType snippet = case mediaType of
+  MTPlan (MTPlanAttrs fmt opts) -> explainF fmt opts snippet
+  _                             -> snippet
+
+-- | We use rowList because when doing EXPLAIN (FORMAT TEXT), the result comes as many rows. FORMAT JSON comes as one.
+planRow :: HD.Result ResultSet
+planRow = RSPlan . BS.unlines <$> HD.rowList (column HD.bytea)
+
 decodeGucHeaders :: HD.Value (Either Error [GucHeader])
 decodeGucHeaders = first (const GucHeadersError) . JSON.eitherDecode . LBS.fromStrict <$> HD.bytea

@@ -51,7 +51,9 @@ import PostgREST.DbStructure.Identifiers (FieldName,
                                          Schema)
 import PostgREST.DbStructure.Proc        (ProcDescription (..),
                                          ProcParam (..), ProcsMap)
-import PostgREST.MediaType               (MediaType (..))
+import PostgREST.MediaType               (MTPlanAttrs (..),
+                                          MTPlanFormat (..),
+                                          MediaType (..))
 import PostgREST.RangeQuery              (NonnegRange, allRange,
                                          hasLimitZero,
                                          limitZeroRange,
@@ -426,7 +428,8 @@ requestMediaTypes conf action path =
        ++ rawMediaTypes
        ++ [MTOpenAPI | pathIsRootSpec path]
    defaultMediaTypes =
-      [MTApplicationJSON, MTSingularJSON, MTGeoJSON, MTTextCSV]
+      [MTApplicationJSON, MTSingularJSON, MTGeoJSON, MTTextCSV] ++
+      [MTPlan $ MTPlanAttrs PlanJSON mempty | configDbPlanEnabled conf]
    rawMediaTypes = configRawMediaTypes conf `union` [MTOctetStream, MTTextPlain, MTTextXML]

 {-|