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

Author: laurenceisla

CHANGELOG.md

@@ -7,6 +7,8 @@ All notable changes to this project will be documented in this file. From versio
 ### Added
 
 - Log error when `db-schemas` config contains schema `pg_catalog` or `information_schema` by @taimoorzaeem in #4359
+- Optimize requests with `Prefer: count=exact` that do not use ranges or `db-max-rows` by @laurenceisla in #3957
+  + The page total is the same as the resource total in these cases, so now it only counts once to build the `Content-Range`.
 
 ## [14.2] - 2025-12-18
 

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

test/spec/Feature/Query/PlanSpec.hs

@@ -455,6 +455,80 @@ spec actualPgVersion = do
         nextValSnip `shouldBe`
           Just [aesonQQ| ["jsonb_agg((jsonb_build_object('id', nextval('\"Surr_Gen_Default_Upsert_id_seq\"'::regclass)) || elem.value))"] |]
 
+  describe "count preference plan costs" $ do
+    context "tables with count=exact" $ do
+      it "shows only 1 count aggregate when no limits/max-rows are set" $ do
+        _ <- request methodPost "/tiobe_pls"
+              [("Prefer","resolution=merge-duplicates"), ("Accept","application/vnd.pgrst.plan+json")]
+              (getInsertDataForTiobePlsTable 1000)
+
+        r <- request methodGet "/tiobe_pls"
+          (("Prefer", "count=exact") : acceptHdrs "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze;") ""
+
+        let resBody = simpleBody r
+            resHeaders = simpleHeaders r
+            totalCost = planCost r
+            aggregateQty = subtract 1 $ length $ T.splitOn "Aggregate" (decodeUtf8 $ LBS.toStrict resBody)
+
+        liftIO $ do
+          resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze; charset=utf-8")
+          totalCost `shouldSatisfy` (< 33.0)
+          aggregateQty `shouldBe` 1
+
+      it "shows relevant count aggregates when limits are set" $ do
+        _ <- request methodPost "/tiobe_pls"
+              [("Prefer","resolution=merge-duplicates"), ("Accept","application/vnd.pgrst.plan+json")]
+              (getInsertDataForTiobePlsTable 1000)
+
+        r <- request methodGet "/tiobe_pls?limit=1000"
+          (("Prefer", "count=exact") : acceptHdrs "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze;") ""
+
+        let resBody = simpleBody r
+            resHeaders = simpleHeaders r
+            totalCost = planCost r
+            aggregateQty = subtract 1 $ length $ T.splitOn "Aggregate" (decodeUtf8 $ LBS.toStrict resBody)
+
+        liftIO $ do
+          resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze; charset=utf-8")
+          totalCost `shouldSatisfy` (> 49.0)
+          aggregateQty `shouldSatisfy` (> 1)
+
+    context "functions with count=exact" $ do
+      it "shows only 1 count aggregate when no limits/max-rows are set" $ do
+        _ <- request methodPost "/tiobe_pls"
+              [("Prefer","resolution=merge-duplicates"), ("Accept","application/vnd.pgrst.plan+json"), ("Prefer", "return=representation")]
+              (getInsertDataForTiobePlsTable 1000)
+
+        r <- request methodGet "/rpc/get_tiobe_pls"
+          (("Prefer", "count=exact") : acceptHdrs "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze;") ""
+
+        let resBody = simpleBody r
+            resHeaders = simpleHeaders r
+            totalCost = planCost r
+            aggregateQty = subtract 1 $ length $ T.splitOn "Aggregate" (decodeUtf8 $ LBS.toStrict resBody)
+
+        liftIO $ do
+          resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze; charset=utf-8")
+          totalCost `shouldSatisfy` (< 38.0)
+          aggregateQty `shouldBe` 1
+
+      it "shows relevant count aggregates when limits are set" $ do
+        _ <- request methodPost "/tiobe_pls"
+              [("Prefer","resolution=merge-duplicates"), ("Accept","application/vnd.pgrst.plan+json")]
+              (getInsertDataForTiobePlsTable 1000)
+
+        r <- request methodGet "/rpc/get_tiobe_pls?limit=1000"
+          (("Prefer", "count=exact") : acceptHdrs "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze;") ""
+
+        let resBody = simpleBody r
+            resHeaders = simpleHeaders r
+            totalCost = planCost r
+            aggregateQty = subtract 1 $ length $ T.splitOn "Aggregate" (decodeUtf8 $ LBS.toStrict resBody)
+
+        liftIO $ do
+          resHeaders `shouldSatisfy` elem ("Content-Type", "application/vnd.pgrst.plan+json; for=\"application/json\"; options=analyze; charset=utf-8")
+          totalCost `shouldSatisfy` (> 67.0)
+          aggregateQty `shouldSatisfy` (> 1)
 
 disabledSpec :: SpecWith ((), Application)
 disabledSpec =

test/spec/fixtures/schema.sql

@@ -3851,3 +3851,7 @@ $$ language sql;
 
 create function do_nothing() returns void as $_$
 $_$ language sql;
+
+create or replace function test.get_tiobe_pls() returns setof test.tiobe_pls as $$
+  select * from test.tiobe_pls;
+$$ language sql;