refactor: move snippet logic inside Statements.hs

Some snippet logic was unnecessarily at the Query.hs level.

This is better as we can see what inputs we need to form the main
queries.

src/PostgREST/Query.hs

@@ -45,8 +45,7 @@ import PostgREST.ApiRequest.Preferences  (PreferCount (..),
                                           PreferHandling (..),
                                           PreferMaxAffected (..),
                                           PreferTransaction (..),
-                                          Preferences (..),
-                                          shouldCount)
+                                          Preferences (..))
 import PostgREST.ApiRequest.Types        (Mutation (..))
 import PostgREST.Auth.Types              (AuthResult (..))
 import PostgREST.Config                  (AppConfig (..),
@@ -59,7 +58,6 @@ import PostgREST.Plan                    (ActionPlan (..),
                                           DbActionPlan (..),
                                           InfoPlan (..),
                                           InspectPlan (..))
-import PostgREST.Plan.MutatePlan         (MutatePlan (..))
 import PostgREST.SchemaCache             (SchemaCache (..))
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
 import PostgREST.SchemaCache.Routine     (Routine (..), RoutineMap)
@@ -139,18 +137,9 @@ actionQuery (DbCrud plan@WrappedReadPlan{..}) conf@AppConfig{..} apiReq@ApiReque
   (mainActionQuery, mainSQLQuery)
   where
     countQuery = QueryBuilder.readPlanToCountQuery wrReadPlan
-    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized (Statements.prepareRead
-      (QueryBuilder.readPlanToQuery wrReadPlan)
-      (if preferCount == Just EstimatedCount then
-         -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
-         QueryBuilder.limitedQuery countQuery ((+ 1) <$> configDbMaxRows)
-       else
-         countQuery
-      )
-      (shouldCount preferCount)
-      wrMedia
-      wrHandler
-      ) decodeIt configDbPreparedStatements
+    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized
+      (Statements.mainRead wrReadPlan countQuery preferCount configDbMaxRows wrMedia wrHandler)
+      decodeIt configDbPreparedStatements
     mainActionQuery = do
       resultSet <- lift $ SQL.statement mempty result
       failNotSingular wrMedia resultSet
@@ -165,19 +154,9 @@ actionQuery (DbCrud plan@WrappedReadPlan{..}) conf@AppConfig{..} apiReq@ApiReque
 actionQuery (DbCrud plan@MutateReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} _ =
   (mainActionQuery, mainSQLQuery)
   where
-    (isPut, isInsert, pkCols) = case mrMutatePlan of
-      Insert{where_,insPkCols} -> ((not . null) where_, True, insPkCols)
-      _ -> (False,False, mempty);
-    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized (Statements.prepareWrite
-      (QueryBuilder.readPlanToQuery mrReadPlan)
-      (QueryBuilder.mutatePlanToQuery mrMutatePlan)
-      isInsert
-      isPut
-      mrMedia
-      mrHandler
-      preferRepresentation
-      preferResolution
-      pkCols) decodeIt configDbPreparedStatements
+    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized
+      (Statements.mainWrite mrReadPlan mrMutatePlan mrMedia mrHandler preferRepresentation preferResolution)
+      decodeIt configDbPreparedStatements
     failMutation resultSet = case mrMutation of
       MutationCreate -> do
         failNotSingular mrMedia resultSet
@@ -203,14 +182,9 @@ actionQuery (DbCrud plan@MutateReadPlan{..}) conf@AppConfig{..} apiReq@ApiReques
 actionQuery (DbCall plan@CallReadPlan{..}) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} _ =
   (mainActionQuery, mainSQLQuery)
   where
-    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized (Statements.prepareCall
-      crProc
-      (QueryBuilder.callPlanToQuery crCallPlan)
-      (QueryBuilder.readPlanToQuery crReadPlan)
-      (QueryBuilder.readPlanToCountQuery crReadPlan)
-      (shouldCount preferCount)
-      crMedia
-      crHandler) decodeIt configDbPreparedStatements
+    result@(SQL.Statement mainSQLQuery _ _ _) = SQL.dynamicallyParameterized
+      (Statements.mainCall crProc crCallPlan crReadPlan preferCount crMedia crHandler)
+      decodeIt configDbPreparedStatements
 
     mainActionQuery = do
       resultSet <- lift $ SQL.statement mempty result
@@ -275,7 +249,7 @@ resultSetWTotal AppConfig{..} ApiRequest{iPreferences=Preferences{..}} rs@RSStan
   where
     explain =
       lift . SQL.statement mempty $
-        SQL.dynamicallyParameterized (Statements.preparePlanRows countQuery)
+        SQL.dynamicallyParameterized (Statements.postExplain countQuery)
         decodeIt
         configDbPreparedStatements
 

src/PostgREST/Query/Statements.hs

@@ -1,12 +1,13 @@
+{-# LANGUAGE NamedFieldPuns #-}
 {-|
 Module      : PostgREST.Query.Statements
-Description : PostgREST SQL statements.
+Description : PostgREST main queries
 -}
 module PostgREST.Query.Statements
-  ( prepareWrite
-  , prepareRead
-  , prepareCall
-  , preparePlanRows
+  ( mainWrite
+  , mainRead
+  , mainCall
+  , postExplain
   ) where
 
 import qualified Hasql.DynamicStatements.Snippet as SQL
@@ -14,15 +15,19 @@ import qualified Hasql.DynamicStatements.Snippet as SQL
 import PostgREST.ApiRequest.Preferences
 import PostgREST.MediaType              (MTVndPlanFormat (..),
                                          MediaType (..))
+import PostgREST.Plan.CallPlan
+import PostgREST.Plan.MutatePlan        as MTPlan
+import PostgREST.Plan.ReadPlan
+import PostgREST.Query.QueryBuilder
 import PostgREST.Query.SqlFragment
 import PostgREST.SchemaCache.Routine    (MediaHandler (..), Routine,
                                          funcReturnsSingle)
 
 import Protolude
 
-prepareWrite :: SQL.Snippet -> SQL.Snippet -> Bool -> Bool -> MediaType -> MediaHandler ->
-                Maybe PreferRepresentation -> Maybe PreferResolution -> [Text] -> SQL.Snippet
-prepareWrite selectQuery mutateQuery isInsert isPut mt handler rep resolution pKeys = mtSnippet mt snippet
+mainWrite :: ReadPlanTree -> MutatePlan -> MediaType -> MediaHandler ->
+             Maybe PreferRepresentation -> Maybe PreferResolution -> SQL.Snippet
+mainWrite rPlan mtplan mt handler rep resolution = mtSnippet mt snippet
  where
   checkUpsert snip = if isInsert && (isPut || resolution == Just MergeDuplicates) then snip else "''"
   pgrstInsertedF = checkUpsert "nullif(current_setting('pgrst.inserted', true),'')::int"
@@ -42,7 +47,7 @@ prepareWrite selectQuery mutateQuery isInsert isPut mt handler rep resolution pK
     if isInsert && rep == Just HeadersOnly
       then
         "CASE WHEN pg_catalog.count(_postgrest_t) = 1 " <>
-          "THEN coalesce(" <> locationF pKeys <> ", " <> noLocationF <> ") " <>
+          "THEN coalesce(" <> locationF pkCols <> ", " <> noLocationF <> ") " <>
           "ELSE " <> noLocationF <> " " <>
         "END"
       else noLocationF
@@ -52,8 +57,15 @@ prepareWrite selectQuery mutateQuery isInsert isPut mt handler rep resolution pK
     | handler == NoAgg = "SELECT * FROM " <> sourceCTE
     | otherwise        = selectQuery
 
-prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> MediaHandler -> SQL.Snippet
-prepareRead selectQuery countQuery countTotal mt handler = mtSnippet mt snippet
+  selectQuery = readPlanToQuery rPlan
+  mutateQuery = mutatePlanToQuery mtplan
+  (isPut, isInsert, pkCols) = case mtplan of
+    MTPlan.Insert{MTPlan.where_,insPkCols} -> ((not . null) where_, True, insPkCols)
+    _ -> (False,False, mempty);
+
+mainRead :: ReadPlanTree -> SQL.Snippet -> Maybe PreferCount -> Maybe Integer ->
+            MediaType -> MediaHandler -> SQL.Snippet
+mainRead rPlan countQuery pCount maxRows mt handler = mtSnippet mt snippet
  where
   snippet =
     "WITH " <> sourceCTE <> " AS ( " <> selectQuery <> " ) " <>
@@ -67,12 +79,18 @@ prepareRead selectQuery countQuery countTotal mt handler = mtSnippet mt snippet
       "''" <> " AS response_inserted " <>
     "FROM ( SELECT * FROM " <> sourceCTE <> " ) _postgrest_t"
 
-  (countCTEF, countResultF) = countF countQuery countTotal
+  (countCTEF, countResultF) = countF countQ $ shouldCount pCount
+  selectQuery = readPlanToQuery rPlan
+  countQ =
+    if pCount == Just EstimatedCount then
+      -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
+      limitedQuery countQuery ((+ 1) <$> maxRows)
+    else
+      countQuery
 
-
-prepareCall :: Routine -> SQL.Snippet -> SQL.Snippet -> SQL.Snippet -> Bool ->
-               MediaType -> MediaHandler -> SQL.Snippet
-prepareCall rout callProcQuery selectQuery countQuery countTotal mt handler = mtSnippet mt snippet
+mainCall :: Routine -> CallPlan -> ReadPlanTree -> Maybe PreferCount ->
+            MediaType -> MediaHandler -> SQL.Snippet
+mainCall rout cPlan rPlan pCount mt handler = mtSnippet mt snippet
   where
     snippet =
       "WITH " <> sourceCTE <> " AS (" <> callProcQuery <> ") " <>
@@ -88,11 +106,14 @@ prepareCall rout callProcQuery selectQuery countQuery countTotal mt handler = mt
         "''" <> " AS response_inserted " <>
       "FROM (" <> selectQuery <> ") _postgrest_t"
 
-    (countCTEF, countResultF) = countF countQuery countTotal
+    (countCTEF, countResultF) = countF countQuery $ shouldCount pCount
+    selectQuery = readPlanToQuery rPlan
+    callProcQuery = callPlanToQuery cPlan
+    countQuery = readPlanToCountQuery rPlan
 
-
-preparePlanRows :: SQL.Snippet -> SQL.Snippet
-preparePlanRows = explainF PlanJSON mempty
+-- This occurs after the main query runs, that's why it's prefixed with "post"
+postExplain :: SQL.Snippet -> SQL.Snippet
+postExplain = explainF PlanJSON mempty
 
 mtSnippet :: MediaType -> SQL.Snippet -> SQL.Snippet
 mtSnippet mediaType snippet = case mediaType of