Add ability to map GUC to http response headers

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) =
-                    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

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)
-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"

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.