Files
postgrest/src/PostgREST/Query/Statements.hs
T
steve-chavez 16f2849724 refactor: wrap proc logic in asJsonF
* add explicit logic for returning setof scalars
2023-04-04 19:55:10 -05:00

205 lines
8.1 KiB
Haskell

{-|
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
( prepareWrite
, prepareRead
, prepareCall
, preparePlanRows
, ResultSet (..)
) where
import qualified Data.Aeson.Lens as L
import qualified Data.ByteString.Char8 as BS
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 PostgREST.ApiRequest.Preferences
import PostgREST.MediaType (MTPlanAttrs (..),
MTPlanFormat (..),
MediaType (..),
getMediaType)
import PostgREST.Query.SqlFragment
import PostgREST.SchemaCache.Identifiers (FieldName)
import Protolude
-- | 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 :: Maybe BS.ByteString
-- ^ the HTTP headers to be added to the response
, rsGucStatus :: Maybe Text
-- ^ the HTTP status to be added to the response
}
| RSPlan BS.ByteString -- ^ the plan of the query
prepareWrite :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType ->
PreferRepresentation -> [Text] -> Bool -> SQL.Statement () ResultSet
prepareWrite selectQuery mutateQuery isInsert mt rep pKeys =
SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt
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 = "''"
| getMediaType mt == MTTextCSV = asCsvF
| getMediaType mt == MTGeoJSON = asGeoJsonF
| getMediaType mt == MTSingularJSON = asJsonSingleF False
| otherwise = asJsonF False False 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
decodeIt :: HD.Result ResultSet
decodeIt = case mt of
MTPlan{} -> planRow
_ -> fromMaybe (RSStandard Nothing 0 mempty mempty Nothing Nothing) <$> HD.rowMaybe (standardRow False)
prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> Maybe FieldName -> Bool -> SQL.Statement () ResultSet
prepareRead selectQuery countQuery countTotal mt binaryField =
SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt
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, " <>
bodyF <> " AS body, " <>
responseHeadersF <> " AS response_headers, " <>
responseStatusF <> " AS response_status " <>
"FROM ( SELECT * FROM " <> sourceCTEName <> " ) _postgrest_t")
(countCTEF, countResultF) = countF countQuery countTotal
bodyF
| getMediaType mt == MTTextCSV = asCsvF
| getMediaType mt == MTSingularJSON = asJsonSingleF False
| getMediaType mt == MTGeoJSON = asGeoJsonF
| isJust binaryField && getMediaType mt == MTTextXML = asXmlF $ fromJust binaryField
| isJust binaryField = asBinaryF $ fromJust binaryField
| otherwise = asJsonF False False False
decodeIt :: HD.Result ResultSet
decodeIt = case mt of
MTPlan{} -> planRow
_ -> HD.singleRow $ standardRow True
prepareCall :: Bool -> Bool -> Bool -> SQL.Snippet -> SQL.Snippet -> SQL.Snippet -> Bool ->
MediaType -> Maybe FieldName -> Bool ->
SQL.Statement () ResultSet
prepareCall returnsScalar returnsSingleComposite returnsSetOfScalar callProcQuery selectQuery countQuery countTotal mt binaryField =
SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt
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
| getMediaType mt == MTSingularJSON = asJsonSingleF returnsScalar
| getMediaType mt == MTTextCSV = asCsvF
| getMediaType mt == MTGeoJSON = asGeoJsonF
| isJust binaryField && getMediaType mt == MTTextXML = asXmlF $ fromJust binaryField
| isJust binaryField = asBinaryF $ fromJust binaryField
| otherwise = asJsonF returnsScalar returnsSetOfScalar returnsSingleComposite
decodeIt :: HD.Result ResultSet
decodeIt = case mt of
MTPlan{} -> planRow
_ -> fromMaybe (RSStandard (Just 0) 0 mempty mempty Nothing Nothing) <$> HD.rowMaybe (standardRow True)
preparePlanRows :: SQL.Snippet -> Bool -> SQL.Statement () (Maybe Int64)
preparePlanRows countQuery =
SQL.dynamicallyParameterized snippet decodeIt
where
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
<*> nullableColumn HD.bytea
<*> nullableColumn HD.text
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)
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