Allow http status override through response.status guc (#1541)

Fixes https://github.com/PostgREST/postgrest/issues/1525

CHANGELOG.md

@@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 
+ - #1525, Allow http status override through response.status guc - @steve-chavez
+
 ### Fixed
 
  - #1530, Fix how the PostgREST version is shown in the help text when the `.git` directory is not available - @monacoremo

nix/tests.nix

@@ -46,7 +46,7 @@ let
 
         EOF
 
-        ${withTmpDb postgresql} ${cabal-install}/bin/cabal v2-test \
+        ${withTmpDb postgresql} ${cabal-install}/bin/cabal v2-test -f FailOnWarn \
           --test-show-detail=direct
 
         cat << EOF

src/PostgREST/App.hs

@@ -136,17 +136,19 @@ app dbStructure proc cols conf apiRequest =
                         (contentType == CTTextCSV) bField pgVer
                   explStm = createExplainStatement cq
               row <- H.statement () stm
-              let (tableTotal, queryTotal, _ , body, gucHeaders) = row
+              let (tableTotal, queryTotal, _ , body, gucHeaders, gucStatus) = row
-              case gucHeaders of
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
-                Left _ -> return . errorResponseFor $ GucHeadersError
+              case gucs of
-                Right ghdrs -> do
+                Left err -> return $ errorResponseFor err
+                Right (ghdrs, gstatus) -> do
                   total <- if | plannedCount   -> H.statement () explStm
                               | estimatedCount -> if tableTotal > (fromIntegral <$> maxRows)
                                                     then do estTotal <- H.statement () explStm
                                                             pure $ if estTotal > tableTotal then estTotal else tableTotal
                                                     else pure tableTotal
                               | otherwise      -> pure tableTotal
-                  let (status, contentRange) = rangeStatusHeader topLevelRange queryTotal total
+                  let (rangeStatus, contentRange) = rangeStatusHeader topLevelRange queryTotal total
+                      status = fromMaybe rangeStatus gstatus
                       headers = addHeadersIfNotIncluded (catMaybes [
                                   Just $ toHeader contentType, Just contentRange,
                                   Just $ contentLocationH tName (iCanonicalQS apiRequest), profileH])
@@ -166,14 +168,16 @@ app dbStructure proc cols conf apiRequest =
                     (contentType == CTSingularJSON) True
                     (contentType == CTTextCSV) (iPreferRepresentation apiRequest) pkCols pgVer
               row <- H.statement (toS $ pjRaw pJson) stm
-              let (_, queryTotal, fields, body, gucHeaders) = row
+              let (_, queryTotal, fields, body, gucHeaders, gucStatus) = row
-              case gucHeaders of
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
-                Left _ -> return . errorResponseFor $ GucHeadersError
+              case gucs of
-                Right ghdrs -> do
+                Left err -> return $ errorResponseFor err
+                Right (ghdrs, gstatus) -> do
                   let
                     (ctHeaders, rBody) = if iPreferRepresentation apiRequest == Full
                                           then ([Just $ toHeader contentType, profileH], toS body)
                                           else ([], mempty)
+                    status = fromMaybe status201 gstatus
                     headers = addHeadersIfNotIncluded (catMaybes ([
                           if null fields
                             then Nothing
@@ -188,7 +192,7 @@ app dbStructure proc cols conf apiRequest =
                       HT.condemn
                       return . errorResponseFor . singularityError $ queryTotal
                     else
-                      return $ responseLBS status201 headers rBody
+                      return $ responseLBS status headers rBody
 
         (ActionUpdate, TargetIdent (QualifiedIdentifier tSchema tName), Just pJson) ->
           case mutateSqlParts tSchema tName of
@@ -198,15 +202,17 @@ app dbStructure proc cols conf apiRequest =
                     (contentType == CTSingularJSON) False (contentType == CTTextCSV)
                     (iPreferRepresentation apiRequest) [] pgVer
               row <- H.statement (toS $ pjRaw pJson) stm
-              let (_, queryTotal, _, body, gucHeaders) = row
+              let (_, queryTotal, _, body, gucHeaders, gucStatus) = row
-              case gucHeaders of
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
-                Left _ -> return . errorResponseFor $ GucHeadersError
+              case gucs of
-                Right ghdrs -> do
+                Left err -> return $ errorResponseFor err
+                Right (ghdrs, gstatus) -> do
                   let
                     updateIsNoOp = S.null cols
-                    status | queryTotal == 0 && not updateIsNoOp      = status404
+                    defStatus | queryTotal == 0 && not updateIsNoOp      = status404
                               | iPreferRepresentation apiRequest == Full = status200
                               | otherwise                                = status204
+                    status = fromMaybe defStatus gstatus
                     contentRangeHeader = contentRangeH 0 (queryTotal - 1) $ if shouldCount then Just queryTotal else Nothing
                     (ctHeaders, rBody) = if iPreferRepresentation apiRequest == Full
                                           then ([Just $ toHeader contentType, profileH], toS body)
@@ -229,12 +235,14 @@ app dbStructure proc cols conf apiRequest =
                 row <- H.statement (toS $ pjRaw pJson) $
                        createWriteStatement sq mq (contentType == CTSingularJSON) False
                                             (contentType == CTTextCSV) (iPreferRepresentation apiRequest) [] pgVer
-                let (_, queryTotal, _, body, gucHeaders) = row
+                let (_, queryTotal, _, body, gucHeaders, gucStatus) = row
-                case gucHeaders of
+                    gucs =  (,) <$> gucHeaders <*> gucStatus
-                  Left _ -> return . errorResponseFor $ GucHeadersError
+                case gucs of
-                  Right ghdrs -> do
+                  Left err -> return $ errorResponseFor err
+                  Right (ghdrs, gstatus) -> do
                     let headers = addHeadersIfNotIncluded (catMaybes [Just $ toHeader contentType, profileH]) (unwrapGucHeader <$> ghdrs)
-                        (status, rBody) = if iPreferRepresentation apiRequest == Full then (status200, toS body) else (status204, mempty)
+                        (defStatus, rBody) = if iPreferRepresentation apiRequest == Full then (status200, toS body) else (status204, mempty)
+                        status = fromMaybe defStatus gstatus
                     -- Makes sure the querystring pk matches the payload pk
                     -- e.g. PUT /items?id=eq.1 { "id" : 1, .. } is accepted, PUT /items?id=eq.14 { "id" : 2, .. } is rejected
                     -- If this condition is not satisfied then nothing is inserted, check the WHERE for INSERT in QueryBuilder.hs to see how it's done
@@ -254,12 +262,14 @@ app dbStructure proc cols conf apiRequest =
                     (contentType == CTTextCSV)
                     (iPreferRepresentation apiRequest) [] pgVer
               row <- H.statement mempty stm
-              let (_, queryTotal, _, body, gucHeaders) = row
+              let (_, queryTotal, _, body, gucHeaders, gucStatus) = row
-              case gucHeaders of
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
-                Left _ -> return . errorResponseFor $ GucHeadersError
+              case gucs of
-                Right ghdrs -> do
+                Left err -> return $ errorResponseFor err
+                Right (ghdrs, gstatus) -> do
                   let
-                    status = if iPreferRepresentation apiRequest == Full then status200 else status204
+                    defStatus = if iPreferRepresentation apiRequest == Full then status200 else status204
+                    status = fromMaybe defStatus gstatus
                     contentRangeHeader = contentRangeH 1 0 $ if shouldCount then Just queryTotal else Nothing
                     (ctHeaders, rBody) = if iPreferRepresentation apiRequest == Full
                                           then ([Just $ toHeader contentType, profileH], toS body)
@@ -294,11 +304,13 @@ app dbStructure proc cols conf apiRequest =
                         (contentType == CTTextCSV) (contentType `elem` rawContentTypes) (preferParams == Just MultipleObjects)
                         bField pgVer
               row <- H.statement (toS $ pjRaw pJson) stm
-              let (tableTotal, queryTotal, body, gucHeaders) = row
+              let (tableTotal, queryTotal, body, gucHeaders, gucStatus) = row
-              case gucHeaders of
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
-                Left _ -> return . errorResponseFor $ GucHeadersError
+              case gucs of
-                Right ghdrs -> do
+                Left err -> return $ errorResponseFor err
-                  let (status, contentRange) = rangeStatusHeader topLevelRange queryTotal tableTotal
+                Right (ghdrs, gstatus) -> do
+                  let (rangeStatus, contentRange) = rangeStatusHeader topLevelRange queryTotal tableTotal
+                      status = fromMaybe rangeStatus gstatus
                       headers = addHeadersIfNotIncluded
                         (catMaybes [Just $ toHeader contentType, Just contentRange, profileH])
                         (unwrapGucHeader <$> ghdrs)

src/PostgREST/Error.hs

@@ -221,6 +221,7 @@ checkIsFatal _ = Nothing
 
 data SimpleError
   = GucHeadersError
+  | GucStatusError
   | BinaryFieldError ContentType
   | ConnectionLostError
   | PutMatchingPkError
@@ -233,6 +234,7 @@ data SimpleError
 
 instance PgrstError SimpleError where
   status GucHeadersError         = HT.status500
+  status GucStatusError          = HT.status500
   status (BinaryFieldError _)    = HT.status406
   status ConnectionLostError     = HT.status503
   status PutMatchingPkError      = HT.status400
@@ -249,6 +251,8 @@ instance PgrstError SimpleError where
 instance JSON.ToJSON SimpleError where
   toJSON GucHeadersError           = JSON.object [
     "message" .= ("response.headers guc must be a JSON array composed of objects with a single key and a string value" :: Text)]
+  toJSON GucStatusError           = JSON.object [
+    "message" .= ("response.status guc must be a valid status code" :: Text)]
   toJSON (BinaryFieldError ct)          = JSON.object [
     "message" .= ((toS (toMime ct) <> " requested but more than one column was selected") :: Text)]
   toJSON ConnectionLostError       = JSON.object [

src/PostgREST/Private/QueryFragment.hs

@@ -202,5 +202,16 @@ returningF qi returnings =
 responseHeadersF :: PgVersion -> SqlFragment
 responseHeadersF pgVer =
   if pgVer >= pgVersion96
-    then "coalesce(nullif(current_setting('response.headers', true), ''), '[]')" :: Text -- nullif is used because of https://gist.github.com/steve-chavez/8d7033ea5655096903f3b52f8ed09a15
+    then currentSettingF "response.headers"
-    else "'[]'" :: Text
+    else "null" :: Text
+
+responseStatusF :: PgVersion -> SqlFragment
+responseStatusF pgVer =
+  if pgVer >= pgVersion96
+    then currentSettingF "response.status"
+    else "null" :: Text
+
+currentSettingF :: SqlFragment -> SqlFragment
+currentSettingF setting =
+  -- nullif is used because of https://gist.github.com/steve-chavez/8d7033ea5655096903f3b52f8ed09a15
+  "nullif(current_setting(" <> pgFmtLit setting <> ", true), '')"

src/PostgREST/Statements.hs

@@ -22,9 +22,12 @@ import           Data.Aeson                      as JSON
 import qualified Data.Aeson.Lens                 as L
 import qualified Data.ByteString.Char8           as BS
 import           Data.Maybe
+import           Data.Text.Read                  (decimal)
 import qualified Hasql.Decoders                  as HD
 import qualified Hasql.Encoders                  as HE
 import qualified Hasql.Statement                 as H
+import           Network.HTTP.Types.Status
+import           PostgREST.Error
 import           PostgREST.Private.Common
 import           PostgREST.Private.QueryFragment
 import           PostgREST.Types
@@ -37,7 +40,7 @@ import           Text.InterpolatedString.Perl6   (qc)
     is represented as a list of strings containing variable bindings like
     @"k1=eq.42"@, or the empty list if there is no location header.
 -}
-type ResultsWithCount = (Maybe Int64, Int64, [BS.ByteString], BS.ByteString, Either Text [GucHeader])
+type ResultsWithCount = (Maybe Int64, Int64, [BS.ByteString], BS.ByteString, Either SimpleError [GucHeader], Either SimpleError (Maybe Status))
 
 createWriteStatement :: SqlQuery -> SqlQuery -> Bool -> Bool -> Bool ->
                         PreferRepresentation -> [Text] -> PgVersion ->
@@ -53,7 +56,8 @@ createWriteStatement selectQuery mutateQuery wantSingle isInsert asCsv rep pKeys
         pg_catalog.count(_postgrest_t) AS page_total,
         {locF} AS header,
         {bodyF} AS body,
-        {responseHeadersF pgVer} AS response_headers
+        {responseHeadersF pgVer} AS response_headers,
+        {responseStatusF pgVer} AS response_status
       FROM ({selectF}) _postgrest_t |]
 
   locF =
@@ -78,7 +82,7 @@ createWriteStatement selectQuery mutateQuery wantSingle isInsert asCsv rep pKeys
 
   decodeStandard :: HD.Result ResultsWithCount
   decodeStandard =
-   fromMaybe (Nothing, 0, [], mempty, Right []) <$> HD.rowMaybe standardRow
+   fromMaybe (Nothing, 0, [], mempty, Right [], Right Nothing) <$> HD.rowMaybe standardRow
 
 createReadStatement :: SqlQuery -> SqlQuery -> Bool -> Bool -> Bool -> Maybe FieldName -> PgVersion ->
                        H.Statement () ResultsWithCount
@@ -94,7 +98,8 @@ createReadStatement selectQuery countQuery isSingle countTotal asCsv binaryField
         pg_catalog.count(_postgrest_t) AS page_total,
         {noLocationF} AS header,
         {bodyF} AS body,
-        {responseHeadersF pgVer} AS response_headers
+        {responseHeadersF pgVer} AS response_headers,
+        {responseStatusF pgVer} AS response_status
       FROM ( SELECT * FROM {sourceCTEName}) _postgrest_t |]
 
   (countCTEF, countResultF) = countF countQuery countTotal
@@ -114,12 +119,14 @@ createReadStatement selectQuery countQuery isSingle countTotal asCsv binaryField
     for that common type of query.
 -}
 standardRow :: HD.Row ResultsWithCount
-standardRow = (,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
+standardRow = (,,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
-                    <*> column header <*> column HD.bytea <*> column decodeGucHeaders
+                      <*> column header <*> column HD.bytea
+                      <*> (fromMaybe (Right []) <$> nullableColumn decodeGucHeaders)
+                      <*> (fromMaybe (Right Nothing) <$> nullableColumn decodeGucStatus)
   where
     header = HD.array $ HD.dimension replicateM $ element HD.bytea
 
-type ProcResults = (Maybe Int64, Int64, ByteString, Either Text [GucHeader])
+type ProcResults = (Maybe Int64, Int64, ByteString, Either SimpleError [GucHeader], Either SimpleError (Maybe Status))
 
 callProcStatement :: Bool -> SqlQuery -> SqlQuery -> SqlQuery -> Bool ->
                      Bool -> Bool -> Bool -> Bool -> Maybe FieldName -> PgVersion ->
@@ -134,7 +141,8 @@ callProcStatement returnsScalar callProcQuery selectQuery countQuery countTotal
         {countResultF} AS total_result_set,
         pg_catalog.count(_postgrest_t) AS page_total,
         {bodyF} AS body,
-        {responseHeadersF pgVer} AS response_headers
+        {responseHeadersF pgVer} AS response_headers,
+        {responseStatusF pgVer} AS response_status
       FROM ({selectQuery}) _postgrest_t;|]
 
     (countCTEF, countResultF) = countF countQuery countTotal
@@ -153,10 +161,14 @@ callProcStatement returnsScalar callProcQuery selectQuery countQuery countTotal
 
     decodeProc :: HD.Result ProcResults
     decodeProc =
-      fromMaybe (Just 0, 0, mempty, Right []) <$> HD.rowMaybe procRow
+      fromMaybe (Just 0, 0, mempty, defGucHeaders, defGucStatus) <$> HD.rowMaybe procRow
       where
-        procRow = (,,,) <$> nullableColumn HD.int8 <*> column HD.int8
+        defGucHeaders = Right []
-                        <*> column HD.bytea <*> column decodeGucHeaders
+        defGucStatus  = Right Nothing
+        procRow = (,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
+                         <*> column HD.bytea
+                         <*> (fromMaybe defGucHeaders <$> nullableColumn decodeGucHeaders)
+                         <*> (fromMaybe defGucStatus <$> nullableColumn decodeGucStatus)
 
 createExplainStatement :: SqlQuery -> H.Statement () (Maybe Int64)
 createExplainStatement countQuery =
@@ -179,5 +191,8 @@ createExplainStatement countQuery =
 unicodeStatement :: Text -> HE.Params a -> HD.Result b -> Bool -> H.Statement a b
 unicodeStatement = H.Statement . encodeUtf8
 
-decodeGucHeaders :: HD.Value (Either Text [GucHeader])
+decodeGucHeaders :: HD.Value (Either SimpleError [GucHeader])
-decodeGucHeaders = first toS . JSON.eitherDecode . toS <$> HD.bytea
+decodeGucHeaders = first (const GucHeadersError) . JSON.eitherDecode . toS <$> HD.bytea
+
+decodeGucStatus :: HD.Value (Either SimpleError (Maybe Status))
+decodeGucStatus = first (const GucStatusError) . fmap (Just . toEnum . fst) . decimal <$> HD.text

test/Feature/PgVersion96Spec.hs

@@ -161,6 +161,27 @@ spec =
           let respHeaders = simpleHeaders r
           respHeaders `shouldSatisfy` noBlankHeader
 
+    context "GUC status override" $ do
+      it "can override the status on RPC" $
+        get "/rpc/send_body_status_403"
+          `shouldRespondWith`
+          [json|{"message" : "invalid user or password"}|]
+          { matchStatus  = 403
+          , matchHeaders = [ matchContentTypeJson ]
+          }
+
+      it "can override the status through trigger" $
+        request methodPatch "/stuff?id=eq.1" [] [json|[{"name": "updated stuff 1"}]|]
+          `shouldRespondWith` 205
+
+      it "fails when setting invalid status guc" $
+        get "/rpc/send_bad_status"
+          `shouldRespondWith`
+          [json|{"message":"response.status guc must be a valid status code"}|]
+          { matchStatus  = 500
+          , matchHeaders = [ matchContentTypeJson ]
+          }
+
     context "Use of the phraseto_tsquery function" $ do
       it "finds matches" $
         get "/tsearch?text_search_vector=phfts.The%20Fat%20Cats" `shouldRespondWith`

test/fixtures/schema.sql

@@ -1086,6 +1086,20 @@ begin
 end;
 $$ language plpgsql;
 
+create or replace function test.send_body_status_403() returns pg_catalog.json as $$
+begin
+  perform set_config('response.status', '403', true);
+  return json_build_object('message', 'invalid user or password');
+end;
+$$ language plpgsql;
+
+create or replace function test.send_bad_status() returns pg_catalog.json as $$
+begin
+  perform set_config('response.status', 'bad', true);
+  return null;
+end;
+$$ language plpgsql;
+
 create or replace function test.get_projects_and_guc_headers() returns setof test.projects as $$
   set local "response.headers" = '[{"X-Test": "key1=val1; someValue; key2=val2"}, {"X-Test-2": "key1=val1"}]';
   select * from test.projects;
@@ -1692,8 +1706,7 @@ create table private.stuff(
 
 create view test.stuff as select * from private.stuff;
 
-create or replace function location_for_stuff() returns trigger
+create or replace function location_for_stuff() returns trigger as $$
-    as $$
 begin
     insert into private.stuff values (new.id, new.name);
     if new.id is not null
@@ -1709,6 +1722,15 @@ end
 $$ language plpgsql security definer;
 create trigger location_for_stuff instead of insert on test.stuff for each row execute procedure test.location_for_stuff();
 
+create or replace function status_205_for_updated_stuff() returns trigger as $$
+begin
+    update private.stuff set id = new.id, name = new.name;
+    perform set_config('response.status' , '205' , true);
+    return new;
+end
+$$ language plpgsql security definer;
+create trigger status_205_for_updated_stuff instead of update on test.stuff for each row execute procedure test.status_205_for_updated_stuff();
+
 create table loc_test (
   id int primary key
 , c text