Add response.headers on GET/POST/PATCH/PUT/DELETE

test/Feature/PgVersion96Spec.hs

@@ -1,6 +1,7 @@
 module Feature.PgVersion96Spec where
 
-import Network.Wai (Application)
+import Network.HTTP.Types
+import Network.Wai        (Application)
 
 import Test.Hspec
 import Test.Hspec.Wai
@@ -12,7 +13,7 @@ import SpecHelper
 spec :: SpecWith Application
 spec =
   describe "features supported on PostgreSQL 9.6" $ do
-    context "GUC headers" $ do
+    context "GUC headers on function calls" $ do
       it "succeeds setting the headers" $ do
         get "/rpc/get_projects_and_guc_headers?id=eq.2&select=id"
           `shouldRespondWith` [json|[{"id": 2}]|]
@@ -67,6 +68,50 @@ spec =
               "Set-Cookie" <:> "sessionid=38afes7a8; HttpOnly; Path=/",
               "Set-Cookie" <:> "id=a3fWa; Expires=Wed, 21 Oct 2015 07:28:00 GMT; Secure; HttpOnly"]}
 
+    context "GUC headers on all other methods via pre-request" $ do
+      it "succeeds setting the headers on GET and HEAD" $ do
+        request methodGet "/items?id=eq.1" [("User-Agent", "MSIE 6.0")] mempty
+          `shouldRespondWith` [json|[{"id": 1}]|]
+          {matchHeaders = [
+              matchContentTypeJson,
+              "Cache-Control" <:> "no-cache, no-store, must-revalidate"]}
+
+        request methodHead "/items?id=eq.1" [("User-Agent", "MSIE 7.0")] mempty
+          `shouldRespondWith` ""
+          {matchHeaders = ["Cache-Control" <:> "no-cache, no-store, must-revalidate"]}
+
+        request methodHead "/projects" [("Accept", "text/csv")] mempty
+          `shouldRespondWith` ""
+          {matchHeaders = ["Content-Disposition" <:> "attachment; filename=projects.csv"]}
+
+      it "succeeds setting the headers on POST" $
+        request methodPost "/items" [] [json|[{"id": 11111}]|]
+          `shouldRespondWith` ""
+          { matchStatus = 201
+          , matchHeaders = ["X-Custom-Header" <:> "mykey=myval"]
+          }
+
+      it "succeeds setting the headers on PATCH" $
+        request methodPatch "/items?id=eq.11111" [] [json|[{"id": 11111}]|]
+          `shouldRespondWith` ""
+          { matchStatus = 204
+          , matchHeaders = ["X-Custom-Header" <:> "mykey=myval"]
+          }
+
+      it "succeeds setting the headers on PUT" $
+        request methodPut "/items?id=eq.11111" [] [json|[{"id": 11111}]|]
+          `shouldRespondWith` ""
+          { matchStatus = 204
+          , matchHeaders = ["X-Custom-Header" <:> "mykey=myval"]
+          }
+
+      it "succeeds setting the headers on DELETE" $
+        request methodDelete "/items?id=eq.11111" [] mempty
+          `shouldRespondWith` ""
+          { matchStatus = 204
+          , matchHeaders = ["X-Custom-Header" <:> "mykey=myval"]
+          }
+
     context "Use of the phraseto_tsquery function" $ do
       it "finds matches" $
         get "/tsearch?text_search_vector=phfts.The%20Fat%20Cats" `shouldRespondWith`

test/Main.hs

@@ -78,6 +78,7 @@ main = do
       extraSearchPathApp   = return $ postgrest (testCfgExtraSearchPath testDbConn)   refDbStructure pool getTime $ pure ()
       rootSpecApp          = return $ postgrest (testCfgRootSpec testDbConn)          refDbStructure pool getTime $ pure ()
       htmlRawOutputApp     = return $ postgrest (testCfgHtmlRawOutput testDbConn)     refDbStructure pool getTime $ pure ()
+      responseHeadersApp   = return $ postgrest (testCfgResponseHeaders testDbConn)   refDbStructure pool getTime $ pure ()
 
   let reset, analyze :: IO ()
       reset = resetDb testDbConn
@@ -88,8 +89,7 @@ main = do
       actualPgVersion = pgVersion dbStructure
       extraSpecs =
         [("Feature.UpsertSpec", Feature.UpsertSpec.spec) | actualPgVersion >= pgVersion95] ++
-        [("Feature.PgVersion95Spec", Feature.PgVersion95Spec.spec) | actualPgVersion >= pgVersion95] ++
-        [("Feature.PgVersion96Spec", Feature.PgVersion96Spec.spec) | actualPgVersion >= pgVersion96]
+        [("Feature.PgVersion95Spec", Feature.PgVersion95Spec.spec) | actualPgVersion >= pgVersion95]
 
       specs = uncurry describe <$> [
           ("Feature.AuthSpec"                , Feature.AuthSpec.spec actualPgVersion)
@@ -165,6 +165,8 @@ main = do
       describe "Feature.ExtraSearchPathSpec" Feature.ExtraSearchPathSpec.spec
 
     -- this test runs with a root spec function override
-    when (actualPgVersion >= pgVersion96) $
+    when (actualPgVersion >= pgVersion96) $ do
       before rootSpecApp $
         describe "Feature.RootSpec" Feature.RootSpec.spec
+      before responseHeadersApp $
+        describe "Feature.PgVersion96Spec" Feature.PgVersion96Spec.spec

test/SpecHelper.hs

@@ -138,6 +138,9 @@ testCfgRootSpec testDbConn = (testCfg testDbConn) { configRootSpec = Just "root"
 testCfgHtmlRawOutput :: Text -> AppConfig
 testCfgHtmlRawOutput testDbConn = (testCfg testDbConn) { configRawMediaTypes = ["text/html"] }
 
+testCfgResponseHeaders :: Text -> AppConfig
+testCfgResponseHeaders testDbConn = (testCfg testDbConn) { configReqCheck = Just "custom_headers" }
+
 setupDb :: Text -> IO ()
 setupDb dbConn = do
   loadFixture dbConn "database"

test/fixtures/schema.sql

@@ -1642,3 +1642,22 @@ add constraint fst_shift foreign key           (fst_shift_activity_id, fst_shift
                          references activities (id, schedule_id),
 add constraint snd_shift foreign key           (snd_shift_activity_id, snd_shift_schedule_id)
                          references activities (id, schedule_id);
+
+-- for a pre-request function
+create or replace function custom_headers() returns void as $$
+declare
+  user_agent text := current_setting('request.header.user-agent', true);
+  req_path   text := current_setting('request.path', true);
+  req_accept text := current_setting('request.header.accept', true);
+begin
+  if user_agent similar to 'MSIE (6.0|7.0)' then
+    perform set_config('response.headers',
+      '[{"Cache-Control": "no-cache, no-store, must-revalidate"}]', false);
+  elsif req_path similar to '/(items|projects)' and req_accept = 'text/csv' then
+    perform set_config('response.headers',
+      format('[{"Content-Disposition": "attachment; filename=%s.csv"}]', trim('/' from req_path)), false);
+  else
+    perform set_config('response.headers',
+      '[{"X-Custom-Header": "mykey=myval"}]', false);
+  end if;
+end; $$ language plpgsql;