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
 
 ### Fixed
 

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]
+
+This is analogous to the ``/ready`` endpoint in :ref:`admin_server`. However, instead of an HTTP response, this exits with 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,6 +13,7 @@ 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 Text.Heredoc (str)
@@ -24,23 +26,36 @@ 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 = Client.ready conf >>= \case
+  status | status >= HTTP.status200 && status < HTTP.status300 -> exitSuccess
+         | otherwise -> 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 +86,13 @@ data CLI = CLI
   }
 
 data Command
+  = Admin AdminCommand
+  | Run RunCommand
+
+data AdminCommand
+  = CmdReady
+
+data RunCommand
   = CmdRun
   | CmdDumpConfig
   | CmdDumpSchema
@@ -105,7 +127,7 @@ readCLIShowHelp =
     cliParser :: O.Parser CLI
     cliParser =
       CLI
-        <$> (dumpConfigFlag <|> dumpSchemaFlag)
+        <$> (dumpConfigFlag <|> dumpSchemaFlag <|> readyFlag)
         <*> O.optional configFileOption
 
     configFileOption =
@@ -114,15 +136,20 @@ 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 "Health check for if PostgREST is running and ready"
+
 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,37 @@
+{-# LANGUAGE NamedFieldPuns #-}
+module PostgREST.Client
+  ( ready
+  ) where
+
+import qualified Data.Text           as T
+import qualified Network.HTTP.Client as HC
+
+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 -> T.unpack (address port) <> "/ready"
+      Nothing   -> panic "Admin Server is not running. Please check your configuration."
+
+  resp <- HC.httpLbs req client     -- HTTP Response
+  return $ HC.responseStatus resp
+    where
+      address port = "http://"
+                   <> escapeHostName configAdminServerHost
+                   <> ":"
+                   <> show port
+
+      escapeHostName :: Text -> Text
+      escapeHostName "*"  = "0.0.0.0"
+      escapeHostName "*4" = "0.0.0.0"
+      escapeHostName "!4" = "0.0.0.0"
+      escapeHostName "*6" = "0.0.0.0"
+      escapeHostName "!6" = "0.0.0.0"
+      escapeHostName h    = h

test/io/test_cli.py

@@ -289,3 +289,78 @@ 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
+    command = [POSTGREST_BIN]
+    env["HPCTIXFILE"] = hpctixfile()
+    pgrst_process = subprocess.Popen(command, env=env)
+
+    time.sleep(5)  # Wait some seconds for the first process to start
+    try:
+        # Run healthcheck process
+        command = [POSTGREST_BIN] + ["--ready"]
+        healthcheck_process = subprocess.run(command, env=env, capture_output=True)
+        assert healthcheck_process.returncode == 0
+    finally:
+        pgrst_process.kill()
+        pgrst_process.wait()
+
+
+def test_cli_ready_flag_fail(defaultenv):
+    "PostgREST ready flag fails when not ready"
+
+    env = {
+        **defaultenv,
+        "PGRST_ADMIN_SERVER_PORT": "3001",
+    }
+
+    # Run postgrest process
+    command = [POSTGREST_BIN]
+    env["HPCTIXFILE"] = hpctixfile()
+    pgrst_process = subprocess.Popen(command, env=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
+    finally:
+        pgrst_process.kill()
+        pgrst_process.wait()
+
+
+def test_cli_ready_flag_fail_no_admin_server(defaultenv):
+    "PostgREST --ready flag fail without admin server running"
+
+    env = {
+        **defaultenv,
+        "PGRST_ADMIN_SERVER_PORT": "",
+    }
+
+    command = [POSTGREST_BIN]
+    env["HPCTIXFILE"] = hpctixfile()
+    pgrst_process = subprocess.Popen(command, env=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:
+        pgrst_process.kill()
+        pgrst_process.wait()