From b2ce322d56687fc67000736876571d8123585502 Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Mon, 1 Sep 2025 19:53:07 -0500 Subject: [PATCH] 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 | 48 ++++++------------------ src/PostgREST/Query/Statements.hs | 61 +++++++++++++++++++++---------- 2 files changed, 52 insertions(+), 57 deletions(-) diff --git a/src/PostgREST/Query.hs b/src/PostgREST/Query.hs index 76d54512d..766b486a7 100644 --- a/src/PostgREST/Query.hs +++ b/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 diff --git a/src/PostgREST/Query/Statements.hs b/src/PostgREST/Query/Statements.hs index 3cb633d2e..e90611aea 100644 --- a/src/PostgREST/Query/Statements.hs +++ b/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