add: --ready flag for postgrest healthcheck

The `--ready` flag is a wrapper around the admin server
`/ready` request. This is done through using an http client
library in postgrest.

Signed-off-by: Taimoor Zaeem <taimoorzaeem@gmail.com>

src/PostgREST/CLI.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase      #-}
 {-# LANGUAGE NamedFieldPuns  #-}
 {-# LANGUAGE QuasiQuotes     #-}
 {-# LANGUAGE RecordWildCards #-}
@@ -24,6 +25,7 @@ import PostgREST.Version     (prettyVersion)
 
 import qualified PostgREST.App      as App
 import qualified PostgREST.AppState as AppState
+import qualified PostgREST.Client   as Client
 import qualified PostgREST.Config   as Config
 
 import Protolude
@@ -31,16 +33,26 @@ import Protolude
 
 main :: CLI -> IO ()
 main CLI{cliCommand, cliPath} = do
-  conf@AppConfig{..} <-
+  conf <-
     either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty
+  case cliCommand of
+    Client adminCmd -> runClientCommand conf adminCmd
+    Run runCmd      -> runAppCommand conf runCmd
 
+-- | Run command using http-client to communicate with an already running postgrest
+runClientCommand :: AppConfig -> ClientCommand -> IO ()
+runClientCommand conf CmdReady = Client.ready conf
+
+-- | Run postgrest with command
+runAppCommand :: AppConfig -> RunCommand -> IO ()
+runAppCommand conf@AppConfig{..} runCmd = do
   -- Per https://github.com/PostgREST/postgrest/issues/268, we want to
   -- explicitly close the connections to PostgreSQL on shutdown.
   -- 'AppState.destroy' takes care of that.
   bracket
     (AppState.init conf)
     AppState.destroy
-    (\appState -> case cliCommand of
+    (\appState -> case runCmd of
       CmdDumpConfig -> do
         when configDbConfig $ AppState.readInDbConfig True appState
         putStr . Config.toText =<< AppState.getConfig appState
@@ -71,6 +83,13 @@ data CLI = CLI
   }
 
 data Command
+  = Client ClientCommand
+  | Run RunCommand
+
+data ClientCommand
+  = CmdReady
+
+data RunCommand
   = CmdRun
   | CmdDumpConfig
   | CmdDumpSchema
@@ -105,7 +124,7 @@ readCLIShowHelp =
     cliParser :: O.Parser CLI
     cliParser =
       CLI
-        <$> (dumpConfigFlag <|> dumpSchemaFlag)
+        <$> (dumpConfigFlag <|> dumpSchemaFlag <|> readyFlag)
         <*> O.optional configFileOption
 
     configFileOption =
@@ -114,15 +133,21 @@ readCLIShowHelp =
         <> O.help "Path to configuration file"
 
     dumpConfigFlag =
-      O.flag CmdRun CmdDumpConfig $
+      O.flag (Run CmdRun) (Run CmdDumpConfig) $
         O.long "dump-config"
         <> O.help "Dump loaded configuration and exit"
 
     dumpSchemaFlag =
-      O.flag CmdRun CmdDumpSchema $
+      O.flag (Run CmdRun) (Run CmdDumpSchema) $
         O.long "dump-schema"
         <> O.help "Dump loaded schema as JSON and exit (for debugging, output structure is unstable)"
 
+    readyFlag =
+      O.flag (Run CmdRun) (Client CmdReady) $
+        O.long "ready"
+        <> O.help "Checks the health of PostgREST by doing a request on the admin server /ready endpoint"
+
+
 exampleConfigFile :: [Char]
 exampleConfigFile =
   [str|## Admin server used for checks. It's disabled by default unless a port is specified.

src/PostgREST/Client.hs

@@ -0,0 +1,100 @@
+{-|
+Module      : PostgREST.Client
+Description : PostgREST HTTP client
+-}
+{-# LANGUAGE NamedFieldPuns #-}
+module PostgREST.Client
+  ( ready
+  ) where
+
+import qualified Data.Text                 as T
+import qualified Network.HTTP.Client       as HC
+import qualified Network.HTTP.Types.Status as HTTP
+
+import Network.HTTP.Client (HttpException (..))
+import System.IO           (hFlush)
+
+import PostgREST.Config  (AppConfig (..))
+import PostgREST.Network (isSpecialHostName)
+
+import Protolude
+
+data PgrstClientError
+  = NoAdminServer
+  | NoSpecialHostNamesAllowed Text
+  | PostgRESTNotReady         Text
+  | HTTPConnectionRefused     Text
+  | HTTPExceptionInvalidURL   Text
+
+-- | This is invoked by the CLI "--ready" flag.
+--   The http-client sends and a request to /ready endpoint
+--   and exits with success or failure.
+ready :: AppConfig -> IO ()
+ready AppConfig{configAdminServerHost, configAdminServerPort} = do
+
+  client <- HC.newManager HC.defaultManagerSettings
+  readyURL <- getURL
+  req <- HC.parseRequest (T.unpack readyURL) `catch` handleHttpException
+  resp <- HC.httpLbs req client `catch` handleHttpException
+
+  let status = HC.responseStatus resp
+
+  if status >= HTTP.status200 && status < HTTP.status300
+    then printAndExitWithSuccess $ "OK: " <> readyURL
+    else printAndExitWithFailure $ clientErrorMsg (PostgRESTNotReady readyURL)
+        where
+          getURL :: IO Text
+          getURL =
+            -- Here, we have three cases:
+            -- 1. If the admin port config is not defined, we exit
+            --      with "no admin server error"
+            -- 2. Otherwise, if admin server is running, then we check if
+            --      postgrest server-host is configured with special hostname like "*4",
+            --      if it is, we fail with "no special hostname allowed with "--ready".
+            --      The reason for this is that we can't know the actual address.
+            -- 3. Finally, if we know the "actual" hostname and the port, then we
+            --      construct the URL and return it.
+            case configAdminServerPort of
+              Nothing   -> printAndExitWithFailure $ clientErrorMsg NoAdminServer
+              Just port ->
+                if isSpecialHostName configAdminServerHost
+                  then printAndExitWithFailure $ clientErrorMsg (NoSpecialHostNamesAllowed configAdminServerHost)
+                  else return $ makeReadyUrl port
+
+          -- NOTE: http-client automatically resolves hostnames
+          makeReadyUrl :: Int -> Text
+          makeReadyUrl p = "http://" <> wrapIfIpv6 configAdminServerHost <> ":" <> (T.pack . show) p <> "/ready"
+            where
+              -- IPv6 needs to wrapped in [], it has ':' as separator
+              wrapIfIpv6 :: Text -> Text
+              wrapIfIpv6 s
+                | T.any (== ':') s = "[" <> s <> "]"
+                | otherwise = s
+
+-- | Handle HTTP exception for "http-client" requests
+handleHttpException :: HttpException -> IO a
+handleHttpException (HttpExceptionRequest req _) = do
+  let url = show (HC.getUri req)
+  printAndExitWithFailure $ clientErrorMsg (HTTPConnectionRefused $ T.pack url)
+handleHttpException (InvalidUrlException url _) = do
+  printAndExitWithFailure $ clientErrorMsg (HTTPExceptionInvalidURL $ T.pack url)
+
+-- | Print the message on stdout and exit with success
+printAndExitWithSuccess :: Text -> IO a
+printAndExitWithSuccess msg = putStrLn (T.unpack msg) >> hFlush stdout >> exitSuccess
+
+-- | Print the message on stderr and exit with failure
+printAndExitWithFailure :: Text -> IO a
+printAndExitWithFailure msg = hPutStrLn stderr (T.unpack msg) >> hFlush stderr >> exitWith (ExitFailure 1)
+
+-- | Pgrst client error to error message
+clientErrorMsg :: PgrstClientError -> Text
+clientErrorMsg err = "ERROR: " <>
+  case err of
+    NoAdminServer -> "Admin server is not running. Please check admin-server-port config."
+    NoSpecialHostNamesAllowed host ->
+      "The `--ready` flag cannot be used when server-host is configured as \"" <> host <> "\". "
+      <> "Please update your server-host config to \"localhost\"."
+    PostgRESTNotReady url -> url
+    HTTPConnectionRefused url -> "connection refused to " <> url
+    HTTPExceptionInvalidURL url -> "invalid url - " <> url

src/PostgREST/Network.hs

@@ -1,6 +1,7 @@
 module PostgREST.Network
   ( resolveSocketToAddress
   , escapeHostName
+  , isSpecialHostName
   ) where
 
 import           Data.String    (IsString (..))
@@ -49,3 +50,12 @@ escapeHostName "!4" = "0.0.0.0"
 escapeHostName "*6" = "0.0.0.0"
 escapeHostName "!6" = "0.0.0.0"
 escapeHostName h    = h
+
+-- | Check if a hostname is special
+isSpecialHostName :: Text -> Bool
+isSpecialHostName "*"  = True
+isSpecialHostName "*4" = True
+isSpecialHostName "!4" = True
+isSpecialHostName "*6" = True
+isSpecialHostName "!6" = True
+isSpecialHostName _    = False