Implements Prefer count=none header

Uses Maybe for total parameter in contentRangeH

src/PostgREST/App.hs

@@ -5,7 +5,7 @@ import Control.Monad (join)
 import Control.Arrow ((***), second)
 import Control.Applicative
 
-import Data.Text hiding (map, find)
+import Data.Text (Text)
 import Data.Maybe (fromMaybe, mapMaybe, isJust, isNothing)
 import Text.Regex.TDFA ((=~))
 import Data.Ord (comparing)
@@ -64,9 +64,12 @@ app conf reqBody req =
       else do
         let qt = qualify table
             from = fromMaybe 0 $ rangeOffset <$> range
+            count = if hasPrefer "count=none" 
+                      then countNone
+                      else countRows qt
             query = B.Stmt "select " V.empty True <>
                 parentheticT (
-                  whereT qt qq $ countRows qt
+                  whereT qt qq count
                 ) <> commaq <> (
                 bodyForAccept contentType qt
                 . limitT range
@@ -76,7 +79,7 @@ app conf reqBody req =
               )
         row <- H.maybeEx query
         let (tableTotal, queryTotal, body) =
-              fromMaybe (0, 0, Just "" :: Maybe Text) row
+              fromMaybe (Just 0, 0, Just "" :: Maybe Text) row
             to = from+queryTotal-1
             contentRange = contentRangeH from to tableTotal
             status = rangeStatus from to tableTotal
@@ -129,7 +132,7 @@ app conf reqBody req =
 
     ([table], "POST") -> do
       let qt = qualify table
-          echoRequested = lookupHeader "Prefer" == Just "return=representation"
+          echoRequested = hasPrefer "return=representation"
           parsed :: Either String (V.Vector Text, V.Vector (V.Vector Value))
           parsed = if lookupHeader "Content-Type" == Just csvMT
             then do
@@ -215,8 +218,8 @@ app conf reqBody req =
         row <- H.maybeEx patch
         let (queryTotal, body) =
               fromMaybe (0 :: Int, Just "" :: Maybe Text) row
-            r = contentRangeH 0 (queryTotal-1) queryTotal
-            echoRequested = lookupHeader "Prefer" == Just "return=representation"
+            r = contentRangeH 0 (queryTotal-1) (Just queryTotal)
+            echoRequested = hasPrefer "return=representation"
             s = case () of _ | queryTotal == 0 -> status404
                              | echoRequested -> status200
                              | otherwise -> status204
@@ -244,6 +247,7 @@ app conf reqBody req =
     qualify       = QualifiedIdentifier schema
     hdrs          = requestHeaders req
     lookupHeader  = flip lookup hdrs
+    hasPrefer val = any (\(h,v) -> h == "Prefer" && v == val) hdrs
     accept        = lookupHeader hAccept
     schema        = requestedSchema (cs $ configV1Schema conf) accept
     authenticator = cs $ configDbUser conf
@@ -259,21 +263,24 @@ sqlError = undefined
 isSqlError :: t
 isSqlError = undefined
 
-rangeStatus :: Int -> Int -> Int -> Status
-rangeStatus from to total
+rangeStatus :: Int -> Int -> Maybe Int -> Status
+rangeStatus _ _ Nothing = status200
+rangeStatus from to (Just total)
   | from > total            = status416
   | (1 + to - from) < total = status206
   | otherwise               = status200
 
-contentRangeH :: Int -> Int -> Int -> Header
+contentRangeH :: Int -> Int -> Maybe Int -> Header
 contentRangeH from to total =
-  ("Content-Range",
-    if total == 0 || from > total
-    then "*/" <> cs (show total)
-    else cs (show from)
-      <> "-" <> cs (show to)
-      <> "/" <> cs (show total)
-  )
+    ("Content-Range", cs headerValue)
+    where
+      headerValue   = rangeString <> "/" <> totalString
+      rangeString
+        | totalNotZero && fromInRange = show from <> "-" <> cs (show to)
+        | otherwise = "*"
+      totalString   = fromMaybe "*" (show <$> total)
+      totalNotZero  = fromMaybe True ((/=) 0 <$> total)
+      fromInRange   = from <= to
 
 requestedSchema :: Text -> Maybe BS.ByteString -> Text
 requestedSchema v1schema accept =

src/PostgREST/PgQuery.hs

@@ -101,6 +101,9 @@ countT s =
 countRows :: QualifiedIdentifier  -> PStmt
 countRows t = B.Stmt ("select pg_catalog.count(1) from " <> fromQi t) empty True
 
+countNone :: PStmt
+countNone = B.Stmt "select null" empty True
+
 asCsvWithCount :: QualifiedIdentifier -> StatementT
 asCsvWithCount table = withCount . asCsv table
 

test/Feature/RangeSpec.hs

@@ -2,6 +2,7 @@ module Feature.RangeSpec where
 
 import Test.Hspec
 import Test.Hspec.Wai
+import Test.Hspec.Wai.JSON
 import Network.HTTP.Types
 import Network.Wai.Test (SResponse(simpleHeaders,simpleStatus))
 
@@ -17,7 +18,7 @@ spec = beforeAll (clearTable "items" >> createItems 15) . afterAll_ (clearTable
         it "returns whole range with status 200" $
           get "/items" `shouldRespondWith` 200
 
-      context "when I don't want the count" $
+      context "when I don't want the count" $ do
         it "returns range Content-Range with /*" $
           request methodGet "/menagerie"
                   [("Prefer", "count=none")] ""
@@ -27,6 +28,15 @@ spec = beforeAll (clearTable "items" >> createItems 15) . afterAll_ (clearTable
             , matchHeaders = ["Content-Range" <:> "*/*"]
             }
 
+        it "returns range Content-Range with range/*" $
+          request methodGet "/items?order=id"
+                  [("Prefer", "count=none")] ""
+            `shouldRespondWith` ResponseMatcher {
+              matchBody    = Just [json| [{"id":1},{"id":2},{"id":3},{"id":4},{"id":5},{"id":6},{"id":7},{"id":8},{"id":9},{"id":10},{"id":11},{"id":12},{"id":13},{"id":14},{"id":15}] |]
+            , matchStatus  = 200
+            , matchHeaders = ["Content-Range" <:> "0-14/*"]
+            }
+
     context "with range headers" $ do
 
       context "of acceptable range" $ do