Add support for Prefer: count=planned/estimated on GET /table (#1386)

src/PostgREST/App.hs

@@ -10,6 +10,7 @@ Some of its functionality includes:
 - Content Negotiation
 -}
 {-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE MultiWayIf          #-}
 {-# LANGUAGE NamedFieldPuns      #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
@@ -55,12 +56,14 @@ import PostgREST.Error            (PgError (..), SimpleError (..),
 import PostgREST.Middleware
 import PostgREST.OpenAPI
 import PostgREST.Parsers          (pRequestColumns)
-import PostgREST.QueryBuilder     (requestToCallProcQuery,
+import PostgREST.QueryBuilder     (limitedQuery,
+                                   requestToCallProcQuery,
                                    requestToCountQuery,
                                    requestToQuery)
 import PostgREST.RangeQuery       (allRange, contentRangeH,
                                    rangeStatusHeader)
 import PostgREST.Statements       (callProcStatement,
+                                   createExplainStatement,
                                    createReadStatement,
                                    createWriteStatement)
 import PostgREST.Types
@@ -128,11 +131,21 @@ app dbStructure proc cols conf apiRequest =
           case partsField of
             Left errorResponse -> return errorResponse
             Right ((q, cq), bField) -> do
-              let stm = createReadStatement q cq (contentType == CTSingularJSON) shouldCount
-                                            (contentType == CTTextCSV) bField
+              let cQuery = if estimatedCount
+                             then limitedQuery cq ((+ 1) <$> maxRows) -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
+                             else cq
+                  stm = createReadStatement q cQuery (contentType == CTSingularJSON) shouldCount
+                        (contentType == CTTextCSV) bField
+                  explStm = createExplainStatement cq
               row <- H.statement () stm
               let (tableTotal, queryTotal, _ , body) = row
-                  (status, contentRange) = rangeStatusHeader topLevelRange queryTotal tableTotal
+              total <- if | plannedCount   -> H.statement () explStm
+                          | estimatedCount -> if tableTotal > (fromIntegral <$> maxRows)
+                                                then do estTotal <- H.statement () explStm
+                                                        pure $ if estTotal > tableTotal then estTotal else tableTotal
+                                                else pure tableTotal
+                          | otherwise      -> pure tableTotal
+              let (status, contentRange) = rangeStatusHeader topLevelRange queryTotal total
               return $
                 if contentType == CTSingularJSON && queryTotal /= 1
                   then errorResponseFor . singularityError $ queryTotal
@@ -317,10 +330,14 @@ app dbStructure proc cols conf apiRequest =
 
     where
       notFound = responseLBS status404 [] ""
-      shouldCount = iPreferCount apiRequest
-      topLevelRange = iTopLevelRange apiRequest
       schema = toS $ configSchema conf
-      readReq = readRequest (configMaxRows conf) (dbRelations dbStructure) proc apiRequest
+      maxRows = configMaxRows conf
+      exactCount = iPreferCount apiRequest == Just ExactCount
+      estimatedCount = iPreferCount apiRequest == Just EstimatedCount
+      plannedCount = iPreferCount apiRequest == Just PlannedCount
+      shouldCount = exactCount || estimatedCount
+      topLevelRange = iTopLevelRange apiRequest
+      readReq = readRequest maxRows (dbRelations dbStructure) proc apiRequest
       fldNames = fieldNames <$> readReq
       readDbRequest = DbRead <$> readReq
       selectQuery = requestToQuery schema False <$> readDbRequest