OverridableAs combinator to enable efficient implementations for some Servant endpoints after the fact.

servant-overrided-as.hs

#! /usr/bin/env nix-shell
#! nix-shell -i runghc -p "haskellPackages.ghcWithPackages (pkgs: with pkgs; [servant-swagger servant-swagger-ui])"
{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
module Main where

import           Control.Monad.Trans                        (liftIO)
import           Control.Monad.Trans.Resource               (runResourceT)
import           Data.Aeson                                 (FromJSON, ToJSON)
import qualified Data.Aeson                                 as Aeson
import qualified Data.ByteString.Lazy.Char8                 as BSL8
import           Data.Swagger                               (Swagger)
import qualified Data.Swagger                               as Swagger
import           Data.Text                                  (Text)
import           GHC.Generics                               (Generic)
import qualified Network.HTTP.Types                         as HTTP
import qualified Network.Wai                                as Wai
import qualified Network.Wai.Handler.Warp                   as Warp
import           Servant
import qualified Servant.Server.Internal.Router             as Servant
import qualified Servant.Server.Internal.RoutingApplication as Servant
import           Servant.Swagger
import           Servant.Swagger.UI

-- * API

type SampleAPI
    = "send" :> OverridableAs Raw SendItem
      -- ^ we want to be able to provide a more efficient
      -- Raw implementation for this endpoint
      -- specifically to avoid unnecessary ToJSON/FromJSON
      -- conversions and validations
 :<|> "list" :> ListItems

type SendItem
  = ReqBody '[JSON] Item              -- ^ An item to save.
 :> PostNoContent '[JSON] NoContent

type ListItems = Get '[JSON] Items

sampleAPI :: Proxy SampleAPI
sampleAPI = Proxy

-- * Model

-- | A sample Item data type that can be encoded/decoded as JSON.
data Item = Item
  { title       :: Text
  , description :: Text
  } deriving (Generic, ToJSON, FromJSON, Swagger.ToSchema)

-- | A bunch of 'Item's.
newtype Items = Items
  { items :: [Item]
  } deriving (Generic, ToJSON, FromJSON, Swagger.ToSchema)

-- * Server handlers

-- | A sample server with standard 'serveSendItem' implementation.
sampleServer :: Server SampleAPI
sampleServer
    = Overridable serveSendItem -- a standard Servant handler
 :<|> serveListItems

-- | Handle sent 'Item' by dumping its JSON encoding to stdout.
serveSendItem :: Item -> Handler NoContent
serveSendItem item = do
  liftIO $ BSL8.putStrLn (Aeson.encode item)
  return NoContent

-- | Serve some list of items.
serveListItems :: Handler Items
serveListItems = return $ Items
  [ Item { title = "Char", description = "Something to sit on" } ]

-- | Like 'sampleServer', but with 'efficientSendItem'.
efficientServer :: Server SampleAPI
efficientServer
    = Overriding efficientSendItem
 :<|> serveListItems

-- | An efficient implementation of SendItem API.
-- Here we bypass Servant's encoding/decoding of JSON
-- and merely dump request body to stdout.
efficientSendItem :: Server Raw
efficientSendItem = Tagged $ \req respond -> do
  body <- Wai.strictRequestBody req
  BSL8.putStrLn body
  respond $ Wai.responseLBS HTTP.status200 [] "Hello World"

sampleSwagger :: Swagger
sampleSwagger = toSwagger sampleAPI

-- | Complete API with 'SampleAPI' and Swagger documentation.
type API
    = SwaggerSchemaUI "swagger-ui" "swagger.json"
 :<|> SampleAPI

main :: IO ()
main = do
  putStrLn "Starting a server at http://localhost:8080"
  putStrLn "Swagger UI available at http://localhost:8080/swagger-ui"
  Warp.run 8080 $ serve (Proxy @API) $
    swaggerSchemaUIServer sampleSwagger :<|> efficientServer

-- * 'OverridableAs' combinator

-- | A value that can be overriden with a value of a different type.
data OverridableAs raw api
  = Overriding  raw
  | Overridable api

-- | For @'OverridableAs' raw api@ handler can implement
-- either handler for @raw@ or for @api@.
instance (HasServer raw ctx, HasServer api ctx)
  => HasServer (OverridableAs raw api) ctx where
  type ServerT (OverridableAs raw api) m
    = OverridableAs (ServerT raw m) (ServerT api m)

  -- FIXME: we can do better if we analyse routers for both raw and api
  route Proxy ctx app = Servant.RawRouter $ \ env request respond ->
    runResourceT $ do
      -- note: a Raw application doesn't register any cleanup
      -- but for the sake of consistency, we nonetheless run
      -- the cleanup once its done
      r <- Servant.runDelayed app env request
      liftIO $ go r request respond

    where go r request respond = case r of
            Servant.Route (Overriding raw)
              -> serveWithContext (Proxy @raw) ctx raw request (respond . Servant.Route)
            Servant.Route (Overridable api)
              -> serveWithContext (Proxy @api) ctx api request (respond . Servant.Route)
            Servant.Fail a    -> respond $ Servant.Fail a
            Servant.FailFatal e       -> respond $ Servant.FailFatal e

  hoistServerWithContext _ ctx phi (Overriding raw)
    = Overriding (hoistServerWithContext (Proxy @raw) ctx phi raw)
  hoistServerWithContext _ ctx phi (Overridable api)
    = Overridable (hoistServerWithContext (Proxy @api) ctx phi api)

-- | For @'OverridableAs' raw api@ we generate Swagger documentation
-- only for @api@.
instance HasSwagger api => HasSwagger (OverridableAs raw api) where
  toSwagger _ = toSwagger (Proxy @api)