Add response header support to UVerb (#1420)

* Use type wrapped in Headers h to generate response

This avoids having to define MimeRender instances for Headers.

servant-client-core/src/Servant/Client/Core/HasClient.hs

@@ -75,7 +75,7 @@ import           Servant.API
                  NoContentVerb, QueryFlag, QueryParam', QueryParams, Raw,
                  ReflectMethod (..), RemoteHost, ReqBody', SBoolI, Stream,
                  StreamBody', Summary, ToHttpApiData, ToSourceIO (..), Vault,
-                 Verb, WithNamedContext, contentType, getHeadersHList,
+                 Verb, WithNamedContext, WithStatus (..), contentType, getHeadersHList,
                  getResponse, toQueryParam, toUrlPiece)
 import           Servant.API.ContentTypes
                  (contentTypes, AllMime (allMime), AllMimeUnrender (allMimeUnrender))
@@ -318,6 +318,25 @@ instance {-# OVERLAPPING #-}
 data ClientParseError = ClientParseError MediaType String | ClientStatusMismatch | ClientNoMatchingStatus
   deriving (Eq, Show)
 
+class UnrenderResponse (cts :: [*]) (a :: *) where
+  unrenderResponse :: Seq.Seq H.Header -> BL.ByteString -> Proxy cts
+                   -> [Either (MediaType, String) a]
+
+instance {-# OVERLAPPABLE #-} AllMimeUnrender cts a => UnrenderResponse cts a where
+  unrenderResponse _ body = map parse . allMimeUnrender
+    where parse (mediaType, parser) = left ((,) mediaType) (parser body)
+
+instance {-# OVERLAPPING #-} forall cts a h . (UnrenderResponse cts a, BuildHeadersTo h)
+  => UnrenderResponse cts (Headers h a) where
+  unrenderResponse hs body = (map . fmap) setHeaders . unrenderResponse hs body
+    where
+      setHeaders :: a -> Headers h a
+      setHeaders x = Headers x (buildHeadersTo (toList hs))
+
+instance {-# OVERLAPPING #-} UnrenderResponse cts a
+  => UnrenderResponse cts (WithStatus n a) where
+  unrenderResponse hs body = (map . fmap) WithStatus . unrenderResponse hs body
+
 instance {-# OVERLAPPING #-}
   ( RunClient m,
     contentTypes ~ (contentType ': otherContentTypes),
@@ -326,7 +345,7 @@ instance {-# OVERLAPPING #-}
     as ~ (a ': as'),
     AllMime contentTypes,
     ReflectMethod method,
-    All (AllMimeUnrender contentTypes) as,
+    All (UnrenderResponse contentTypes) as,
     All HasStatus as, HasStatuses as',
     Unique (Statuses as)
   ) =>
@@ -349,7 +368,8 @@ instance {-# OVERLAPPING #-}
 
     let status = responseStatusCode response
         body = responseBody response
-        res = tryParsers status $ mimeUnrenders (Proxy @contentTypes) body
+        headers = responseHeaders response
+        res = tryParsers status $ mimeUnrenders (Proxy @contentTypes) headers body
     case res of
       Left errors -> throwClientError $ DecodeFailure (T.pack (show errors)) response
       Right x -> return x
@@ -370,13 +390,14 @@ instance {-# OVERLAPPING #-}
       -- | Given a list of types, parses the given response body as each type
       mimeUnrenders ::
         forall cts xs.
-        All (AllMimeUnrender cts) xs =>
+        All (UnrenderResponse cts) xs =>
         Proxy cts ->
+        Seq.Seq H.Header ->
         BL.ByteString ->
         NP ([] :.: Either (MediaType, String)) xs
-      mimeUnrenders ctp body = cpure_NP
+      mimeUnrenders ctp headers body = cpure_NP
-        (Proxy @(AllMimeUnrender cts))
+        (Proxy @(UnrenderResponse cts))
-        (Comp . map (\(mediaType, parser) -> left ((,) mediaType) (parser body)) . allMimeUnrender $ ctp)
+        (Comp . unrenderResponse headers body $ ctp)
 
   hoistClientMonad _ _ nt s = nt s
 

servant-client/test/Servant/ClientTestUtils.hs

@@ -32,6 +32,7 @@ import           Data.Char
                  (chr, isPrint)
 import           Data.Monoid ()
 import           Data.Proxy
+import           Data.SOP
 import           Data.Text
                  (Text)
 import qualified Data.Text                        as Text
@@ -121,6 +122,7 @@ type Api =
             ReqBody '[JSON] [(String, [Rational])] :>
             Get '[JSON] (String, Maybe Int, Bool, [(String, [Rational])])
   :<|> "headers" :> Get '[JSON] (Headers TestHeaders Bool)
+  :<|> "uverb-headers" :> UVerb 'GET '[JSON] '[ WithStatus 200 (Headers TestHeaders Bool), WithStatus 204 String ]
   :<|> "deleteContentType" :> DeleteNoContent
   :<|> "redirectWithCookie" :> Raw
   :<|> "empty" :> EmptyAPI
@@ -150,6 +152,7 @@ getRawFailure   :: HTTP.Method -> ClientM Response
 getMultiple     :: String -> Maybe Int -> Bool -> [(String, [Rational])]
   -> ClientM (String, Maybe Int, Bool, [(String, [Rational])])
 getRespHeaders  :: ClientM (Headers TestHeaders Bool)
+getUVerbRespHeaders  :: ClientM (Union '[ WithStatus 200 (Headers TestHeaders Bool), WithStatus 204 String ])
 getDeleteContentType :: ClientM NoContent
 getRedirectWithCookie :: HTTP.Method -> ClientM Response
 uverbGetSuccessOrRedirect :: Bool
@@ -172,6 +175,7 @@ getRoot
   :<|> getRawFailure
   :<|> getMultiple
   :<|> getRespHeaders
+  :<|> getUVerbRespHeaders
   :<|> getDeleteContentType
   :<|> getRedirectWithCookie
   :<|> EmptyClient
@@ -198,6 +202,7 @@ server = serve api (
   :<|> (Tagged $ \ _request respond -> respond $ Wai.responseLBS HTTP.badRequest400 [] "rawFailure")
   :<|> (\ a b c d -> return (a, b, c, d))
   :<|> (return $ addHeader 1729 $ addHeader "eg2" True)
+  :<|> (pure . Z . I . WithStatus $ addHeader 1729 $ addHeader "eg2" True)
   :<|> return NoContent
   :<|> (Tagged $ \ _request respond -> respond $ Wai.responseLBS HTTP.found302 [("Location", "testlocation"), ("Set-Cookie", "testcookie=test")] "")
   :<|> emptyServer

servant-client/test/Servant/SuccessSpec.hs

@@ -32,6 +32,7 @@ import           Data.Foldable
 import           Data.Maybe
                  (listToMaybe)
 import           Data.Monoid ()
+import           Data.SOP (NS (..), I (..))
 import           Data.Text
                  (Text)
 import qualified Network.HTTP.Client                as C
@@ -129,6 +130,14 @@ successSpec = beforeAll (startWaiApp server) $ afterAll endWaiApp $ do
         Left e -> assertFailure $ show e
         Right val -> getHeaders val `shouldBe` [("X-Example1", "1729"), ("X-Example2", "eg2")]
 
+    it "Returns headers on UVerb requests" $ \(_, baseUrl) -> do
+      res <- runClient getUVerbRespHeaders baseUrl
+      case res of
+        Left e -> assertFailure $ show e
+        Right (Z (I (WithStatus val))) ->
+          getHeaders val `shouldBe` [("X-Example1", "1729"), ("X-Example2", "eg2")]
+        Right (S _) -> assertFailure "expected first alternative of union"
+
     it "Stores Cookie in CookieJar after a redirect" $ \(_, baseUrl) -> do
       mgr <- C.newManager C.defaultManagerSettings
       cj <- atomically . newTVar $ C.createCookieJar []

servant-server/servant-server.cabal

@@ -147,6 +147,7 @@ test-suite spec
     , safe
     , servant
     , servant-server
+    , sop-core
     , string-conversions
     , text
     , transformers

servant-server/src/Servant/Server/UVerb.hs

@@ -22,15 +22,17 @@ module Servant.Server.UVerb
   )
 where
 
+import qualified Data.ByteString as B
 import Data.Proxy (Proxy (Proxy))
 import Data.SOP (I (I))
 import Data.SOP.Constraint (All, And)
 import Data.String.Conversions (LBS, cs)
-import Network.HTTP.Types (Status, hContentType)
+import Network.HTTP.Types (Status, HeaderName, hContentType)
-import Network.Wai (responseLBS)
+import Network.Wai (responseLBS, Request)
 import Servant.API (ReflectMethod, reflectMethod)
 import Servant.API.ContentTypes (AllCTRender (handleAcceptH), AllMime)
-import Servant.API.UVerb (HasStatus, IsMember, Statuses, UVerb, Union, Unique, foldMapUnion, inject, statusOf)
+import Servant.API.ResponseHeaders (GetHeaders (..), Headers (..))
+import Servant.API.UVerb (HasStatus, IsMember, Statuses, UVerb, Union, Unique, WithStatus (..), foldMapUnion, inject, statusOf)
 import Servant.Server.Internal (Context, Delayed, Handler, HasServer (..), RouteResult (FailFatal, Route), Router, Server, ServerT, acceptCheck, addAcceptCheck, addMethodCheck, allowedMethodHead, err406, getAcceptHeader, leafRouter, methodCheck, runAction)
 
 
@@ -43,13 +45,38 @@ respond ::
   f (Union xs)
 respond = pure . inject . I
 
--- | Helper constraint used in @instance 'HasServer' 'UVerb'@.
+class IsServerResource (cts :: [*]) a where
-type IsServerResource contentTypes = AllCTRender contentTypes `And` HasStatus
+  resourceResponse :: Request -> Proxy cts -> a -> Maybe (LBS, LBS)
+  resourceHeaders :: Proxy cts -> a -> [(HeaderName, B.ByteString)]
+
+instance {-# OVERLAPPABLE #-} AllCTRender cts a
+  => IsServerResource cts a where
+  resourceResponse request p res = handleAcceptH p (getAcceptHeader request) res
+  resourceHeaders _ _ = []
+
+instance {-# OVERLAPPING #-} (IsServerResource cts a, GetHeaders (Headers h a))
+  => IsServerResource cts (Headers h a) where
+  resourceResponse request p res = resourceResponse request p (getResponse res)
+  resourceHeaders cts res = getHeaders res ++ resourceHeaders cts (getResponse res)
+
+instance {-# OVERLAPPING #-} IsServerResource cts a
+  => IsServerResource cts (WithStatus n a) where
+  resourceResponse request p (WithStatus x) = resourceResponse request p x
+  resourceHeaders cts (WithStatus x) = resourceHeaders cts x
+
+encodeResource :: forall a cts . (IsServerResource cts a, HasStatus a)
+               => Request -> Proxy cts -> a
+               -> (Status, Maybe (LBS, LBS), [(HeaderName, B.ByteString)])
+encodeResource request cts res = (statusOf (Proxy @a),
+                                  resourceResponse request cts res,
+                                  resourceHeaders cts res)
+
+type IsServerResourceWithStatus cts = IsServerResource cts `And` HasStatus
 
 instance
   ( ReflectMethod method,
     AllMime contentTypes,
-    All (IsServerResource contentTypes) as,
+    All (IsServerResourceWithStatus contentTypes) as,
     Unique (Statuses as) -- for consistency with servant-swagger (server would work fine
         -- without; client is a bit of a corner case, because it dispatches
         -- the parser based on the status code.  with this uniqueness
@@ -77,20 +104,13 @@ instance
               action
                 `addMethodCheck` methodCheck method request
                 `addAcceptCheck` acceptCheck (Proxy @contentTypes) (getAcceptHeader request)
-            mkProxy :: a -> Proxy a
-            mkProxy _ = Proxy
 
         runAction action' env request cont $ \(output :: Union as) -> do
-          let encodeResource :: (AllCTRender contentTypes a, HasStatus a) => a -> (Status, Maybe (LBS, LBS))
+          let cts = Proxy @contentTypes
-              encodeResource res =
+              pickResource :: Union as -> (Status, Maybe (LBS, LBS), [(HeaderName, B.ByteString)])
-                ( statusOf $ mkProxy res,
+              pickResource = foldMapUnion (Proxy @(IsServerResourceWithStatus contentTypes)) (encodeResource request cts)
-                  handleAcceptH (Proxy @contentTypes) (getAcceptHeader request) res
-                )
-              pickResource :: Union as -> (Status, Maybe (LBS, LBS))
-              pickResource = foldMapUnion (Proxy @(IsServerResource contentTypes)) encodeResource
-
           case pickResource output of
-            (_, Nothing) -> FailFatal err406 -- this should not happen (checked before), so we make it fatal if it does
+            (_, Nothing, _) -> FailFatal err406 -- this should not happen (checked before), so we make it fatal if it does
-            (status, Just (contentT, body)) ->
+            (status, Just (contentT, body), headers) ->
               let bdy = if allowedMethodHead method request then "" else body
-               in Route $ responseLBS status ((hContentType, cs contentT) : []) bdy
+               in Route $ responseLBS status ((hContentType, cs contentT) : headers) bdy

servant-server/test/Servant/ServerSpec.hs

@@ -28,6 +28,8 @@ import           Data.Maybe
                  (fromMaybe)
 import           Data.Proxy
                  (Proxy (Proxy))
+import           Data.SOP
+                 (I (..), NS (..))
 import           Data.String
                  (fromString)
 import           Data.String.Conversions
@@ -53,7 +55,7 @@ import           Servant.API
                  NoContent (..), NoContentVerb, NoFraming, OctetStream, Patch,
                  PlainText, Post, Put, QueryFlag, QueryParam, QueryParams, Raw,
                  RemoteHost, ReqBody, SourceIO, StdMethod (..), Stream, Strict,
-                 UVerb, Union, Verb, addHeader)
+                 UVerb, Union, Verb, WithStatus (..), addHeader)
 import           Servant.Server
                  (Context ((:.), EmptyContext), Handler, Server, Tagged (..),
                  emptyServer, err401, err403, err404, respond, serve,
@@ -98,6 +100,7 @@ spec = do
   rawSpec
   alternativeSpec
   responseHeadersSpec
+  uverbResponseHeadersSpec
   miscCombinatorSpec
   basicAuthSpec
   genAuthSpec
@@ -684,6 +687,31 @@ responseHeadersSpec = describe "ResponseHeaders" $ do
         THW.request method "" [(hAccept, "crazy/mime")] ""
           `shouldRespondWith` 406
 
+-- }}}
+------------------------------------------------------------------------------
+-- * uverbResponseHeaderSpec {{{
+------------------------------------------------------------------------------
+type UVerbHeaderResponse = '[
+  WithStatus 200 (Headers '[Header "H1" Int] String),
+  WithStatus 404 String ]
+
+type UVerbResponseHeadersApi =
+       Capture "ok" Bool :> UVerb 'GET '[JSON] UVerbHeaderResponse
+
+uverbResponseHeadersServer :: Server UVerbResponseHeadersApi
+uverbResponseHeadersServer True = pure . Z . I . WithStatus $ addHeader 5 "foo"
+uverbResponseHeadersServer False = pure . S . Z . I . WithStatus $ "bar"
+
+uverbResponseHeadersSpec :: Spec
+uverbResponseHeadersSpec = describe "UVerbResponseHeaders" $ do
+  with (return $ serve (Proxy :: Proxy UVerbResponseHeadersApi) uverbResponseHeadersServer) $ do
+
+    it "includes the headers in the response" $
+        THW.request methodGet "/true" [] ""
+          `shouldRespondWith` "\"foo\"" { matchHeaders = ["H1" <:> "5"]
+                                        , matchStatus  = 200
+                                        }
+
 -- }}}
 ------------------------------------------------------------------------------
 -- * miscCombinatorSpec {{{

servant/src/Servant/API/ResponseHeaders.hs

@@ -51,6 +51,9 @@ import           Web.HttpApiData
 
 import           Prelude ()
 import           Prelude.Compat
+import           Servant.API.ContentTypes
+                 (JSON, PlainText, FormUrlEncoded, OctetStream,
+                  MimeRender(..))
 import           Servant.API.Header
                  (Header)