Add ability to map GUC to http response headers

2017-10-18 17:41:55 -05:00
parent 38de56de4a
commit b9a591aecb
7 changed files with 109 additions and 17 deletions
@@ -10,6 +10,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 - #887, #601, Allow specifying dictionary and plain/phrase tsquery in full text search - @steve-chavez
 - #328, Allow doing GET on rpc - @steve-chavez
 - #917, Add ability to map RAISE errorcode/message to http status - @steve-chavez
+- #940, Add ability to map GUC to http response headers - @steve-chavez

 ### Fixed

@@ -6,7 +6,7 @@ module PostgREST.App (
 ) where

 import           Control.Applicative
-import           Data.Aeson                (toJSON)
+import           Data.Aeson                (toJSON, eitherDecode)
 import qualified Data.ByteString.Char8     as BS
 import           Data.Maybe
 import           Data.IORef                (IORef, readIORef)
@@ -44,7 +44,7 @@ import           PostgREST.DbRequestBuilder( readRequest
 import           PostgREST.Error           ( simpleError, pgError
                                           , apiRequestError
                                           , singularityError, binaryFieldError
-                                           , connectionLostError
+                                           , connectionLostError, gucHeadersError
                                           )
 import           PostgREST.RangeQuery      (allRange, rangeOffset)
 import           PostgREST.Middleware
@@ -248,18 +248,22 @@ app dbStructure conf apiRequest =
                  singular = contentType == CTSingularJSON
                  paramsAsSingleObject = iPreferSingleObjectParameter apiRequest
              row <- H.query () $
-                callProc qi prms returnsScalar q cq topLevelRange shouldCount
+                callProc qi prms returnsScalar q cq shouldCount
                         singular paramsAsSingleObject
                         (contentType == CTTextCSV)
                         (contentType == CTOctetStream) _isReadOnly bField
-              let (tableTotal, queryTotal, body) =
-                    fromMaybe (Just 0, 0, "[]") row
+              let (tableTotal, queryTotal, body, jsonHeaders) =
+                    fromMaybe (Just 0, 0, "[]", "[]") row
                  (status, contentRange) = rangeHeader queryTotal tableTotal
-              if singular && queryTotal /= 1
-                then do
-                  HT.condemn
-                  return $ singularityError (toInteger queryTotal)
-                else return $ responseLBS status [toHeader contentType, contentRange] (toS body)
+                  decodedHeaders = first toS $ eitherDecode $ toS jsonHeaders :: Either Text [GucHeader]
+              case decodedHeaders of
+                Left _ -> return gucHeadersError
+                Right hs ->
+                  if singular && queryTotal /= 1
+                    then do
+                      HT.condemn
+                      return $ singularityError (toInteger queryTotal)
+                    else return $ responseLBS status ([toHeader contentType, contentRange] ++ toHeaders hs) (toS body)

        (ActionInspect, TargetRoot, Nothing) -> do
          let host = configHost conf
@@ -10,6 +10,7 @@ module PostgREST.Error (
 , binaryFieldError
 , connectionLostError
 , encodeError
+, gucHeadersError
 ) where

 import           Protolude
@@ -79,6 +80,11 @@ binaryFieldError =
  simpleError HT.status406 [] (toS (toMime CTOctetStream) <>
  " requested but a single column was not selected")

+gucHeadersError :: Response
+gucHeadersError =
+ simpleError HT.status500 []
+ "response.headers guc must be a JSON array composed of objects with a single key and a string value"
+
 connectionLostError :: Response
 connectionLostError =
  simpleError HT.status503 [] "Database connection lost, retrying the connection."
@@ -142,10 +142,10 @@ createWriteStatement selectQuery mutateQuery wantSingle wantHdrs asCsv rep pKeys
    | wantSingle = asJsonSingleF
    | otherwise = asJsonF

-type ProcResults = (Maybe Int64, Int64, ByteString)
-callProc :: QualifiedIdentifier -> JSON.Object -> Bool -> SqlQuery -> SqlQuery -> NonnegRange ->
+type ProcResults = (Maybe Int64, Int64, ByteString, ByteString)
+callProc :: QualifiedIdentifier -> JSON.Object -> Bool -> SqlQuery -> SqlQuery ->
  Bool -> Bool -> Bool -> Bool -> Bool -> Bool -> Maybe FieldName -> H.Query () (Maybe ProcResults)
-callProc qi params returnsScalar selectQuery countQuery _ countTotal isSingle paramsAsJson asCsv asBinary isReadOnly binaryField =
+callProc qi params returnsScalar selectQuery countQuery countTotal isSingle paramsAsJson asCsv asBinary isReadOnly binaryField =
  unicodeStatement sql HE.unit decodeProc True
  where
    sql =
@@ -154,14 +154,16 @@ callProc qi params returnsScalar selectQuery countQuery _ countTotal isSingle pa
       SELECT
         {countResultF} AS total_result_set,
         1 AS page_total,
-         {scalarBodyF} as body
+         {scalarBodyF} AS body,
+         {responseHeaders} AS headers
       FROM ({selectQuery}) _postgrest_t;|]
     else [qc|
       WITH {sourceCTEName} AS (select * from {fromQi qi}({_args}))
       SELECT
         {countResultF} AS total_result_set,
         pg_catalog.count(_postgrest_t) AS page_total,
-         {bodyF} as body
+         {bodyF} AS body,
+         {responseHeaders} AS headers
       FROM ({selectQuery}) _postgrest_t;|]

    countResultF = if countTotal then "( "<> countQuery <> ")" else "null::bigint" :: Text
@@ -170,9 +172,10 @@ callProc qi params returnsScalar selectQuery countQuery _ countTotal isSingle pa
                else intercalate "," $ map _assignment (HM.toList params)
    _procName = qiName qi
    _assignment (n,v) = pgFmtIdent n <> ":=" <> insertableValue v
+    responseHeaders = "coalesce(nullif(current_setting('response.headers', true), ''), '[]')" :: Text -- nullif is used because of https://gist.github.com/steve-chavez/8d7033ea5655096903f3b52f8ed09a15
    decodeProc = HD.maybeRow procRow
-    procRow = (,,) <$> HD.nullableValue HD.int8 <*> HD.value HD.int8
-                   <*> HD.value HD.bytea
+    procRow = (,,,) <$> HD.nullableValue HD.int8 <*> HD.value HD.int8
+                    <*> HD.value HD.bytea <*> HD.value HD.bytea
    scalarBodyF
     | asBinary = asBinaryF _procName
     | otherwise = "(row_to_json(_postgrest_t)->" <> pgFmtLit _procName <> ")::character varying"
@@ -4,6 +4,7 @@ import           Protolude
 import qualified GHC.Show
 import           Data.Aeson
 import qualified Data.ByteString.Lazy as BL
+import qualified Data.CaseInsensitive as CI
 import qualified Data.HashMap.Strict  as M
 import           Data.Tree
 import qualified Data.Vector          as V
@@ -216,6 +217,22 @@ type NodeName = Text
 -- Rpc query param, only used for GET rpcs
 type RpcQParam = (Text, Text)

+{-|
+  Custom guc header, it's obtained by parsing the json in a:
+  `SET LOCAL "response.headers" = '[{"Set-Cookie": ".."}]'
+-}
+newtype GucHeader = GucHeader (Text, Text)
+
+instance FromJSON GucHeader where
+  parseJSON (Object o) = case headMay (M.toList o) of
+    Just (k, String s) | M.size o == 1 -> pure $ GucHeader (k, s)
+                       | otherwise     -> mzero
+    _ -> mzero
+  parseJSON _          = mzero
+
+toHeaders :: [GucHeader] -> [Header]
+toHeaders = map $ \(GucHeader (k, v)) -> (CI.mk $ toS k, toS v)
+
 {-|
  This type will hold information about which particular 'Relation' between two tables to choose when there are multiple ones.
  Specifically, it will contain the name of the foreign key or the join table in many to many relations.
@@ -286,6 +286,41 @@ spec =
    it "defaults to status 500 if RAISE code is PT not followed by a number" $
      get "/rpc/raise_bad_pt" `shouldRespondWith` 500

+    context "GUC headers" $ do
+      it "succeeds setting the headers" $ do
+        get "/rpc/get_projects_and_guc_headers?id=eq.2&select=id"
+          `shouldRespondWith` [json|[{"id": 2}]|]
+          {matchHeaders = [
+              matchContentTypeJson,
+              "X-Test"   <:> "key1=val1; someValue; key2=val2",
+              "X-Test-2" <:> "key1=val1"]}
+        get "/rpc/get_int_and_guc_headers?num=1"
+          `shouldRespondWith` [json|1|]
+          {matchHeaders = [
+              matchContentTypeJson,
+              "X-Test"   <:> "key1=val1; someValue; key2=val2",
+              "X-Test-2" <:> "key1=val1"]}
+        post "/rpc/get_int_and_guc_headers" [json|{"num": 1}|]
+          `shouldRespondWith` [json|1|]
+          {matchHeaders = [
+              matchContentTypeJson,
+              "X-Test"   <:> "key1=val1; someValue; key2=val2",
+              "X-Test-2" <:> "key1=val1"]}
+
+      it "fails when setting headers with wrong json structure" $ do
+        get "/rpc/bad_guc_headers_1" `shouldRespondWith` 500
+        get "/rpc/bad_guc_headers_2" `shouldRespondWith` 500
+        get "/rpc/bad_guc_headers_3" `shouldRespondWith` 500
+        post "/rpc/bad_guc_headers_1" [json|{}|] `shouldRespondWith` 500
+
+      it "can set the same http header twice" $
+        get "/rpc/set_cookie_twice"
+          `shouldRespondWith` "null"
+          {matchHeaders = [
+              matchContentTypeJson,
+              "Set-Cookie" <:> "sessionid=38afes7a8; HttpOnly; Path=/",
+              "Set-Cookie" <:> "id=a3fWa; Expires=Wed, 21 Oct 2015 07:28:00 GMT; Secure; HttpOnly"]}
+
    context "only for POST rpc" $ do
      context "expects a single json object" $ do
        it "does not expand posted json into parameters" $
@@ -1306,6 +1306,32 @@ begin
  raise sqlstate 'PT40A' using message = 'Wrong';
 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;
+$$ language sql;
+
+create or replace function test.get_int_and_guc_headers(num int) returns integer as $$
+  set local "response.headers" = '[{"X-Test":"key1=val1; someValue; key2=val2"},{"X-Test-2":"key1=val1"}]';
+  select num;
+$$ language sql;
+
+create or replace function test.bad_guc_headers_1() returns void as $$
+  set local "response.headers" = '{"X-Test": "invalid structure for headers"}';
+$$ language sql;
+
+create or replace function test.bad_guc_headers_2() returns void as $$
+  set local "response.headers" = '["invalid", "structure", "for", "headers"]';
+$$ language sql;
+
+create or replace function test.bad_guc_headers_3() returns void as $$
+  set local "response.headers" = '{"X-Test": "invalid", "X-Test-2": "structure", "X-Test-3": "for headers"}';
+$$ language sql;
+
+create or replace function test.set_cookie_twice() returns void as $$
+  set local "response.headers" = '[{"Set-Cookie": "sessionid=38afes7a8; HttpOnly; Path=/"}, {"Set-Cookie": "id=a3fWa; Expires=Wed, 21 Oct 2015 07:28:00 GMT; Secure; HttpOnly"}]';
+$$ language sql;
 --
 -- PostgreSQL database dump complete
 --