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

src/PostgREST/QueryBuilder.hs

@@ -14,6 +14,7 @@ module PostgREST.QueryBuilder (
     requestToQuery
   , requestToCountQuery
   , requestToCallProcQuery
+  , limitedQuery
   , setLocalQuery
   , setLocalSearchPathQuery
   ) where
@@ -32,17 +33,6 @@ import PostgREST.Types
 import Protolude                      hiding (cast, intercalate,
                                        replace)
 
-requestToCountQuery :: Schema -> DbRequest -> SqlQuery
-requestToCountQuery _ (DbMutate _) = witness
-requestToCountQuery schema (DbRead (Node (Select{where_=logicForest}, (mainTbl, _, _, _, _)) _)) =
- unwords [
-   "SELECT pg_catalog.count(*)",
-   "FROM ", fromQi qi,
-    ("WHERE " <> intercalate " AND " (map (pgFmtLogicTree qi) logicForest)) `emptyOnFalse` null logicForest
-   ]
- where
-   qi = removeSourceCTESchema schema mainTbl
-
 requestToQuery :: Schema -> Bool -> DbRequest -> SqlQuery
 requestToQuery schema isParent (DbRead (Node (Select colSelects tbl tblAlias implJoins logicForest joinConditions_ ordts range, _) forest)) =
   unwords [
@@ -177,6 +167,25 @@ requestToCallProcQuery qi pgArgs returnsScalar preferParams =
     callIt :: SqlFragment
     callIt = fromQi qi <> "(" <> args <> ")"
 
+
+-- | SQL query meant for COUNTing the root node of the DbRead Tree.
+-- It only takes WHERE into account and doesn't include LIMIT/OFFSET because it would reduce the COUNT.
+-- SELECT 1 is done instead of SELECT * to prevent doing expensive operations(like functions based on the columns)
+-- inside the FROM target.
+requestToCountQuery :: Schema -> DbRequest -> SqlQuery
+requestToCountQuery _ (DbMutate _) = witness
+requestToCountQuery schema (DbRead (Node (Select{where_=logicForest}, (mainTbl, _, _, _, _)) _)) =
+ unwords [
+   "SELECT 1",
+   "FROM " <> fromQi qi,
+   ("WHERE " <> intercalate " AND " (map (pgFmtLogicTree qi) logicForest)) `emptyOnFalse` null logicForest
+   ]
+ where
+   qi = removeSourceCTESchema schema mainTbl
+
+limitedQuery :: SqlQuery -> Maybe Integer -> SqlQuery
+limitedQuery query maxRows = query <> maybe mempty (\x -> " LIMIT " <> show x) maxRows
+
 setLocalQuery :: Text -> (Text, Text) -> SqlQuery
 setLocalQuery prefix (k, v) =
   "SET LOCAL " <> pgFmtIdent (prefix <> k) <> " = " <> pgFmtLit v <> ";"