Ask for confirmation to quit cli if a transaction is ongoing. (#1400)

If user tries to quit the cli while a transaction is ongoing (i.e.
begun, but not committed or rolled back yet), pgcli now asks for a
confirmation. The user can choose to commit, rollback or cancel the
exit. If the user chooses to commit or rollback, we exit only if the
commit/rollback succeeds.

Fixes #1071.

4 files changed, 88 insertions, 3 deletions

diff --git a/changelog.rst b/changelog.rst
index f6e99520..db9a39a9 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -4,6 +4,7 @@ Upcoming
 Features:
 ---------
 
+* Ask for confirmation when quitting cli while a transaction is ongoing.
 * New `destructive_statements_require_transaction` config option to refuse to execute a
   destructive SQL statement if outside a transaction. This option is off by default.
 * Changed the `destructive_warning` config to be a list of commands that are considered
diff --git a/pgcli/main.py b/pgcli/main.py
index dfffc5fe..a9d53f68 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -800,6 +800,34 @@ class PGCli:
                 logger.debug("Search path: %r", self.completer.search_path)
         return query
 
+    def _check_ongoing_transaction_and_allow_quitting(self):
+        """Return whether we can really quit, possibly by asking the
+        user to confirm so if there is an ongoing transaction.
+        """
+        if not self.pgexecute.valid_transaction():
+            return True
+        while 1:
+            try:
+                choice = click.prompt(
+                    "A transaction is ongoing. Choose `c` to COMMIT, `r` to ROLLBACK, `a` to abort exit.",
+                    default="a",
+                )
+            except click.Abort:
+                # Print newline if user aborts with `^C`, otherwise
+                # pgcli's prompt will be printed on the same line
+                # (just after the confirmation prompt).
+                click.echo(None, err=False)
+                choice = "a"
+            choice = choice.lower()
+            if choice == "a":
+                return False  # do not quit
+            if choice == "c":
+                query = self.execute_command("commit")
+                return query.successful  # quit only if query is successful
+            if choice == "r":
+                query = self.execute_command("rollback")
+                return query.successful  # quit only if query is successful
+
     def run_cli(self):
         logger = self.logger
 
@@ -822,6 +850,10 @@ class PGCli:
                     text = self.prompt_app.prompt()
                 except KeyboardInterrupt:
                     continue
+                except EOFError:
+                    if not self._check_ongoing_transaction_and_allow_quitting():
+                        continue
+                    raise
 
                 try:
                     text = self.handle_editor_command(text)
@@ -831,7 +863,12 @@ class PGCli:
                     click.secho(str(e), err=True, fg="red")
                     continue
 
-                self.handle_watch_command(text)
+                try:
+                    self.handle_watch_command(text)
+                except PgCliQuitError:
+                    if not self._check_ongoing_transaction_and_allow_quitting():
+                        continue
+                    raise
 
                 self.now = dt.datetime.today()
 
diff --git a/tests/features/basic_commands.feature b/tests/features/basic_commands.feature
index da27bb96..ee497b98 100644
--- a/tests/features/basic_commands.feature
+++ b/tests/features/basic_commands.feature
@@ -23,6 +23,23 @@ Feature: run the cli,
      When we send "ctrl + d"
       then dbcli exits
 
+  Scenario: confirm exit when a transaction is ongoing
+     When we begin transaction
+      and we try to send "ctrl + d"
+      then we see ongoing transaction message
+      when we send "c"
+      then dbcli exits
+
+  Scenario: cancel exit when a transaction is ongoing
+     When we begin transaction
+      and we try to send "ctrl + d"
+      then we see ongoing transaction message
+      when we send "a"
+      then we see dbcli prompt
+      when we rollback transaction
+      when we send "ctrl + d"
+      then dbcli exits
+
   Scenario: interrupt current query via "ctrl + c"
      When we send sleep query
       and we send "ctrl + c"
diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py
index 3b5c82b3..8f184bbd 100644
--- a/tests/features/steps/basic_commands.py
+++ b/tests/features/steps/basic_commands.py
@@ -64,13 +64,22 @@ def step_ctrl_d(context):
     """
     Send Ctrl + D to hopefully exit.
     """
+    step_try_to_ctrl_d(context)
+    context.cli.expect(pexpect.EOF, timeout=5)
+    context.exit_sent = True
+
+
+@when('we try to send "ctrl + d"')
+def step_try_to_ctrl_d(context):
+    """
+    Send Ctrl + D, perhaps exiting, perhaps not (if a transaction is
+    ongoing).
+    """
     # turn off pager before exiting
     context.cli.sendcontrol("c")
     context.cli.sendline(r"\pset pager off")
     wrappers.wait_prompt(context)
     context.cli.sendcontrol("d")
-    context.cli.expect(pexpect.EOF, timeout=5)
-    context.exit_sent = True
 
 
 @when('we send "ctrl + c"')
@@ -87,6 +96,14 @@ def step_see_cancelled_query_warning(context):
     wrappers.expect_exact(context, "cancelled query", timeout=2)
 
 
+@then("we see ongoing transaction message")
+def step_see_ongoing_transaction_error(context):
+    """
+    Make sure we receive the warning that a transaction is ongoing.
+    """
+    context.cli.expect("A transaction is ongoing.", timeout=2)
+
+
 @when("we send sleep query")
 def step_send_sleep_15_seconds(context):
     """
@@ -199,3 +216,16 @@ def step_resppond_to_destructive_command(context, response):
 def step_send_password(context):
     wrappers.expect_exact(context, "Password for", timeout=5)
     context.cli.sendline(context.conf["pass"] or "DOES NOT MATTER")
+
+
+@when('we send "{text}"')
+def step_send_text(context, text):
+    context.cli.sendline(text)
+    # Try to detect whether we are exiting. If so, set `exit_sent`
+    # so that `after_scenario` correctly cleans up.
+    try:
+        context.cli.expect(pexpect.EOF, timeout=0.2)
+    except pexpect.TIMEOUT:
+        pass
+    else:
+        context.exit_sent = True