Merge pull request #404 from begriffs/max-rows

Add --max-rows option

CHANGELOG.md

@@ -11,6 +11,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 - Allow order by computed columns - @diogob
+- Set max rows in response with --max-rows - @begriffs
 
 ## [0.3.0.1] - 2015-11-27
 

docs/install/server.md

@@ -92,6 +92,10 @@ The possible flags are:
 <dd>Max connections to use in db pool. Defaults to to 10, but you
     should find an optimal value for your db by running the SQL
     command <code>show max_connections;</code></dd>
+
+<dt>-m, --max-rows</dt>
+<dd>Max number of rows to return in a read request. The default is
+    no limit.</dd>
 </dl>
 
 <div class="admonition note">

postgrest.cabal

@@ -45,6 +45,7 @@ executable postgrest
                      , parsec
                      , postgrest
                      , regex-tdfa
+                     , safe >= 0.3 && < 0.4
                      , scientific
                      , string-conversions
                      , text
@@ -98,6 +99,7 @@ library
                      , optparse-applicative
                      , parsec
                      , regex-tdfa
+                     , safe
                      , scientific
                      , string-conversions
                      , text
@@ -181,6 +183,7 @@ Test-Suite spec
                      , parsec
                      , process
                      , regex-tdfa
+                     , safe
                      , scientific
                      , string-conversions
                      , text

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 =

test/Feature/AuthSpec.hs

@@ -12,7 +12,7 @@ import SpecHelper
 spec :: Spec
 spec = beforeAll
   (clearTable "postgrest.auth") . afterAll_ (clearTable "postgrest.auth")
-  $ around withApp
+  $ around (withApp cfgDefault)
   $ describe "authorization" $ do
 
   it "hides tables that anonymous does not own" $

test/Feature/CorsSpec.hs

@@ -12,7 +12,7 @@ import Network.HTTP.Types
 -- }}}
 
 spec :: Spec
-spec = around withApp $ describe "CORS" $ do
+spec = around (withApp cfgDefault) $ describe "CORS" $ do
     let preflightHeaders = [
           ("Accept", "*/*"),
           ("Origin", "http://example.com"),

test/Feature/DeleteSpec.hs

@@ -8,7 +8,7 @@ import Network.HTTP.Types
 
 spec :: Spec
 spec = beforeAll (clearTable "items" >> createItems 15) . afterAll_ (clearTable "items")
-  . around withApp $
+  . around (withApp cfgDefault) $
   describe "Deleting" $ do
     context "existing record" $ do
       it "succeeds with 204 and deletion count" $

test/Feature/InsertSpec.hs

@@ -17,7 +17,7 @@ import Control.Monad (replicateM_)
 import TestTypes(IncPK(..), CompoundPK(..))
 
 spec :: Spec
-spec = afterAll_ resetDb $ around withApp $ do
+spec = afterAll_ resetDb $ around (withApp cfgDefault) $ do
   describe "Posting new record" $ do
     after_ (clearTable "menagerie") . context "disparate csv types" $ do
       it "accepts disparate json types" $ do

test/Feature/QueryLimitedSpec.hs

@@ -0,0 +1,31 @@
+module Feature.QueryLimitedSpec where
+
+import Test.Hspec hiding (pendingWith)
+import Test.Hspec.Wai
+import Test.Hspec.Wai.JSON
+import Network.HTTP.Types
+import Network.Wai.Test (SResponse(simpleHeaders, simpleStatus))
+
+import SpecHelper
+
+spec :: Spec
+spec =
+  beforeAll (clearTable "items" >> createItems 15)
+   . afterAll_ (clearTable "items")
+   . around (withApp $ cfgLimitRows 3) $
+  describe "Requesting many items with server limits enabled" $ do
+    it "restricts results" $
+      get "/items"
+        `shouldRespondWith` ResponseMatcher {
+          matchBody    = Just [json| [{"id":1},{"id":2},{"id":3}] |]
+        , matchStatus  = 206
+        , matchHeaders = ["Content-Range" <:> "0-2/15"]
+        }
+
+    it "respects additional client limiting" $ do
+      r <- request methodGet  "/items"
+                   (rangeHdrs $ ByteRangeFromTo 0 1) ""
+      liftIO $ do
+        simpleHeaders r `shouldSatisfy`
+          matchHeader "Content-Range" "0-1/15"
+        simpleStatus r `shouldBe` partialContent206

test/Feature/QuerySpec.hs

@@ -22,7 +22,7 @@ spec =
        createLikableStrings >>
        createJsonData)
    . afterAll_ (clearTable "items" >> clearTable "complex_items" >> clearTable "no_pk" >> clearTable "simple_pk")
-   . around withApp $ do
+   . around (withApp cfgDefault) $ do
 
   describe "Querying a table with a column called count" $
     it "should not confuse count column with pg_catalog.count aggregate" $

test/Feature/RangeSpec.hs

@@ -10,7 +10,7 @@ import SpecHelper
 
 spec :: Spec
 spec = beforeAll (clearTable "items" >> createItems 15) . afterAll_ (clearTable "items")
-  . around withApp $
+  . around (withApp cfgDefault) $
   describe "GET /items" $ do
 
     context "without range headers" $ do

test/Feature/StructureSpec.hs

@@ -9,7 +9,7 @@ import SpecHelper
 import Network.HTTP.Types
 
 spec :: Spec
-spec = around withApp $ do
+spec = around (withApp cfgDefault) $ do
   describe "GET /" $ do
     it "lists views in schema" $
       request methodGet "/" [] ""

test/SpecHelper.hs

@@ -41,8 +41,13 @@ isLeft :: Either a b -> Bool
 isLeft (Left _ ) = True
 isLeft _ = False
 
-cfg :: AppConfig
-cfg = AppConfig dbString 3000 "postgrest_anonymous" "test" (secret "safe") 10
+cfgDefault :: AppConfig
+cfgDefault = AppConfig dbString 3000 "postgrest_anonymous" "test" (secret "safe") 10 Nothing
+
+cfgLimitRows :: Int -> AppConfig
+cfgLimitRows limit =
+  AppConfig dbString 3000 "postgrest_anonymous" "test"
+    (secret "safe") 10 (Just limit)
 
 testPoolOpts :: PoolSettings
 testPoolOpts = fromMaybe (error "bad settings") $ H.poolSettings 1 30
@@ -50,20 +55,20 @@ testPoolOpts = fromMaybe (error "bad settings") $ H.poolSettings 1 30
 pgSettings :: P.Settings
 pgSettings = P.StringSettings $ cs dbString
 
-withApp :: ActionWith Application -> IO ()
-withApp perform = do
+withApp :: AppConfig -> ActionWith Application -> IO ()
+withApp config perform = do
   pool :: H.Pool P.Postgres
     <- H.acquirePool pgSettings testPoolOpts
 
   let txSettings = Just (H.ReadCommitted, Just True)
-  dbOrError <- H.session pool $ H.tx txSettings $ getDbStructure (cs $ configSchema cfg)
+  dbOrError <- H.session pool $ H.tx txSettings $ getDbStructure (cs $ configSchema config)
   db <- either (fail . show) return dbOrError
 
   perform $ middle $ \req resp -> do
     time <- getPOSIXTime
     body <- strictRequestBody req
     result <- liftIO $ H.session pool $ H.tx txSettings
-      $ runWithClaims cfg time (app db cfg body) req
+      $ runWithClaims config time (app db config body) req
     either (resp . pgErrResponse) resp result
 
   where middle = defaultMiddle