postgrest/src/PostgREST/Query/Statements.hs

{-|
Module      : PostgREST.Query.Statements
Description : PostgREST single SQL statements.

This module constructs single SQL statements that can be parametrized and prepared.

- It consumes the SqlQuery types generated by the QueryBuilder module.
- It generates the body format and some headers of the final HTTP response.
-}
module PostgREST.Query.Statements
  ( createWriteStatement
  , createReadStatement
  , callProcStatement
  , createExplainStatement
  ) where

import qualified Data.Aeson                        as JSON
import qualified Data.Aeson.Lens                   as L
import qualified Data.ByteString.Char8             as BS
import qualified Data.ByteString.Lazy              as LBS
import qualified Hasql.Decoders                    as HD
import qualified Hasql.DynamicStatements.Snippet   as SQL
import qualified Hasql.DynamicStatements.Statement as SQL
import qualified Hasql.Statement                   as SQL

import Control.Lens              ((^?))
import Data.Maybe                (fromJust)
import Data.Text.Read            (decimal)
import Network.HTTP.Types.Status (Status)

import PostgREST.Error     (Error (..))
import PostgREST.GucHeader (GucHeader)

import PostgREST.DbStructure.Identifiers (FieldName)
import PostgREST.MediaType               (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))

createWriteStatement :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType ->
                        PreferRepresentation -> [Text] -> Bool ->
                        SQL.Statement () ResultsWithCount
createWriteStatement selectQuery mutateQuery isInsert mediaType rep pKeys =
  SQL.dynamicallyParameterized snippet decodeStandard
 where
  snippet =
    "WITH " <> SQL.sql sourceCTEName <> " AS (" <> mutateQuery <> ") " <>
    SQL.sql (
    "SELECT " <>
      "'' AS total_result_set, " <>
      "pg_catalog.count(_postgrest_t) AS page_total, " <>
      locF <> " AS header, " <>
      bodyF <> " AS body, " <>
      responseHeadersF <> " AS response_headers, " <>
      responseStatusF  <> " AS response_status "
    ) <>
    "FROM (" <> selectF <> ") _postgrest_t"

  locF =
    if isInsert && rep == HeadersOnly
      then BS.unwords [
        "CASE WHEN pg_catalog.count(_postgrest_t) = 1",
          "THEN coalesce(" <> locationF pKeys <> ", " <> noLocationF <> ")",
          "ELSE " <> noLocationF,
        "END"]
      else noLocationF

  bodyF
    | rep /= Full                 = "''"
    | mediaType == MTTextCSV      = asCsvF
    | mediaType == MTGeoJSON      = asGeoJsonF
    | mediaType == MTSingularJSON = asJsonSingleF False
    | otherwise                   = asJsonF False

  selectF
    -- prevent using any of the column names in ?select= when no response is returned from the CTE
    | rep /= Full = SQL.sql ("SELECT * FROM " <> sourceCTEName)
    | otherwise   = selectQuery

  decodeStandard :: HD.Result ResultsWithCount
  decodeStandard =
   fromMaybe (Nothing, 0, [], mempty, Right [], Right Nothing) <$> HD.rowMaybe standardRow

createReadStatement :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> Maybe FieldName -> Bool ->
                       SQL.Statement () ResultsWithCount
createReadStatement selectQuery countQuery countTotal mediaType binaryField =
  SQL.dynamicallyParameterized snippet decodeStandard
 where
  snippet =
    "WITH " <>
    SQL.sql sourceCTEName <> " AS ( " <> selectQuery <> " ) " <>
    countCTEF <> " " <>
    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 " <>
    "FROM ( SELECT * FROM " <> sourceCTEName <> " ) _postgrest_t")

  (countCTEF, countResultF) = countF countQuery countTotal

  bodyF
    | mediaType == MTTextCSV                       = asCsvF
    | mediaType == MTSingularJSON                  = asJsonSingleF False
    | mediaType == MTGeoJSON                       = asGeoJsonF
    | isJust binaryField && mediaType == MTTextXML = asXmlF $ fromJust binaryField
    | isJust binaryField                           = asBinaryF $ fromJust binaryField
    | otherwise                                    = asJsonF False

  decodeStandard :: HD.Result ResultsWithCount
  decodeStandard =
    HD.singleRow standardRow

{-| 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
  where
    snippet =
      "WITH " <> SQL.sql sourceCTEName <> " AS (" <> callProcQuery <> ") " <>
      countCTEF <>
      SQL.sql (
      "SELECT " <>
        countResultF <> " AS total_result_set, " <>
        "pg_catalog.count(_postgrest_t) AS page_total, " <>
        bodyF <> " AS body, " <>
        responseHeadersF <> " AS response_headers, " <>
        responseStatusF <> " AS response_status ") <>
      "FROM (" <> selectQuery <> ") _postgrest_t"

    (countCTEF, countResultF) = countF countQuery countTotal

    bodyF
     | mediaType == MTSingularJSON                  = asJsonSingleF returnsScalar
     | mediaType == MTTextCSV                       = asCsvF
     | mediaType == MTGeoJSON                       = asGeoJsonF
     | isJust binaryField && mediaType == MTTextXML = asXmlF $ fromJust binaryField
     | isJust binaryField                           = asBinaryF $ fromJust binaryField
     | 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)

createExplainStatement :: SQL.Snippet -> Bool -> SQL.Statement () (Maybe Int64)
createExplainStatement countQuery =
  SQL.dynamicallyParameterized snippet decodeExplain
  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 =
      let row = HD.singleRow $ column HD.bytea in
      (^? L.nth 0 . L.key "Plan" .  L.key "Plan Rows" . L._Integral) <$> row

decodeGucHeaders :: HD.Value (Either Error [GucHeader])
decodeGucHeaders = first (const GucHeadersError) . JSON.eitherDecode . LBS.fromStrict <$> HD.bytea

decodeGucStatus :: HD.Value (Either Error (Maybe Status))
decodeGucStatus = first (const GucStatusError) . fmap (Just . toEnum . fst) . decimal <$> HD.text

column :: HD.Value a -> HD.Row a
column = HD.column . HD.nonNullable

nullableColumn :: HD.Value a -> HD.Row (Maybe a)
nullableColumn = HD.column . HD.nullable

arrayColumn :: HD.Value a -> HD.Row [a]
arrayColumn = column . HD.listArray . HD.nonNullable