From d37e14c4db32ff1c4b83df894be5549fa2dd1a90 Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavezast@gmail.com>
Date: Tue, 4 Oct 2022 18:44:34 -0500
Subject: [PATCH] refactor: readTotal into ResultSet

---
 src/PostgREST/App.hs      |  4 ++--
 src/PostgREST/Query.hs    | 29 ++++++++++++++++-------------
 src/PostgREST/Response.hs |  9 ++++-----
 3 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
index d2f0290bd..baaaa8b33 100644
--- a/src/PostgREST/App.hs
+++ b/src/PostgREST/App.hs
@@ -225,9 +225,9 @@ handleRead :: Bool -> QualifiedIdentifier -> RequestContext -> DbHandler Wai.Res
 handleRead headersOnly identifier context@RequestContext{..} = do
   req <- liftEither $ readPlan identifier context
 
-  (resultSet, total) <- Query.readQuery req ctxConfig ctxApiRequest
+  resultSet <- Query.readQuery req ctxConfig ctxApiRequest
 
-  pure $ Response.readResponse headersOnly identifier ctxApiRequest total resultSet
+  pure $ Response.readResponse headersOnly identifier ctxApiRequest resultSet
 
 handleCreate :: QualifiedIdentifier -> RequestContext -> DbHandler Wai.Response
 handleCreate identifier context@RequestContext{..} = do
diff --git a/src/PostgREST/Query.hs b/src/PostgREST/Query.hs
index 3d27e40da..6efb999e1 100644
--- a/src/PostgREST/Query.hs
+++ b/src/PostgREST/Query.hs
@@ -67,7 +67,7 @@ import Protolude hiding (Handler)
 
 type DbHandler = ExceptT Error SQL.Transaction
 
-readQuery :: ReadPlanTree -> AppConfig -> ApiRequest -> DbHandler (ResultSet, Maybe Int64)
+readQuery :: ReadPlanTree -> AppConfig -> ApiRequest -> DbHandler ResultSet
 readQuery req conf@AppConfig{..} apiReq@ApiRequest{..} = do
   let countQuery = QueryBuilder.readPlanToCountQuery req
   resultSet <-
@@ -85,22 +85,25 @@ readQuery req conf@AppConfig{..} apiReq@ApiRequest{..} = do
         iBinaryField
         configDbPreparedStatements
   failNotSingular iAcceptMediaType resultSet
-  total <- readTotal conf apiReq resultSet countQuery
-  pure (resultSet, total)
+  resultSetWTotal conf apiReq resultSet countQuery
 
-readTotal :: AppConfig -> ApiRequest -> ResultSet -> SQL.Snippet -> DbHandler (Maybe Int64)
-readTotal _ _ RSPlan{} _ = pure Nothing
-readTotal AppConfig{..} ApiRequest{..} RSStandard{rsTableTotal=tableTotal} countQuery =
+resultSetWTotal :: AppConfig -> ApiRequest -> ResultSet -> SQL.Snippet -> DbHandler ResultSet
+resultSetWTotal _ _ rs@RSPlan{} _ = return rs
+resultSetWTotal AppConfig{..} ApiRequest{..} rs@RSStandard{rsTableTotal=tableTotal} countQuery =
   case iPreferCount of
-    Just PlannedCount ->
-      explain
+    Just PlannedCount -> do
+      total <- explain
+      return rs{rsTableTotal=total}
     Just EstimatedCount ->
-      if tableTotal > (fromIntegral <$> configDbMaxRows) then
-        max tableTotal <$> explain
+      if tableTotal > (fromIntegral <$> configDbMaxRows) then do
+        total <- max tableTotal <$> explain
+        return rs{rsTableTotal=total}
       else
-        return tableTotal
-    _ ->
-      return tableTotal
+        return rs
+    Just ExactCount ->
+      return rs
+    Nothing ->
+      return rs
   where
     explain =
       lift . SQL.statement mempty . Statements.preparePlanRows countQuery $
diff --git a/src/PostgREST/Response.hs b/src/PostgREST/Response.hs
index aea0d23c8..6cbabfc5f 100644
--- a/src/PostgREST/Response.hs
+++ b/src/PostgREST/Response.hs
@@ -56,12 +56,11 @@ import qualified PostgREST.Request.Types    as ApiRequestTypes
 import Protolude      hiding (Handler, toS)
 import Protolude.Conv (toS)
 
-
-readResponse :: Bool -> QualifiedIdentifier -> ApiRequest -> Maybe Int64 -> ResultSet -> Wai.Response
-readResponse headersOnly identifier ctxApiRequest@ApiRequest{..} total resultSet = case resultSet of
+readResponse :: Bool -> QualifiedIdentifier -> ApiRequest -> ResultSet -> Wai.Response
+readResponse headersOnly identifier ctxApiRequest@ApiRequest{..} resultSet = case resultSet of
   RSStandard{..} -> do
     let
-      (status, contentRange) = RangeQuery.rangeStatusHeader iTopLevelRange rsQueryTotal total
+      (status, contentRange) = RangeQuery.rangeStatusHeader iTopLevelRange rsQueryTotal rsTableTotal
       response = gucResponse rsGucStatus rsGucHeaders
       headers =
         [ contentRange
@@ -74,7 +73,7 @@ readResponse headersOnly identifier ctxApiRequest@ApiRequest{..} total resultSet
         ++ contentTypeHeaders ctxApiRequest
       rsOrErrBody = if status == HTTP.status416
         then Error.errorPayload $ Error.ApiRequestError $ ApiRequestTypes.InvalidRange
-          $ ApiRequestTypes.OutOfBounds (show $ RangeQuery.rangeOffset iTopLevelRange) (maybe "0" show total)
+          $ ApiRequestTypes.OutOfBounds (show $ RangeQuery.rangeOffset iTopLevelRange) (maybe "0" show rsTableTotal)
         else LBS.fromStrict rsBody
 
     response status headers $ if headersOnly then mempty else rsOrErrBody