refactor: put callProc core query to QueryBuilder

* Move set local queries to QueryBuilder

* Move unquoted to Middleware

src/PostgREST/App.hs

@@ -55,7 +55,8 @@ import PostgREST.Error            (PgError (..), SimpleError (..),
 import PostgREST.Middleware
 import PostgREST.OpenAPI
 import PostgREST.Parsers          (pRequestColumns)
-import PostgREST.QueryBuilder     (requestToCountQuery,
+import PostgREST.QueryBuilder     (requestToCallProcQuery,
+                                   requestToCountQuery,
                                    requestToQuery)
 import PostgREST.RangeQuery       (allRange, contentRangeH,
                                    rangeStatusHeader)
@@ -270,9 +271,7 @@ app dbStructure proc cols conf apiRequest =
               return $ responseLBS status200 [allOrigins, allowH] mempty
 
         (ActionInvoke invMethod, TargetProc qi _, Just pJson) ->
-          let returnsScalar = case proc of
+          let returnsScalar = maybe False procReturnsScalar proc
-                Just ProcDescription{pdReturnType = (Single (Scalar _))} -> True
-                _ -> False
               rpcBinaryField = if returnsScalar
                                  then Right Nothing
                                  else binaryField contentType rawContentTypes =<< fldNames
@@ -280,19 +279,18 @@ app dbStructure proc cols conf apiRequest =
           case parts of
             Left errorResponse -> return errorResponse
             Right ((q, cq), bField) -> do
-              let singular = contentType == CTSingularJSON
+              let
-              row <- H.statement (toS $ pjRaw pJson) $
+                pq = requestToCallProcQuery qi (specifiedProcArgs cols proc) returnsScalar $
-                callProcStatement qi (specifiedProcArgs cols proc) returnsScalar q cq shouldCount
+                        iPreferSingleObjectParameter apiRequest
-                         singular (iPreferSingleObjectParameter apiRequest)
+                stm = callProcStatement returnsScalar pq q cq shouldCount (contentType == CTSingularJSON)
-                         (contentType == CTTextCSV)
+                        (contentType == CTTextCSV) (contentType `elem` rawContentTypes) bField (pgVersion dbStructure)
-                         (contentType `elem` rawContentTypes) bField
+              row <- H.statement (toS $ pjRaw pJson) stm
-                         (pgVersion dbStructure)
               let (tableTotal, queryTotal, body, gucHeaders) = row
                   (status, contentRange) = rangeStatusHeader topLevelRange queryTotal tableTotal
               case gucHeaders of
                 Left _ -> return . errorResponseFor $ GucHeadersError
                 Right hs ->
-                  if singular && queryTotal /= 1
+                  if contentType == CTSingularJSON && queryTotal /= 1
                     then do
                       HT.condemn
                       return . errorResponseFor . singularityError $ queryTotal

src/PostgREST/Middleware.hs

@@ -10,6 +10,8 @@ module PostgREST.Middleware where
 
 import qualified Data.Aeson          as JSON
 import qualified Data.HashMap.Strict as M
+import           Data.Scientific     (FPFormat (..), formatScientific,
+                                      isInteger)
 import qualified Hasql.Transaction   as H
 
 import Network.Wai                   (Application, Response)
@@ -24,8 +26,7 @@ import PostgREST.Auth         (JWTAttempt (..))
 import PostgREST.Config       (AppConfig (..), corsPolicy)
 import PostgREST.Error        (SimpleError (JwtTokenInvalid, JwtTokenMissing),
                                errorResponseFor)
-import PostgREST.QueryBuilder (pgFmtSetLocal, pgFmtSetLocalSearchPath,
+import PostgREST.QueryBuilder (setLocalQuery, setLocalSearchPathQuery)
-                               unquoted)
 import Protolude
 
 runWithClaims :: AppConfig -> JWTAttempt ->
@@ -41,13 +42,13 @@ runWithClaims conf eClaims app req =
       mapM_ H.sql customReqCheck
       app req
       where
-        headersSql = pgFmtSetLocal "request.header." <$> iHeaders req
+        headersSql = setLocalQuery "request.header." <$> iHeaders req
-        cookiesSql = pgFmtSetLocal "request.cookie." <$> iCookies req
+        cookiesSql = setLocalQuery "request.cookie." <$> iCookies req
-        claimsSql = pgFmtSetLocal "request.jwt.claim." <$> [(c,unquoted v) | (c,v) <- M.toList claimsWithRole]
+        claimsSql = setLocalQuery "request.jwt.claim." <$> [(c,unquoted v) | (c,v) <- M.toList claimsWithRole]
-        appSettingsSql = pgFmtSetLocal mempty <$> configSettings conf
+        appSettingsSql = setLocalQuery mempty <$> configSettings conf
         setRoleSql = maybeToList $ (\x ->
-          pgFmtSetLocal mempty ("role", unquoted x)) <$> M.lookup "role" claimsWithRole
+          setLocalQuery mempty ("role", unquoted x)) <$> M.lookup "role" claimsWithRole
-        setSearchPathSql = pgFmtSetLocalSearchPath $ configSchema conf : configExtraSearchPath conf
+        setSearchPathSql = setLocalSearchPathQuery $ configSchema conf : configExtraSearchPath conf
         -- role claim defaults to anon if not specified in jwt
         claimsWithRole = M.union claims (M.singleton "role" anon)
         anon = JSON.String . toS $ configAnonRole conf
@@ -58,3 +59,10 @@ defaultMiddle =
     gzip def
   . cors corsPolicy
   . staticPolicy (only [("favicon.ico", "static/favicon.ico")])
+
+unquoted :: JSON.Value -> Text
+unquoted (JSON.String t) = t
+unquoted (JSON.Number n) =
+  toS $ formatScientific Fixed (if isInteger n then Just 0 else Nothing) n
+unquoted (JSON.Bool b) = show b
+unquoted v = toS $ JSON.encode v

src/PostgREST/QueryBuilder.hs

@@ -4,20 +4,18 @@
 {-# OPTIONS_GHC -fno-warn-orphans  #-}
 {-|
 Module      : PostgREST.QueryBuilder
-Description : PostgREST SQL fragments generating functions.
+Description : PostgREST SQL queries generating functions.
 
 This module provides functions to consume data types that
-represent database objects (e.g. Relation, Schema, SqlQuery)
+represent database objects (e.g. Relation, Schema) and SqlFragment
-and produces SQL fragments.
+to produce SqlQuery type outputs.
-
-Any function that outputs a SQL fragment should be in this module.
 -}
 module PostgREST.QueryBuilder (
     requestToQuery
   , requestToCountQuery
-  , unquoted
+  , requestToCallProcQuery
-  , pgFmtSetLocal
+  , setLocalQuery
-  , pgFmtSetLocalSearchPath
+  , setLocalSearchPathQuery
   ) where
 
 import qualified Data.Set as S
@@ -135,3 +133,47 @@ requestToQuery schema _ (DbMutate (Delete mainTbl logicForest returnings)) =
     ]
   where
     qi = QualifiedIdentifier schema mainTbl
+
+requestToCallProcQuery :: QualifiedIdentifier -> [PgArg] -> Bool -> Bool -> SqlQuery
+requestToCallProcQuery qi pgArgs returnsScalar paramsAsSingleObject =
+  unwords [
+    "WITH",
+    argsRecord,
+    sourceBody ]
+  where
+    (argsRecord, args)
+      | null pgArgs = (ignoredBody, "")
+      | paramsAsSingleObject = ("_args_record AS (SELECT NULL)", "$1::json")
+      | otherwise = (
+          unwords [
+            normalizedBody <> ",",
+            "_args_record AS (",
+              "SELECT * FROM json_to_recordset(" <> selectBody <> ") AS _(" <>
+                intercalate ", " ((\a -> pgFmtIdent (pgaName a) <> " " <> pgaType a) <$> pgArgs) <> ")",
+            ")"]
+         , intercalate ", " ((\a -> pgFmtIdent (pgaName a) <> " := _args_record." <> pgFmtIdent (pgaName a)) <$> pgArgs))
+
+    sourceBody :: SqlFragment
+    sourceBody
+      | paramsAsSingleObject || null pgArgs =
+          if returnsScalar
+            then "SELECT " <> callIt <> " AS _scalar_res"
+            else "SELECT * FROM " <> callIt
+      | otherwise =
+          if returnsScalar
+            then "SELECT " <> callIt <> " AS _scalar_res FROM _args_record"
+            else unwords [
+              "SELECT _.*",
+              "FROM _args_record,",
+              "LATERAL ( SELECT * FROM " <> callIt <> " ) _" ]
+
+    callIt :: SqlFragment
+    callIt = fromQi qi <> "(" <> args <> ")"
+
+setLocalQuery :: Text -> (Text, Text) -> SqlQuery
+setLocalQuery prefix (k, v) =
+  "SET LOCAL " <> pgFmtIdent (prefix <> k) <> " = " <> pgFmtLit v <> ";"
+
+setLocalSearchPathQuery :: [Text] -> SqlQuery
+setLocalSearchPathQuery vals =
+  "SET LOCAL search_path = " <> intercalate ", " (pgFmtLit <$> vals) <> ";"

src/PostgREST/QueryBuilder/Private.hs

@@ -2,15 +2,13 @@
 {-|
 Module      : PostgREST.QueryBuilder.Private
 Description : Helper functions for PostgREST.QueryBuilder.
+
+Any function that outputs a SqlFragment should be in this module.
 -}
 module PostgREST.QueryBuilder.Private where
 
-import qualified Data.Aeson                    as JSON
 import qualified Data.HashMap.Strict           as HM
 import           Data.Maybe
-import           Data.Scientific               (FPFormat (..),
-                                                formatScientific,
-                                                isInteger)
 import           Data.Text                     (intercalate,
                                                 isInfixOf, replace,
                                                 toLower, unwords)
@@ -21,12 +19,9 @@ import           Protolude                     hiding (cast,
                                                 intercalate, replace)
 import           Text.InterpolatedString.Perl6 (qc)
 
-noLocationF :: Text
+noLocationF :: SqlFragment
 noLocationF = "array[]::text[]"
 
-removeSourceCTESchema :: Schema -> TableName -> QualifiedIdentifier
-removeSourceCTESchema schema tbl = QualifiedIdentifier (if tbl == sourceCTEName then "" else schema) tbl
-
 -- Due to the use of the `unknown` encoder we need to cast '$1' when the value is not used in the main query
 -- otherwise the query will err with a `could not determine data type of parameter $1`.
 -- This happens because `unknown` relies on the context to determine the value type.
@@ -185,20 +180,8 @@ pgFmtAs fName jp Nothing = case jOp <$> lastMay jp of
   Nothing -> ""
 pgFmtAs _ _ (Just alias) = " AS " <> pgFmtIdent alias
 
-pgFmtSetLocal :: Text -> (Text, Text) -> SqlFragment
-pgFmtSetLocal prefix (k, v) =
-  "SET LOCAL " <> pgFmtIdent (prefix <> k) <> " = " <> pgFmtLit v <> ";"
-
-pgFmtSetLocalSearchPath :: [Text] -> SqlFragment
-pgFmtSetLocalSearchPath vals =
-  "SET LOCAL search_path = " <> intercalate ", " (pgFmtLit <$> vals) <> ";"
-
 trimNullChars :: Text -> Text
 trimNullChars = T.takeWhile (/= '\x0')
 
-unquoted :: JSON.Value -> Text
+removeSourceCTESchema :: Schema -> TableName -> QualifiedIdentifier
-unquoted (JSON.String t) = t
+removeSourceCTESchema schema tbl = QualifiedIdentifier (if tbl == sourceCTEName then "" else schema) tbl
-unquoted (JSON.Number n) =
-  toS $ formatScientific Fixed (if isInteger n then Just 0 else Nothing) n
-unquoted (JSON.Bool b) = show b
-unquoted v = toS $ JSON.encode v

src/PostgREST/Statements.hs

@@ -7,7 +7,7 @@ This module constructs single SQL statements that can be parametrized and prepar
 - It consumes the SqlQuery types generated by the QueryBuilder module.
 - It generates the body format and some headers of the final HTTP response.
 
-TODO: Currently, createReadStatement is not using prepared statements, see https://github.com/PostgREST/postgrest/issues/718.
+TODO: Currently, createReadStatement is not using prepared statements. See https://github.com/PostgREST/postgrest/issues/718.
 -}
 module PostgREST.Statements (
     createWriteStatement
@@ -117,18 +117,14 @@ standardRow = (,,,) <$> nullableColumn HD.int8 <*> column HD.int8
 
 type ProcResults = (Maybe Int64, Int64, ByteString, Either Text [GucHeader])
 
-callProcStatement :: QualifiedIdentifier -> [PgArg] -> Bool -> SqlQuery -> SqlQuery -> Bool ->
+callProcStatement :: Bool -> SqlQuery -> SqlQuery -> SqlQuery -> Bool ->
-                     Bool -> Bool -> Bool -> Bool -> Maybe FieldName -> PgVersion ->
+                     Bool -> Bool -> Bool -> Maybe FieldName -> PgVersion ->
                      H.Statement ByteString ProcResults
-callProcStatement qi pgArgs returnsScalar selectQuery countQuery countTotal isSingle paramsAsSingleObject asCsv asBinary binaryField pgVer =
+callProcStatement returnsScalar callProcQuery selectQuery countQuery countTotal isSingle asCsv asBinary binaryField pgVer =
   unicodeStatement sql (param HE.unknown) decodeProc True
   where
-    sql =[qc|
+    sql = [qc|
-      WITH
+      WITH {sourceCTEName} AS ({callProcQuery})
-      {argsRecord},
-      {sourceCTEName} AS (
-        {sourceBody}
-      )
       SELECT
         {countResultF} AS total_result_set,
         pg_catalog.count(_postgrest_t) AS page_total,
@@ -136,31 +132,6 @@ callProcStatement qi pgArgs returnsScalar selectQuery countQuery countTotal isSi
         {responseHeaders} AS response_headers
       FROM ({selectQuery}) _postgrest_t;|]
 
-    (argsRecord, args)
-      | paramsAsSingleObject = ("_args_record AS (SELECT NULL)", "$1::json")
-      | null pgArgs = (ignoredBody, "")
-      | otherwise = (
-          unwords [
-            normalizedBody <> ",",
-            "_args_record AS (",
-              "SELECT * FROM json_to_recordset(" <> selectBody <> ") AS _(" <>
-                intercalate ", " ((\a -> pgFmtIdent (pgaName a) <> " " <> pgaType a) <$> pgArgs) <> ")",
-            ")"]
-         , intercalate ", " ((\a -> pgFmtIdent (pgaName a) <> " := _args_record." <> pgFmtIdent (pgaName a)) <$> pgArgs))
-
-    sourceBody :: SqlFragment
-    sourceBody
-      | paramsAsSingleObject || null pgArgs =
-          if returnsScalar
-            then [qc| SELECT {fromQi qi}({args}) |]
-            else [qc| SELECT * FROM {fromQi qi}({args}) |]
-      | otherwise =
-          if returnsScalar
-            then [qc| SELECT {fromQi qi}({args}) FROM _args_record |]
-            else [qc| SELECT _.*
-                      FROM _args_record,
-                      LATERAL ( SELECT * FROM {fromQi qi}({args}) ) _ |]
-
     bodyF
      | returnsScalar = scalarBodyF
      | isSingle = asJsonSingleF
@@ -169,16 +140,15 @@ callProcStatement qi pgArgs returnsScalar selectQuery countQuery countTotal isSi
      | otherwise = asJsonF
 
     scalarBodyF
-     | asBinary = asBinaryF _procName
+     | asBinary = asBinaryF "_scalar_res"
      | otherwise = unwords [
         "CASE",
           "WHEN pg_catalog.count(_postgrest_t) = 1",
-            "THEN (json_agg(_postgrest_t." <> pgFmtIdent _procName <> ")->0)::character varying",
+            "THEN (json_agg(_postgrest_t._scalar_res)->0)::character varying",
-            "ELSE (json_agg(_postgrest_t." <> pgFmtIdent _procName <> "))::character varying",
+            "ELSE (json_agg(_postgrest_t._scalar_res))::character varying",
         "END"]
 
     countResultF = if countTotal then "( "<> countQuery <> ")" else "null::bigint" :: Text
-    _procName = qiName qi
     responseHeaders =
       if pgVer >= pgVersion96
         then "coalesce(nullif(current_setting('response.headers', true), ''), '[]')" :: Text -- nullif is used because of https://gist.github.com/steve-chavez/8d7033ea5655096903f3b52f8ed09a15

src/PostgREST/Types.hs

@@ -56,6 +56,12 @@ decodeContentType ct = case BS.takeWhile (/= BS.c2w ';') ct of
   "*/*"                               -> CTAny
   ct'                                 -> CTOther ct'
 
+-- | A SQL query that can be executed independently
+type SqlQuery = Text
+
+-- | A part of a SQL query that cannot be executed independently
+type SqlFragment = Text
+
 data PreferResolution = MergeDuplicates | IgnoreDuplicates deriving Eq
 instance Show PreferResolution where
   show MergeDuplicates  = "resolution=merge-duplicates"
@@ -135,10 +141,13 @@ specifiedProcArgs keys proc =
   in
   (\k -> fromMaybe (PgArg k "text" True) (find ((==) k . pgaName) args)) <$> S.toList keys
 
+procReturnsScalar :: ProcDescription -> Bool
+procReturnsScalar proc = case proc of
+  ProcDescription{pdReturnType = (Single (Scalar _))} -> True
+  _                                                   -> False
+
 type Schema = Text
 type TableName = Text
-type SqlQuery = Text
-type SqlFragment = Text
 
 data Table = Table {
   tableSchema      :: Schema
@@ -441,8 +450,6 @@ type JSPath = [JSPathExp]
 -- | jspath expression, e.g. .property, .property[0] or ."property-dash"
 data JSPathExp = JSPKey Text | JSPIdx Int deriving (Eq, Show)
 
-
-
 -- | Current database connection status data ConnectionStatus
 data ConnectionStatus
   = NotConnected