Add ability to map GUC to http response headers

CHANGELOG.md

@@ -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
 

src/PostgREST/App.hs

@@ -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) =
+              let (tableTotal, queryTotal, body, jsonHeaders) =
-                    fromMaybe (Just 0, 0, "[]") row
+                    fromMaybe (Just 0, 0, "[]", "[]") row
                   (status, contentRange) = rangeHeader queryTotal tableTotal
-              if singular && queryTotal /= 1
+                  decodedHeaders = first toS $ eitherDecode $ toS jsonHeaders :: Either Text [GucHeader]
-                then do
+              case decodedHeaders of
-                  HT.condemn
+                Left _ -> return gucHeadersError
-                  return $ singularityError (toInteger queryTotal)
+                Right hs ->
-                else return $ responseLBS status [toHeader contentType, contentRange] (toS body)
+                  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

src/PostgREST/Error.hs

@@ -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."

src/PostgREST/QueryBuilder.hs

@@ -142,10 +142,10 @@ createWriteStatement selectQuery mutateQuery wantSingle wantHdrs asCsv rep pKeys
     | wantSingle = asJsonSingleF
     | otherwise = asJsonF
 
-type ProcResults = (Maybe Int64, Int64, ByteString)
+type ProcResults = (Maybe Int64, Int64, ByteString, ByteString)
-callProc :: QualifiedIdentifier -> JSON.Object -> Bool -> SqlQuery -> SqlQuery -> NonnegRange ->
+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
+    procRow = (,,,) <$> HD.nullableValue HD.int8 <*> HD.value HD.int8
-                   <*> HD.value HD.bytea
+                    <*> HD.value HD.bytea <*> HD.value HD.bytea
     scalarBodyF
      | asBinary = asBinaryF _procName
      | otherwise = "(row_to_json(_postgrest_t)->" <> pgFmtLit _procName <> ")::character varying"

src/PostgREST/Types.hs

@@ -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.

test/Feature/RpcSpec.hs

@@ -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" $

test/fixtures/schema.sql

@@ -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
 --