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

Fixes https://github.com/PostgREST/postgrest/issues/1525
2020-06-05 13:01:36 -05:00
parent 69b09e312a
commit 0f0d617951
8 changed files with 137 additions and 50 deletions
@@ -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
-              case gucHeaders of
-                Left _ -> return . errorResponseFor $ GucHeadersError
-                Right ghdrs -> do
+              let (tableTotal, queryTotal, _ , body, gucHeaders, gucStatus) = row
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
+              case gucs of
+                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
-              case gucHeaders of
-                Left _ -> return . errorResponseFor $ GucHeadersError
-                Right ghdrs -> do
+              let (_, queryTotal, fields, body, gucHeaders, gucStatus) = row
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
+              case gucs of
+                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
-              case gucHeaders of
-                Left _ -> return . errorResponseFor $ GucHeadersError
-                Right ghdrs -> do
+              let (_, queryTotal, _, body, gucHeaders, gucStatus) = row
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
+              case gucs of
+                Left err -> return $ errorResponseFor err
+                Right (ghdrs, gstatus) -> do
                  let
                    updateIsNoOp = S.null cols
-                    status | queryTotal == 0 && not updateIsNoOp      = status404
-                           | iPreferRepresentation apiRequest == Full = status200
-                           | otherwise                                = status204
+                    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
-                case gucHeaders of
-                  Left _ -> return . errorResponseFor $ GucHeadersError
-                  Right ghdrs -> do
+                let (_, queryTotal, _, body, gucHeaders, gucStatus) = row
+                    gucs =  (,) <$> gucHeaders <*> gucStatus
+                case gucs of
+                  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
-              case gucHeaders of
-                Left _ -> return . errorResponseFor $ GucHeadersError
-                Right ghdrs -> do
+              let (_, queryTotal, _, body, gucHeaders, gucStatus) = row
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
+              case gucs of
+                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
-              case gucHeaders of
-                Left _ -> return . errorResponseFor $ GucHeadersError
-                Right ghdrs -> do
-                  let (status, contentRange) = rangeStatusHeader topLevelRange queryTotal tableTotal
+              let (tableTotal, queryTotal, body, gucHeaders, gucStatus) = row
+                  gucs =  (,) <$> gucHeaders <*> gucStatus
+              case gucs of
+                Left err -> return $ errorResponseFor err
+                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)
@@ -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 [
@@ -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
-    else "'[]'" :: Text
+    then currentSettingF "response.headers"
+    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), '')"
@@ -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
-                    <*> column header <*> column HD.bytea <*> column decodeGucHeaders
+standardRow = (,,,,,) <$> nullableColumn HD.int8 <*> column HD.int8
+                      <*> 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
-                        <*> column HD.bytea <*> column decodeGucHeaders
+        defGucHeaders = Right []
+        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 = first toS . JSON.eitherDecode . toS <$> HD.bytea
+decodeGucHeaders :: HD.Value (Either SimpleError [GucHeader])
+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