From b979180621abb7dad5805fc9734e9695fd888e34 Mon Sep 17 00:00:00 2001
From: Chris Rose <offbyone@github.com>
Date: Wed, 19 Jun 2024 21:45:50 -0700
Subject: [PATCH] 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 <damien@damienbaty.com>

* Unused variable

---------

Co-authored-by: Irina Truong <637013+j-bennet@users.noreply.github.com>
Co-authored-by: Damien Baty <damien@damienbaty.com>
---
 AUTHORS                                |  1 +
 changelog.rst                          |  7 +++++++
 pgcli/main.py                          | 26 ++++++++++++++++++++++++--
 tests/features/basic_commands.feature  |  4 ++++
 tests/features/steps/basic_commands.py | 13 +++++++++++++
 5 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/AUTHORS b/AUTHORS
index 9f33ff551..fefddefb1 100644
--- 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
index 744e9031e..5113a3095 100644
--- 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
index 056a9403f..f07255e56 100644
--- 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
index ee497b98a..7e975dcc8 100644
--- 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
index 687bdc0a9..00ac277f5 100644
--- 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)