From b9a591aecbbc573742f28842bc62c58932f014be Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Fri, 22 Sep 2017 18:33:08 -0500 Subject: [PATCH] Add ability to map GUC to http response headers --- CHANGELOG.md | 1 + src/PostgREST/App.hs | 24 ++++++++++++++---------- src/PostgREST/Error.hs | 6 ++++++ src/PostgREST/QueryBuilder.hs | 17 ++++++++++------- src/PostgREST/Types.hs | 17 +++++++++++++++++ test/Feature/RpcSpec.hs | 35 +++++++++++++++++++++++++++++++++++ test/fixtures/schema.sql | 26 ++++++++++++++++++++++++++ 7 files changed, 109 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e6905ba01..543bb0ce6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs index 622e3c9a5..639d3d788 100644 --- a/src/PostgREST/App.hs +++ b/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 diff --git a/src/PostgREST/Error.hs b/src/PostgREST/Error.hs index 126a2d166..b7b331927 100644 --- a/src/PostgREST/Error.hs +++ b/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." diff --git a/src/PostgREST/QueryBuilder.hs b/src/PostgREST/QueryBuilder.hs index 148eb787f..e4121908a 100644 --- a/src/PostgREST/QueryBuilder.hs +++ b/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" diff --git a/src/PostgREST/Types.hs b/src/PostgREST/Types.hs index 3e9efd3cf..8e94d0252 100644 --- a/src/PostgREST/Types.hs +++ b/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. diff --git a/test/Feature/RpcSpec.hs b/test/Feature/RpcSpec.hs index b6f03a616..35413feef 100644 --- a/test/Feature/RpcSpec.hs +++ b/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" $ diff --git a/test/fixtures/schema.sql b/test/fixtures/schema.sql index 0227345db..f85db180f 100755 --- a/test/fixtures/schema.sql +++ b/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 --