From 893b7a71264847deca6eaafa7550f5d4f0fd0c88 Mon Sep 17 00:00:00 2001
From: Diogo Biazus <diogob@gmail.com>
Date: Wed, 10 Feb 2016 22:45:52 -0500
Subject: [PATCH] Applies range headers to /rpc calls using LIMIT/OFFSET.

---
 src/PostgREST/App.hs          | 27 ++++++++++++++-------------
 src/PostgREST/Main.hs         |  2 --
 src/PostgREST/QueryBuilder.hs |  6 +++---
 test/Feature/QuerySpec.hs     |  5 ++++-
 4 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
index 000d66544..812379c83 100644
--- a/src/PostgREST/App.hs
+++ b/src/PostgREST/App.hs
@@ -71,13 +71,10 @@ app dbStructure conf reqBody req =
       case readSqlParts of
         Left e -> return $ responseLBS status400 [jsonH] $ cs e
         Right (q, cq) -> do
-          let range = restrictRange (configMaxRows conf) $ iRange apiRequest
-              singular = iPreferSingular apiRequest
+          let singular = iPreferSingular apiRequest
               stm = createReadStatement q cq range singular
                     (iPreferCount apiRequest) (contentType == TextCSV)
-          if range == emptyRange
-          then return $ errResponse status416 "HTTP Range error"
-          else do
+          respondToRange $ do
             row <- H.query () stm
             let (tableTotal, queryTotal, _ , body) = row
             if singular
@@ -165,14 +162,14 @@ app dbStructure conf reqBody req =
         then do
           let p = V.head payload
               jwtSecret = configJwtSecret conf
-
-          bodyJson <- H.query () (callProc qi p)
-          returnJWT <- H.query qi doesProcReturnJWT
-          return $ responseLBS status200 [jsonH]
-                 (let body = fromMaybe emptyArray bodyJson in
-                    if returnJWT
-                    then "{\"token\":\"" <> cs (tokenJWT jwtSecret body) <> "\"}"
-                    else cs $ encode body)
+          respondToRange $ do
+            bodyJson <- H.query () (callProc qi p range)
+            returnJWT <- H.query qi doesProcReturnJWT
+            return $ responseLBS status200 [jsonH]
+                  (let body = fromMaybe emptyArray bodyJson in
+                      if returnJWT
+                      then "{\"token\":\"" <> cs (tokenJWT jwtSecret body) <> "\"}"
+                      else cs $ encode body)
         else return notFound
 
     (ActionRead, TargetRoot, Nothing) -> do
@@ -196,6 +193,7 @@ app dbStructure conf reqBody req =
   allOrigins = ("Access-Control-Allow-Origin", "*") :: Header
   schema = cs $ configSchema conf
   apiRequest = userApiRequest schema req reqBody
+  range = restrictRange (configMaxRows conf) $ iRange apiRequest
   readDbRequest = DbRead <$> buildReadRequest (dbRelations dbStructure) apiRequest
   mutateDbRequest = DbMutate <$> buildMutateRequest apiRequest
   selectQuery = requestToQuery schema <$> readDbRequest
@@ -203,6 +201,9 @@ app dbStructure conf reqBody req =
   mutateQuery = requestToQuery schema <$> mutateDbRequest
   readSqlParts = (,) <$> selectQuery <*> countQuery
   mutateSqlParts = (,) <$> selectQuery <*> mutateQuery
+  respondToRange response = if range == emptyRange
+                            then return $ errResponse status416 "HTTP Range error"
+                            else response
 
 rangeStatus :: Integer -> Integer -> Maybe Integer -> Status
 rangeStatus _ _ Nothing = status200
diff --git a/src/PostgREST/Main.hs b/src/PostgREST/Main.hs
index ae840249d..c3aa95563 100644
--- a/src/PostgREST/Main.hs
+++ b/src/PostgREST/Main.hs
@@ -30,7 +30,6 @@ import           System.IO                            (BufferMode (..),
                                                        hSetBuffering, stderr,
                                                        stdin, stdout)
 import           Web.JWT                              (secret)
-
 #ifndef mingw32_HOST_OS
 import           System.Posix.Signals
 import           Control.Concurrent                   (myThreadId)
@@ -64,7 +63,6 @@ main = do
   Prelude.putStrLn $ "Listening on port " ++
     (show $ configPort conf :: String)
 
-
   pool <- P.acquire (configPool conf, 10, pgSettings)
 
 #ifndef mingw32_HOST_OS
diff --git a/src/PostgREST/QueryBuilder.hs b/src/PostgREST/QueryBuilder.hs
index 046040909..f551196df 100644
--- a/src/PostgREST/QueryBuilder.hs
+++ b/src/PostgREST/QueryBuilder.hs
@@ -203,8 +203,8 @@ addJoinConditions schema (Node (query, (n, r)) forest) =
     updatedForest = mapM (addJoinConditions schema) forest
     addCond query' con = query'{flt_=con ++ flt_ query'}
 
-callProc :: QualifiedIdentifier -> JSON.Object -> H.Query () (Maybe JSON.Value)
-callProc qi params =
+callProc :: QualifiedIdentifier -> JSON.Object -> NonnegRange -> H.Query () (Maybe JSON.Value)
+callProc qi params range =
   H.statement sql HE.unit decodeObj True
   where
     sql = [qc| SELECT array_to_json(
@@ -213,7 +213,7 @@ callProc qi params =
               from ({_callSql}) t |]
     _args = intercalate "," $ map _assignment (HM.toList params)
     _assignment (n,v) = pgFmtIdent n <> ":=" <> insertableValue v
-    _callSql = [qc| select * from {fromQi qi}({_args}) |] :: BS.ByteString
+    _callSql = [qc| select * from {fromQi qi}({_args}) {limitF range} |] :: BS.ByteString
     decodeObj = HD.maybeRow (HD.value HD.json)
 
 operators :: [(Text, SqlFragment)]
diff --git a/test/Feature/QuerySpec.hs b/test/Feature/QuerySpec.hs
index 6ae4194e7..b70109a04 100644
--- a/test/Feature/QuerySpec.hs
+++ b/test/Feature/QuerySpec.hs
@@ -370,7 +370,10 @@ spec = do
         [json| [{"data": {"id": 1, "foo": {"bar": "baz"}}}] |]
 
   describe "remote procedure call" $ do
-    context "a proc that returns a set" $
+    context "a proc that returns a set" $ do
+      it "returns paginated results" $
+        request methodPost "/rpc/getitemrange" (rangeHdrs (ByteRangeFromTo 0 0))  [json| { "min": 2, "max": 4 } |] `shouldRespondWith`
+          [json| [ {"id": 3} ] |]
       it "returns proper json" $
         post "/rpc/getitemrange" [json| { "min": 2, "max": 4 } |] `shouldRespondWith`
           [json| [ {"id": 3}, {"id":4} ] |]