diff --git a/changelog.rst b/changelog.rst index f6e99520c..db9a39a96 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 dfffc5fe1..a9d53f689 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -800,6 +800,34 @@ def execute_command(self, text, handle_closed_connection=True): 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 @@ def run_cli(self): 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 @@ def run_cli(self): 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 da27bb96a..ee497b98a 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 3b5c82b39..8f184bbdf 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