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>

CHANGELOG.md

@@ -14,6 +14,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Added
 
 - Bounded JWT cache using the SIEVE algorithm by @mkleczek in #4084
+- Add `--ready` flag for postgrest healthcheck by @taimoorzaeem in #4239
 
 ### Changed
 

docs/explanations/install.rst

@@ -142,6 +142,7 @@ To avoid having to install the database at all, you can run both it and the serv
       ports:
         - "3000:3000"
       environment:
+        PGRST_SERVER_HOST: localhost # necessary for `postgrest --ready` flag to work
         PGRST_DB_URI: postgres://app_user:password@db:5432/app_db
         PGRST_OPENAPI_SERVER_PROXY_URI: http://127.0.0.1:3000
       depends_on:

docs/references/cli.rst

@@ -7,7 +7,7 @@ PostgREST provides a CLI with the options listed below:
 
 .. code:: text
 
-  Usage: postgrest [-v|--version] [-e|--example] [--dump-config | --dump-schema] 
+  Usage: postgrest [-v|--version] [-e|--example] [--dump-config | --dump-schema | --ready]
                  [FILENAME]
 
     PostgREST / create a REST API to an existing Postgres
@@ -20,6 +20,8 @@ PostgREST provides a CLI with the options listed below:
     --dump-config            Dump loaded configuration and exit
     --dump-schema            Dump loaded schema as JSON and exit (for debugging,
                              output structure is unstable)
+    --ready                  Checks the health of PostgREST by doing a request on
+                             the admin server /ready endpoint
     FILENAME                 Path to configuration file
 
 FILENAME
@@ -71,3 +73,17 @@ Dump Schema
   $ postgrest --dump-schema
 
 Dumps the schema cache in JSON format.
+
+Ready Flag
+----------
+
+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.
+
+.. code-block:: bash
+
+  $ postgrest --ready
+  OK: http://localhost:3001/ready
+
+.. note::
+
+  The ``--ready`` flag cannot be used when :ref:`server-host` is configured with special hostnames. We suggest to change it to ``localhost``.

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
@@ -118,6 +119,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 #-}
@@ -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

test/io/config.py

@@ -95,3 +95,9 @@ def hpctixfile():
     # astronomically low.
     test = uuid.uuid4().hex[:12]
     return tixfile.with_suffix(f".{test}.tix")
+
+
+def get_admin_host_and_port_from_config(config):
+    admin_host = config.get("PGRST_ADMIN_SERVER_HOST", config["PGRST_SERVER_HOST"])
+    admin_port = config["PGRST_ADMIN_SERVER_PORT"]
+    return (admin_host, admin_port)

test/io/postgrest.py

@@ -57,6 +57,7 @@ class PostgrestProcess:
     admin: object
     process: object
     session: object
+    config: object
 
     def read_stdout(self, nlines=1):
         "Wait for line(s) on standard output."
@@ -143,6 +144,7 @@ def run(
                 process=process,
                 session=PostgrestSession(baseurl),
                 admin=PostgrestSession(adminurl),
+                config=env,
             )
         finally:
             remaining_output = process.stdout.read()

test/io/test_cli.py

@@ -1,12 +1,14 @@
 "Unit tests for Input/Ouput of PostgREST seen as a black box."
 
 from operator import attrgetter
+import signal
 import subprocess
 import pytest
 from syrupy.extensions.json import SingleFileSnapshotExtension
 import yaml
 
 from config import *
+from postgrest import *
 
 
 class ExtraNewLinesDumper(yaml.SafeDumper):
@@ -286,3 +288,105 @@ def test_jwt_secret_min_length(defaultenv):
 
     error = cli(["--dump-config"], env=env, expect_error=True)
     assert "The JWT secret must be at least 32 characters long." in error
+
+
+@pytest.mark.parametrize("host", ["127.0.0.1", "::1"], ids=["IPv4", "IPv6"])
+def test_cli_ready_flag_success(host, defaultenv):
+    "test PostgREST ready flag succeeds when ready"
+
+    port = freeport()
+
+    with run(env=defaultenv, host=host, port=port) as postgrest:
+        output = cli(["--ready"], env=postgrest.config)
+
+        (admin_host, admin_port) = get_admin_host_and_port_from_config(postgrest.config)
+
+        if is_ipv6(host):
+            assert f"OK: http://[{admin_host}]:{admin_port}/ready" in output
+        else:
+            assert f"OK: http://{admin_host}:{admin_port}/ready" in output
+
+
+def test_cli_ready_flag_fail_when_schema_cache_not_loaded(defaultenv, metapostgrest):
+    "test PosgREST ready flag fail when schema cache not loaded"
+
+    role = "timeout_authenticator"
+
+    env = {
+        **defaultenv,
+        "PGUSER": role,
+        "PGRST_DB_ANON_ROLE": role,
+        "PGRST_INTERNAL_SCHEMA_CACHE_SLEEP": "500",
+    }
+
+    port = freeport()
+
+    with run(env=env, port=port) as postgrest:
+        # The schema cache query takes at least 500ms, due to PGRST_INTERNAL_SCHEMA_CACHE_SLEEP above.
+        # Make it impossible to load the schema cache, by setting statement timeout to 400ms.
+        set_statement_timeout(metapostgrest, role, 400)
+
+        # force a reconnection so the new role setting is picked up
+        postgrest.process.send_signal(signal.SIGUSR1)
+
+        postgrest.wait_until_scache_starts_loading()
+
+        output = cli(["--ready"], env=postgrest.config, expect_error=True)
+        (admin_host, admin_port) = get_admin_host_and_port_from_config(postgrest.config)
+
+        assert f"ERROR: http://{admin_host}:{admin_port}/ready" in output
+
+
+def test_cli_ready_flag_fail_with_http_exception(defaultenv):
+    "test PostgREST ready flag fail when http exception occurs"
+
+    port = freeport()
+
+    # when healthcheck process sends the request to a wrong endpoint
+    with run(env=defaultenv, port=port) as postgrest:
+        # we set it to some freeport where admin is not running
+        postgrest.config["PGRST_ADMIN_SERVER_PORT"] = str(freeport())
+        output = cli(["--ready"], env=postgrest.config, expect_error=True)
+        (admin_host, admin_port) = get_admin_host_and_port_from_config(postgrest.config)
+
+        assert (
+            f"ERROR: connection refused to http://{admin_host}:{admin_port}/ready"
+            in output
+        )
+
+    # When client sends the request to invalid URL
+    with run(env=defaultenv, port=port) as postgrest:
+        postgrest.config["PGRST_ADMIN_SERVER_PORT"] = str(-1)
+        output = cli(["--ready"], env=postgrest.config, expect_error=True)
+        (admin_host, admin_port) = get_admin_host_and_port_from_config(postgrest.config)
+
+        assert f"ERROR: invalid url - http://{admin_host}:{admin_port}/ready" in output
+
+
+def test_cli_ready_flag_fail_with_special_hostname(defaultenv):
+    "test PostgREST ready flag fail when http exception occurs"
+
+    port = freeport()
+    host = "*4"
+
+    with run(env=defaultenv, host=host, port=port) as postgrest:
+        output = cli(["--ready"], env=postgrest.config, expect_error=True)
+
+        assert (
+            f'ERROR: The `--ready` flag cannot be used when server-host is configured as "{host}". Please update your server-host config to "localhost".'
+            in output
+        )
+
+
+def test_cli_ready_flag_fail_when_no_admin_server(defaultenv):
+    "test PostgREST ready flag fail when admin server not running"
+
+    with run(env=defaultenv) as postgrest:
+        # We set admin-server-port to <empty> to disable admin server
+        postgrest.config["PGRST_ADMIN_SERVER_PORT"] = ""
+        output = cli(["--ready"], env=postgrest.config, expect_error=True)
+
+        assert (
+            "ERROR: Admin server is not running. Please check admin-server-port config."
+            in output
+        )