Add --max-rows option (no automated tests yet)

Fixes #288

src/PostgREST/ApiRequest.hs

@@ -20,7 +20,7 @@ import           PostgREST.RangeQuery    (NonnegRange, rangeRequested)
 import           PostgREST.Types         (QualifiedIdentifier (..),
                                           Schema, Payload(..),
                                           UniformObjects(..))
-import           Data.Ranged.Ranges      (singletonRange)                                          
+import           Data.Ranged.Ranges      (singletonRange)
 
 type RequestBody = BL.ByteString
 
@@ -51,7 +51,7 @@ data ApiRequest = ApiRequest {
   -- | Set to Nothing for unknown HTTP verbs
     iAction :: Action
   -- | Set to Nothing for malformed range
-  , iRange  :: Maybe NonnegRange
+  , iRange  :: NonnegRange
   -- | Set to Nothing for strangely nested urls
   , iTarget :: Target
   -- | The content type the client most desires (or JSON if undecided)
@@ -115,7 +115,7 @@ userApiRequest schema req reqBody =
 
   ApiRequest {
     iAction = action
-  , iRange  = if singular then Just (singletonRange 0) else rangeRequested hdrs
+  , iRange  = if singular then singletonRange 0 else rangeRequested hdrs
   , iTarget = target
   , iAccepts = pickContentType $ lookupHeader "accept"
   , iPayload = relevantPayload

src/PostgREST/App.hs

@@ -73,11 +73,11 @@ app dbStructure conf reqBody req =
       case selectQuery of
         Left e -> return $ responseLBS status400 [jsonH] $ cs e
         Right q -> do
-          let range = iRange apiRequest
+          let range = restrictRange (configMaxRows conf) $ iRange apiRequest
               singular = iPreferSingular apiRequest
               stm = createReadStatement q range singular
                     (iPreferCount apiRequest) (contentType == TextCSV)
-          if range == Just emptyRange
+          if range == emptyRange
           then return $ errResponse status416 "HTTP Range error"
           else do
             row <- H.maybeEx stm
@@ -87,7 +87,7 @@ app dbStructure conf reqBody req =
               then responseLBS status404 [] ""
               else responseLBS status200 [contentTypeH] (fromMaybe "{}" body)
             else do
-              let frm = fromMaybe 0 $ rangeOffset <$> range
+              let frm = rangeOffset range
                   to = frm+queryTotal-1
                   contentRange = contentRangeH frm to tableTotal
                   status = rangeStatus frm to tableTotal

src/PostgREST/Config.hs

@@ -30,6 +30,7 @@ import           Network.Wai
 import           Network.Wai.Middleware.Cors (CorsResourcePolicy (..))
 import           Options.Applicative
 import           Paths_postgrest             (version)
+import           Safe                        (readMay)
 import           Web.JWT                     (Secret, secret)
 import           Prelude
 
@@ -41,6 +42,7 @@ data AppConfig = AppConfig {
   , configSchema    :: String
   , configJwtSecret :: Secret
   , configPool      :: Int
+  , configMaxRows   :: Maybe Int
   }
 
 argParser :: Parser AppConfig
@@ -53,6 +55,7 @@ argParser = AppConfig
   <*> (secret . cs <$>
       strOption    (long "jwt-secret" <> short 'j' <> help "secret used to encrypt and decrypt JWT tokens" <> metavar "SECRET" <> value "secret" <> showDefault))
   <*> option auto  (long "pool"       <> short 'o' <> help "max connections in database pool" <> metavar "COUNT" <> value 10 <> showDefault)
+  <*> (readMay <$> strOption  (long "max-rows"   <> short 'm' <> help "max rows in response" <> metavar "COUNT" <> value "infinity" <> showDefault))
 
 defaultCorsPolicy :: CorsResourcePolicy
 defaultCorsPolicy =  CorsResourcePolicy Nothing

src/PostgREST/QueryBuilder.hs

@@ -34,7 +34,6 @@ import qualified Data.Aeson              as JSON
 
 import           PostgREST.RangeQuery    (NonnegRange, rangeLimit, rangeOffset)
 import           Control.Error           (note, fromMaybe, mapMaybe)
-import           Control.Monad           (join)
 import qualified Data.HashMap.Strict     as HM
 import           Data.List               (find)
 import           Data.Monoid             ((<>))
@@ -61,10 +60,10 @@ instance Monoid PStmt where
   mempty = B.Stmt "" empty True
 type StatementT = PStmt -> PStmt
 
-createReadStatement :: SqlQuery -> Maybe NonnegRange -> Bool -> Bool -> Bool -> B.Stmt P.Postgres
+createReadStatement :: SqlQuery -> NonnegRange -> Bool -> Bool -> Bool -> B.Stmt P.Postgres
 createReadStatement selectQuery range isSingle countTable asCsv =
   B.Stmt (
-    wrapQuery selectQuery [
+    wrapLimitedQuery selectQuery [
       if countTable then countAllF else countNoneF,
       countF,
       "null", -- location header can not be calucalted
@@ -91,7 +90,7 @@ createWriteStatement selectQuery mutateQuery isSingle echoRequested
         else if isSingle then asJsonSingleF else asJsonF
       else "null"
 
-    ] selectQuery Nothing
+    ] selectQuery
   ) (V.singleton . B.encodeValue . JSON.Array . V.map JSON.Object $ rows) True
 
 addRelations :: Schema -> [Relation] -> Maybe ReadRequest -> ReadRequest -> Either Text ReadRequest
@@ -426,19 +425,27 @@ withSourceF s = "WITH " <> sourceSubqueryName <> " AS (" <> s <>")"
 fromF :: SqlFragment -> SqlFragment -> SqlFragment
 fromF sel limit = "FROM (" <> sel <> " " <> limit <> ") t"
 
-limitF :: Maybe NonnegRange -> SqlFragment
+limitF :: NonnegRange -> SqlFragment
 limitF r  = "LIMIT " <> limit <> " OFFSET " <> offset
   where
-    limit  = maybe "ALL" (cs . show) $ join $ rangeLimit <$> r
-    offset = cs . show $ fromMaybe 0 $ rangeOffset <$> r
+    limit  = maybe "ALL" (cs . show) $ rangeLimit r
+    offset = cs . show $ rangeOffset r
 
 selectStarF :: SqlFragment
 selectStarF = "SELECT * FROM " <> sourceSubqueryName
 
-wrapQuery :: SqlQuery -> [Text] -> Text ->  Maybe NonnegRange -> SqlQuery
-wrapQuery source selectColumns returnSelect range =
+wrapLimitedQuery :: SqlQuery -> [Text] -> Text -> NonnegRange -> SqlQuery
+wrapLimitedQuery source selectColumns returnSelect range =
   withSourceF source <>
   " SELECT " <>
   intercalate ", " selectColumns <>
   " " <>
   fromF returnSelect ( limitF range )
+
+wrapQuery :: SqlQuery -> [Text] -> Text -> SqlQuery
+wrapQuery source selectColumns returnSelect =
+  withSourceF source <>
+  " SELECT " <>
+  intercalate ", " selectColumns <>
+  " " <>
+  fromF returnSelect ""

src/PostgREST/RangeQuery.hs

@@ -3,6 +3,7 @@ module PostgREST.RangeQuery (
 , rangeRequested
 , rangeLimit
 , rangeOffset
+, restrictRange
 , NonnegRange
 ) where
 
@@ -25,20 +26,26 @@ import           Prelude
 
 type NonnegRange = Range Int
 
-rangeParse :: BS.ByteString -> Maybe NonnegRange
+rangeParse :: BS.ByteString -> NonnegRange
 rangeParse range = do
   let rangeRegex = "^([0-9]+)-([0-9]*)$" :: BS.ByteString
 
-  parsedRange <- listToMaybe (range =~ rangeRegex :: [[BS.ByteString]])
+  case listToMaybe (range =~ rangeRegex :: [[BS.ByteString]]) of
+    Just parsedRange ->
+      let [_, from, to] = readMaybe . cs <$> parsedRange
+          lower         = fromMaybe emptyRange   (rangeGeq <$> from)
+          upper         = fromMaybe (rangeGeq 0) (rangeLeq <$> to) in
+      rangeIntersection lower upper
+    Nothing -> rangeGeq 0
 
-  let [_, from, to] = readMaybe . cs <$> parsedRange
-  let lower = fromMaybe emptyRange   (rangeGeq <$> from)
-  let upper = fromMaybe (rangeGeq 0) (rangeLeq <$> to)
+rangeRequested :: RequestHeaders -> NonnegRange
+rangeRequested = rangeParse . fromMaybe "" . lookup hRange
 
-  return $ rangeIntersection lower upper
-
-rangeRequested :: RequestHeaders -> Maybe NonnegRange
-rangeRequested = (rangeParse =<<) . lookup hRange
+restrictRange :: Maybe Int -> NonnegRange -> NonnegRange
+restrictRange Nothing r = r
+restrictRange (Just limit) r =
+  rangeIntersection r $
+    Range BoundaryBelowAll (BoundaryAbove $ rangeOffset r + limit - 1)
 
 rangeLimit :: NonnegRange -> Maybe Int
 rangeLimit range =