Add --ping as a CLI option (#1471)

* Add `--ping` as a CLI option This builds in the functionality asked for in #1470; it allows pgcli to replace `pg_isready` from the postgresql command line toolchain * Update the changelog and AUTHORS * Use click.echo pgcli.echo doesn't seem to actually _echo_ anything * Add integ tests for the ping feature * Refactor try / else Co-authored-by: Damien Baty <[email protected]> * Unused variable --------- Co-authored-by: Irina Truong <[email protected]> Co-authored-by: Damien Baty <[email protected]>
dbcli · Jun 20, 2024 · b979180 · b979180
1 parent 46fe654
commit b979180
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 2 deletions.
diff --git a/AUTHORS b/AUTHORS
@@ -134,6 +134,7 @@ Contributors:
     * Antonio Aguilar (crazybolillo)
     * Andrew M. MacFie (amacfie)
     * saucoide
+    * Chris Rose (offbyone/offby1)
 
 Creator:
 --------

diff --git a/changelog.rst b/changelog.rst
@@ -1,3 +1,10 @@
+Dev
+===
+
+Features
+--------
+* Add a `--ping` command line option; allows pgcli to replace `pg_isready`
+
 4.1.0 (2024-03-09)
 ==================
 

diff --git a/pgcli/main.py b/pgcli/main.py
@@ -1463,6 +1463,13 @@ def echo_via_pager(self, text, color=None):
     is_flag=True,
     help="list available databases, then exit.",
 )
+@click.option(
+    "--ping",
+    "ping_database",
+    is_flag=True,
+    default=False,
+    help="Check database connectivity, then exit.",
+)
 @click.option(
     "--auto-vertical-output",
     is_flag=True,
@@ -1504,6 +1511,7 @@ def cli(
     prompt,
     prompt_dsn,
     list_databases,
+    ping_database,
     auto_vertical_output,
     list_dsn,
     warn,
@@ -1581,8 +1589,8 @@ def cli(
         service = database[8:]
     elif os.getenv("PGSERVICE") is not None:
         service = os.getenv("PGSERVICE")
-    # because option --list or -l are not supposed to have a db name
-    if list_databases:
+    # because option --ping, --list or -l are not supposed to have a db name
+    if list_databases or ping_database:
         database = "postgres"
 
     if dsn != "":
@@ -1626,6 +1634,20 @@ def cli(
 
         sys.exit(0)
 
+    if ping_database:
+        try:
+            list(pgcli.pgexecute.run("SELECT 1"))
+        except Exception:
+            click.secho(
+                "Could not connect to the database. Please check that the database is running.",
+                err=True,
+                fg="red",
+            )
+            sys.exit(1)
+        else:
+            click.echo("PONG")
+            sys.exit(0)
+
     pgcli.logger.debug(
         "Launch Params: \n" "\tdatabase: %r" "\tuser: %r" "\thost: %r" "\tport: %r",
         database,

diff --git a/tests/features/basic_commands.feature b/tests/features/basic_commands.feature
@@ -51,6 +51,10 @@ Feature: run the cli,
       When we list databases
       then we see list of databases
 
+  Scenario: ping databases
+      When we ping the database
+      then we get a pong response
+
   Scenario: run the cli with --username
     When we launch dbcli using --username
       and we send "\?" command

diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py
@@ -26,6 +26,19 @@ def step_see_list_databases(context):
     context.cmd_output = None
 
 
+@when("we ping the database")
+def step_ping_database(context):
+    cmd = ["pgcli", "--ping"]
+    context.cmd_output = subprocess.check_output(cmd, cwd=context.package_root)
+
+
+@then("we get a pong response")
+def step_get_pong_response(context):
+    # exit code 0 is implied by the presence of cmd_output here, which
+    # is only set on a successful run.
+    assert context.cmd_output.strip() == b"PONG", f"Output was {context.cmd_output}"
+
+
 @when("we run dbcli")
 def step_run_cli(context):
     wrappers.run_cli(context)