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>

Author: taimoorzaeem

CHANGELOG.md

@@ -12,6 +12,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
   + The exposed schemas are now listed in the `hint` instead of the `message` field.
 - Improve error details of `PGRST301` error by @taimoorzaeem in #4051
 - Bounded JWT cache using the SIEVE algorithm by @mkleczek in #4084
+- Add `--ready` flag for postgrest healthcheck by @taimoorzaeem in #4239
 
 ### Changed
 

docs/references/cli.rst

@@ -49,3 +49,12 @@ Dump Schema
   $ postgrest [--dump-schema]
 
 Dumps the schema cache in JSON format.
+
+Ready Flag
+----------
+
+.. code:: bash
+
+  $ postgrest [--ready]
+
+Makes a request to the ``/ready`` endpoint of the :ref:`admin_server`. It exits with a return code of ``0`` on success and ``1`` on failure.

postgrest.cabal

@@ -52,6 +52,7 @@ library
                       PostgREST.Auth.Types
                       PostgREST.Cache.Sieve
                       PostgREST.CLI
+                      PostgREST.Client
                       PostgREST.Config
                       PostgREST.Config.Database
                       PostgREST.Config.JSPath
@@ -117,6 +118,7 @@ library
                     , hasql-pool                >= 1.0.1 && < 1.1
                     , hasql-transaction         >= 1.0.1 && < 1.2
                     , heredoc                   >= 0.2 && < 0.3
+                    , http-client               >= 0.7.19 && < 0.8
                     , http-types                >= 0.12.2 && < 0.13
                     , insert-ordered-containers >= 0.2.2 && < 0.3
                     , iproute                   >= 1.7.0 && < 1.8

src/PostgREST/CLI.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase      #-}
 {-# LANGUAGE NamedFieldPuns  #-}
 {-# LANGUAGE QuasiQuotes     #-}
 {-# LANGUAGE RecordWildCards #-}
@@ -12,8 +13,10 @@ import qualified Data.Aeson                 as JSON
 import qualified Data.ByteString.Char8      as BS
 import qualified Data.ByteString.Lazy       as LBS
 import qualified Hasql.Transaction.Sessions as SQL
+import qualified Network.HTTP.Types.Status  as HTTP
 import qualified Options.Applicative        as O
 
+import Data.Maybe   (fromJust)
 import Text.Heredoc (str)
 
 import PostgREST.AppState    (AppState)
@@ -24,23 +27,41 @@ 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
 
 
 main :: CLI -> IO ()
 main CLI{cliCommand, cliPath} = do
-  conf@AppConfig{..} <-
+  conf <-
     either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty
+  case cliCommand of
+    Admin adminCmd -> runAdminCommand conf adminCmd
+    Run runCmd     -> runAppCommand conf runCmd
 
+-- | Run command using http-client to communicate with an already running postgrest
+runAdminCommand :: AppConfig -> AdminCommand -> IO ()
+runAdminCommand conf CmdReady = do
+  status <- Client.ready conf
+  if status >= HTTP.status200 && status < HTTP.status300
+    then do
+      putStrLn $ BS.pack $ "OK: http://:" <> show (fromJust $ configAdminServerPort conf) <> "/ready"
+      exitSuccess
+    else
+      exitWith $ ExitFailure 1
+
+-- | 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 +92,13 @@ data CLI = CLI
   }
 
 data Command
+  = Admin AdminCommand
+  | Run RunCommand
+
+data AdminCommand
+  = CmdReady
+
+data RunCommand
   = CmdRun
   | CmdDumpConfig
   | CmdDumpSchema
@@ -105,7 +133,7 @@ readCLIShowHelp =
     cliParser :: O.Parser CLI
     cliParser =
       CLI
-        <$> (dumpConfigFlag <|> dumpSchemaFlag)
+        <$> (dumpConfigFlag <|> dumpSchemaFlag <|> readyFlag)
         <*> O.optional configFileOption
 
     configFileOption =
@@ -114,15 +142,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) (Admin 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,47 @@
+{-|
+Module      : PostgREST.Client
+Description : PostgREST HTTP client
+-}
+{-# LANGUAGE LambdaCase          #-}
+{-# LANGUAGE NamedFieldPuns      #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module PostgREST.Client
+  ( ready
+  ) where
+
+import qualified Data.Text           as T
+import qualified Network.HTTP.Client as HC
+
+import Network.HTTP.Client       (HttpException (..))
+import Network.HTTP.Types.Status (Status (..))
+import PostgREST.Config          (AppConfig (..))
+
+import Protolude
+
+ready :: AppConfig -> IO Status
+ready AppConfig{configAdminServerHost, configAdminServerPort} = do
+
+  client <- HC.newManager HC.defaultManagerSettings
+  req <- HC.parseRequest $ -- Create HTTP Request
+    case configAdminServerPort of
+      Just port -> "http://" <> host <> ":" <> show port <> "/ready"
+      Nothing   -> panic "Admin Server is not running. Please check your configuration."
+
+  resp <- HC.httpLbs req client `catch` \(_ :: HttpException) -> do
+    let url = show (HC.getUri req)
+    hPutStrLn stderr $ T.pack $ "postgrest: health check request failed - connection refused to " <> url
+    exitWith $ ExitFailure 1
+
+  return $ HC.responseStatus resp
+    where
+      host = escapeHostName $ T.unpack configAdminServerHost
+      -- We don't need to resolve the host name except the postgrest special
+      -- addresses. The http-client package automatically resolves the other
+      -- host names for us.
+      escapeHostName = \case
+        "*"  -> "0.0.0.0"
+        "*4" -> "0.0.0.0"
+        "!4" -> "0.0.0.0"
+        "*6" -> "[::1]" -- TODO: Not covered by tests
+        "!6" -> "[::1]" -- TODO: Not covered by tests
+        h    -> h

test/io/postgrest.py

@@ -230,3 +230,19 @@ def is_ipv6(addr):
         return True
     except OSError:
         return False
+
+
+def run_postgrest_binary(env):
+    "Runs a the postgrest using the binary and return the running process"
+
+    command = [POSTGREST_BIN]
+    env["HPCTIXFILE"] = hpctixfile()
+    pgrst_process = subprocess.Popen(command, env=env)
+    return pgrst_process
+
+
+def kill_pgrst_process(process):
+    "Kill the given postgrest process"
+
+    process.kill()  # Send SIGKILL
+    process.wait()  # Wait for termination

test/io/test_cli.py

@@ -24,6 +24,7 @@
 import yaml
 
 from config import *
+from postgrest import *
 
 
 class ExtraNewLinesDumper(yaml.SafeDumper):
@@ -289,3 +290,75 @@ def test_jwt_aud_config_set_to_invalid_uri(defaultenv):
     with pytest.raises(PostgrestError):
         dump = cli(["--dump-config"], env=env).split("\n")
         assert "jwt-aud should be a string or a valid URI" in dump
+
+
+def test_cli_ready_flag_success(defaultenv):
+    "PostgREST ready flag succeeds when ready"
+
+    env = {
+        **defaultenv,
+        "PGRST_ADMIN_SERVER_PORT": "3001",
+    }
+
+    # Run postgrest process
+    pgrst_process = run_postgrest_binary(env)
+
+    time.sleep(3)  # Wait so postgrest gets ready
+    try:
+        # Run healthcheck process
+        command = [POSTGREST_BIN] + ["--ready"]
+        healthcheck_process = subprocess.run(command, env=env, capture_output=True)
+        assert healthcheck_process.returncode == 0
+        assert f"OK: http://:3001/ready" in healthcheck_process.stdout.decode()
+    finally:
+        kill_pgrst_process(pgrst_process)
+
+
+def test_cli_ready_flag_fail(defaultenv):
+    "PostgREST ready flag fails when not ready"
+
+    env = {
+        **defaultenv,
+        "PGRST_ADMIN_SERVER_PORT": "3001",
+    }
+
+    # Run postgrest process
+    pgrst_process = run_postgrest_binary(env)
+
+    time.sleep(0)  # No waiting and instantly run the healthcheck process.
+    try:
+        # Run healthcheck process
+        command = [POSTGREST_BIN] + ["--ready"]
+        healthcheck_process = subprocess.run(command, env=env, capture_output=True)
+        assert healthcheck_process.returncode == 1
+        assert (
+            "postgrest: health check request failed - connection refused to"
+            in healthcheck_process.stderr.decode()
+        )
+    finally:
+        kill_pgrst_process(pgrst_process)
+
+
+def test_cli_ready_flag_fail_no_admin_server(defaultenv):
+    "PostgREST --ready flag fail without admin server running"
+
+    env = {
+        **defaultenv,
+        "PGRST_ADMIN_SERVER_PORT": "",
+    }
+
+    # Run postgrest process
+    pgrst_process = run_postgrest_binary(env)
+
+    time.sleep(3)  # Wait so postgrest gets ready
+    try:
+        # Run healthcheck process
+        command = [POSTGREST_BIN] + ["--ready"]
+        healthcheck_process = subprocess.run(command, env=env, capture_output=True)
+        assert healthcheck_process.returncode == 1
+        assert (
+            "Admin Server is not running. Please check your configuration."
+            in healthcheck_process.stderr.decode()
+        )
+    finally:
+        kill_pgrst_process(pgrst_process)