perf: optimize count=exact when there's no limits, offsets or db-max-rows

Author: laurenceisla

src/PostgREST/Query.hs

@@ -43,16 +43,16 @@ data MainQuery = MainQuery
 
 mainQuery :: ActionPlan -> AppConfig -> ApiRequest -> AuthResult -> Maybe QualifiedIdentifier -> MainQuery
 mainQuery (NoDb _) _ _ _ _ = MainQuery mempty Nothing mempty (mempty, mempty, mempty) mempty
-mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}} authRes preReq =
+mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iTopLevelRange=range, iPreferences=Preferences{..}} authRes preReq =
   let genQ = MainQuery (PreQuery.txVarQuery plan conf authRes apiReq) (PreQuery.preReqQuery <$> preReq) in
   case plan of
     DbCrud _ WrappedReadPlan{..} ->
       let countQuery = QueryBuilder.readPlanToCountQuery wrReadPlan in
-      genQ (Statements.mainRead wrReadPlan countQuery preferCount configDbMaxRows pMedia wrHandler) (mempty, mempty, mempty)
+      genQ (Statements.mainRead wrReadPlan countQuery preferCount configDbMaxRows range pMedia wrHandler) (mempty, mempty, mempty)
       (if shouldExplainCount preferCount then Just (Statements.postExplain countQuery) else Nothing)
     DbCrud _ MutateReadPlan{..} ->
       genQ (Statements.mainWrite mrReadPlan mrMutatePlan pMedia mrHandler preferRepresentation preferResolution) (mempty, mempty, mempty) mempty
     DbCrud _ CallReadPlan{..} ->
-      genQ (Statements.mainCall crProc crCallPlan crReadPlan preferCount pMedia crHandler) (mempty, mempty, mempty) mempty
+      genQ (Statements.mainCall crProc crCallPlan crReadPlan preferCount configDbMaxRows range pMedia crHandler) (mempty, mempty, mempty) mempty
     MayUseDb InspectPlan{ipSchema=tSchema} ->
       genQ mempty (SqlFragment.accessibleTables tSchema, SqlFragment.accessibleFuncs tSchema, SqlFragment.schemaDescription tSchema) mempty

src/PostgREST/Query/SqlFragment.hs

@@ -487,15 +487,15 @@ pgFmtGroup _  CoercibleSelectField{csAggFunction=Just _} = Nothing
 pgFmtGroup _  CoercibleSelectField{csAlias=Just alias, csAggFunction=Nothing} = Just $ pgFmtIdent alias
 pgFmtGroup qi CoercibleSelectField{csField=fld, csAlias=Nothing, csAggFunction=Nothing} = Just $ pgFmtField qi fld
 
-countF :: SQL.Snippet -> Bool -> (SQL.Snippet, SQL.Snippet)
-countF countQuery shouldCount =
-  if shouldCount
-    then (
-        ", pgrst_source_count AS (" <> countQuery <> ")"
-      , "(SELECT pg_catalog.count(*) FROM pgrst_source_count)" )
-    else (
-        mempty
-      , "null::bigint")
+countF :: SQL.Snippet -> SQL.Snippet -> Bool -> Maybe Integer -> NonnegRange -> (SQL.Snippet, SQL.Snippet)
+countF countQuery pageCountSelect shouldCount maxRows range
+  | shouldCount = if isJust maxRows || range /= allRange
+      then ( ", pgrst_source_count AS (" <> countQuery <> ")"
+           , "(SELECT pg_catalog.count(*) FROM pgrst_source_count)" )
+      -- When there are no db-max-rows and limits/offsets, the total count will be the same as the page count,
+      -- so we use the same page count here to avoid doing a separate aggregated count.
+      else ( mempty, pageCountSelect )
+  | otherwise = ( mempty, "null::bigint" )
 
 pageCountSelectF :: Maybe Routine -> SQL.Snippet
 pageCountSelectF rout =

src/PostgREST/Query/Statements.hs

@@ -20,6 +20,7 @@ import PostgREST.Plan.MutatePlan        as MTPlan
 import PostgREST.Plan.ReadPlan
 import PostgREST.Query.QueryBuilder
 import PostgREST.Query.SqlFragment
+import PostgREST.RangeQuery             (NonnegRange)
 import PostgREST.SchemaCache.Routine    (MediaHandler (..), Routine)
 
 import Protolude
@@ -63,50 +64,52 @@ mainWrite rPlan mtplan mt handler rep resolution = mtSnippet mt snippet
     _ -> (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
+            NonnegRange -> MediaType -> MediaHandler -> SQL.Snippet
+mainRead rPlan countQuery pCount maxRows range mt handler = mtSnippet mt snippet
  where
   snippet =
     "WITH " <> sourceCTE <> " AS ( " <> selectQuery <> " ) " <>
     countCTEF <> " " <>
     "SELECT " <>
       countResultF <> " AS total_result_set, " <>
-      pageCountSelectF Nothing <> " AS page_total, " <>
+      pageCountSelect <> " AS page_total, " <>
       handlerF Nothing handler <> " AS body, " <>
       responseHeadersF <> " AS response_headers, " <>
       responseStatusF <> " AS response_status, " <>
       "''" <> " AS response_inserted " <>
     "FROM ( SELECT * FROM " <> sourceCTE <> " ) _postgrest_t"
 
-  (countCTEF, countResultF) = countF countQ $ shouldCount pCount
+  (countCTEF, countResultF) = countF countQ pageCountSelect (shouldCount pCount) maxRows range
   selectQuery = readPlanToQuery rPlan
+  pageCountSelect = pageCountSelectF Nothing
   countQ =
     if pCount == Just EstimatedCount then
       -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
       limitedQuery countQuery ((+ 1) <$> maxRows)
     else
       countQuery
 
-mainCall :: Routine -> CallPlan -> ReadPlanTree -> Maybe PreferCount ->
-            MediaType -> MediaHandler -> SQL.Snippet
-mainCall rout cPlan rPlan pCount mt handler = mtSnippet mt snippet
+mainCall :: Routine -> CallPlan -> ReadPlanTree -> Maybe PreferCount -> Maybe Integer ->
+            NonnegRange-> MediaType -> MediaHandler -> SQL.Snippet
+mainCall rout cPlan rPlan pCount maxRows range mt handler = mtSnippet mt snippet
   where
     snippet =
       "WITH " <> sourceCTE <> " AS (" <> callProcQuery <> ") " <>
       countCTEF <>
       "SELECT " <>
         countResultF <> " AS total_result_set, " <>
-        pageCountSelectF (Just rout) <> " AS page_total, " <>
+        pageCountSelect <> " AS page_total, " <>
         handlerF (Just rout) handler <> " AS body, " <>
         responseHeadersF <> " AS response_headers, " <>
         responseStatusF <> " AS response_status, " <>
         "''" <> " AS response_inserted " <>
       "FROM (" <> selectQuery <> ") _postgrest_t"
 
-    (countCTEF, countResultF) = countF countQuery $ shouldCount pCount
+    (countCTEF, countResultF) = countF countQuery pageCountSelect (shouldCount pCount) maxRows range
     selectQuery = readPlanToQuery rPlan
     callProcQuery = callPlanToQuery cPlan
     countQuery = readPlanToCountQuery rPlan
+    pageCountSelect = pageCountSelectF (Just rout)
 
 -- This occurs after the main query runs, that's why it's prefixed with "post"
 postExplain :: SQL.Snippet -> SQL.Snippet