Add config option to not automatically restart connection on destructive warning abort; defaults to not restarting. (#1379)

changelog.rst

@@ -10,6 +10,8 @@ Features:
 * pgcli.magic will now work with connection URLs that use TLS client certificates for authentication
 * Have config option to retry queries on operational errors like connections being lost.
   Also prevents getting stuck in a retry loop.
+* Config option to not restart connection when cancelling a `destructive_warning` query. By default,
+  it will now not restart.
 
 3.5.0 (2022/09/15):
 ===================

pgcli/main.py

@@ -231,6 +231,9 @@ class PGCli:
         self.destructive_warning = parse_destructive_warning(
             warn or c["main"].as_list("destructive_warning")
         )
+        self.destructive_warning_restarts_connection = c["main"].as_bool(
+            "destructive_warning_restarts_connection"
+        )
 
         self.less_chatty = bool(less_chatty) or c["main"].as_bool("less_chatty")
         self.null_string = c["main"].get("null_string", "<null>")
@@ -711,10 +714,16 @@ class PGCli:
 
             output, query = self._evaluate_command(text)
         except KeyboardInterrupt:
-            # Restart connection to the database
-            self.pgexecute.connect()
-            logger.debug("cancelled query, sql: %r", text)
-            click.secho("cancelled query", err=True, fg="red")
+            if self.destructive_warning_restarts_connection:
+                # Restart connection to the database
+                self.pgexecute.connect()
+                logger.debug("cancelled query and restarted connection, sql: %r", text)
+                click.secho(
+                    "cancelled query and restarted connection", err=True, fg="red"
+                )
+            else:
+                logger.debug("cancelled query, sql: %r", text)
+                click.secho("cancelled query", err=True, fg="red")
         except NotImplementedError:
             click.secho("Not Yet Implemented.", fg="yellow")
         except OperationalError as e:

pgcli/pgclirc

@@ -29,6 +29,12 @@ multi_line_mode = psql
 # "unconditional_update" will warn you of update statements that don't have a where clause
 destructive_warning = drop, shutdown, delete, truncate, alter, update, unconditional_update
 
+# Destructive warning can restart the connection if this is enabled and the
+# user declines. This means that any current uncommitted transaction can be
+# aborted if the user doesn't want to proceed with a destructive_warning 
+# statement.
+destructive_warning_restarts_connection = False
+
 # Enables expand mode, which is similar to `\x` in psql.
 expand = False
 

tests/features/basic_commands.feature

@@ -23,6 +23,13 @@ Feature: run the cli,
      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"
+      then we see cancelled query warning
+      when we check for any non-idle sleep queries
+      then we don't see any non-idle sleep queries
+
   Scenario: list databases
       When we list databases
       then we see list of databases

tests/features/crud_database.feature

@@ -5,7 +5,7 @@ Feature: manipulate databases:
      When we create database
       then we see database created
       when we drop database
-      then we confirm the destructive warning
+      then we respond to the destructive warning: y
       then we see database dropped
       when we connect to dbserver
       then we see database connected

tests/features/crud_table.feature

@@ -8,15 +8,38 @@ Feature: manipulate tables:
       then we see table created
       when we insert into table
       then we see record inserted
+      when we select from table
+      then we see data selected: initial
       when we update table
       then we see record updated
       when we select from table
-      then we see data selected
+      then we see data selected: updated
       when we delete from table
-      then we confirm the destructive warning
+      then we respond to the destructive warning: y
       then we see record deleted
       when we drop table
-      then we confirm the destructive warning
+      then we respond to the destructive warning: y
       then we see table dropped
       when we connect to dbserver
       then we see database connected
+
+  Scenario: transaction handling, with cancelling on a destructive warning.
+    When we connect to test database
+      then we see database connected
+      when we create table
+      then we see table created
+      when we begin transaction
+      then we see transaction began
+      when we insert into table
+      then we see record inserted
+      when we delete from table
+      then we respond to the destructive warning: n
+      when we select from table
+      then we see data selected: initial
+      when we rollback transaction
+      then we see transaction rolled back
+      when we select from table
+      then we see select output without data
+      when we drop table
+      then we respond to the destructive warning: y
+      then we see table dropped

tests/features/expanded.feature

@@ -7,7 +7,7 @@ Feature: expanded mode:
       and we select from table
       then we see expanded data selected
       when we drop table
-      then we confirm the destructive warning
+      then we respond to the destructive warning: y
       then we see table dropped
 
   Scenario: expanded off
@@ -16,7 +16,7 @@ Feature: expanded mode:
       and we select from table
       then we see nonexpanded data selected
       when we drop table
-      then we confirm the destructive warning
+      then we respond to the destructive warning: y
       then we see table dropped
 
   Scenario: expanded auto
@@ -25,5 +25,5 @@ Feature: expanded mode:
       and we select from table
       then we see auto data selected
       when we drop table
-      then we confirm the destructive warning
+      then we respond to the destructive warning: y
       then we see table dropped

tests/features/steps/basic_commands.py

@@ -73,6 +73,59 @@ def step_ctrl_d(context):
     context.exit_sent = True
 
 
+@when('we send "ctrl + c"')
+def step_ctrl_c(context):
+    """Send Ctrl + c to hopefully interrupt."""
+    context.cli.sendcontrol("c")
+
+
+@then("we see cancelled query warning")
+def step_see_cancelled_query_warning(context):
+    """
+    Make sure we receive the warning that the current query was cancelled.
+    """
+    wrappers.expect_exact(context, "cancelled query", timeout=2)
+
+
+@when("we send sleep query")
+def step_send_sleep_15_seconds(context):
+    """
+    Send query to sleep for 15 seconds.
+    """
+    context.cli.sendline("select pg_sleep(15)")
+
+
+@when("we check for any non-idle sleep queries")
+def step_check_for_active_sleep_queries(context):
+    """
+    Send query to check for any non-idle pg_sleep queries.
+    """
+    context.cli.sendline(
+        "select state from pg_stat_activity where query not like '%pg_stat_activity%' and query like '%pg_sleep%' and state != 'idle';"
+    )
+
+
+@then("we don't see any non-idle sleep queries")
+def step_no_active_sleep_queries(context):
+    """Confirm that any pg_sleep queries are either idle or not active."""
+    wrappers.expect_exact(
+        context,
+        context.conf["pager_boundary"]
+        + "\r"
+        + dedent(
+            """
+            +-------+\r
+            | state |\r
+            |-------|\r
+            +-------+\r
+            SELECT 0\r
+        """
+        )
+        + context.conf["pager_boundary"],
+        timeout=5,
+    )
+
+
 @when(r'we send "\?" command')
 def step_send_help(context):
     r"""
@@ -131,15 +184,15 @@ def step_see_found(context):
     )
 
 
-@then("we confirm the destructive warning")
-def step_confirm_destructive_command(context):
-    """Confirm destructive command."""
+@then("we respond to the destructive warning: {response}")
+def step_resppond_to_destructive_command(context, response):
+    """Respond to destructive command."""
     wrappers.expect_exact(
         context,
         "You're about to run a destructive command.\r\nDo you want to proceed? (y/n):",
         timeout=2,
     )
-    context.cli.sendline("y")
+    context.cli.sendline(response.strip())
 
 
 @then("we send password")

tests/features/steps/crud_table.py

@@ -9,6 +9,10 @@ from textwrap import dedent
 import wrappers
 
 
+INITIAL_DATA = "xxx"
+UPDATED_DATA = "yyy"
+
+
 @when("we create table")
 def step_create_table(context):
     """
@@ -22,7 +26,7 @@ def step_insert_into_table(context):
     """
     Send insert into table.
     """
-    context.cli.sendline("""insert into a(x) values('xxx');""")
+    context.cli.sendline(f"""insert into a(x) values('{INITIAL_DATA}');""")
 
 
 @when("we update table")
@@ -30,7 +34,9 @@ def step_update_table(context):
     """
     Send insert into table.
     """
-    context.cli.sendline("""update a set x = 'yyy' where x = 'xxx';""")
+    context.cli.sendline(
+        f"""update a set x = '{UPDATED_DATA}' where x = '{INITIAL_DATA}';"""
+    )
 
 
 @when("we select from table")
@@ -46,7 +52,7 @@ def step_delete_from_table(context):
     """
     Send deete from table.
     """
-    context.cli.sendline("""delete from a where x = 'yyy';""")
+    context.cli.sendline(f"""delete from a where x = '{UPDATED_DATA}';""")
 
 
 @when("we drop table")
@@ -57,6 +63,30 @@ def step_drop_table(context):
     context.cli.sendline("drop table a;")
 
 
+@when("we alter the table")
+def step_alter_table(context):
+    """
+    Alter the table by adding a column.
+    """
+    context.cli.sendline("""alter table a add column y varchar;""")
+
+
+@when("we begin transaction")
+def step_begin_transaction(context):
+    """
+    Begin transaction
+    """
+    context.cli.sendline("begin;")
+
+
+@when("we rollback transaction")
+def step_rollback_transaction(context):
+    """
+    Rollback transaction
+    """
+    context.cli.sendline("rollback;")
+
+
 @then("we see table created")
 def step_see_table_created(context):
     """
@@ -81,21 +111,42 @@ def step_see_record_updated(context):
     wrappers.expect_pager(context, "UPDATE 1\r\n", timeout=2)
 
 
-@then("we see data selected")
-def step_see_data_selected(context):
+@then("we see data selected: {data}")
+def step_see_data_selected(context, data):
     """
-    Wait to see select output.
+    Wait to see select output with initial or updated data.
+    """
+    x = UPDATED_DATA if data == "updated" else INITIAL_DATA
+    wrappers.expect_pager(
+        context,
+        dedent(
+            f"""\
+            +-----+\r
+            | x   |\r
+            |-----|\r
+            | {x} |\r
+            +-----+\r
+            SELECT 1\r
+        """
+        ),
+        timeout=1,
+    )
+
+
+@then("we see select output without data")
+def step_see_no_data_selected(context):
+    """
+    Wait to see select output without data.
     """
     wrappers.expect_pager(
         context,
         dedent(
             """\
-            +-----+\r
-            | x   |\r
-            |-----|\r
-            | yyy |\r
-            +-----+\r
-            SELECT 1\r
+            +---+\r
+            | x |\r
+            |---|\r
+            +---+\r
+            SELECT 0\r
         """
         ),
         timeout=1,
@@ -116,3 +167,19 @@ def step_see_table_dropped(context):
     Wait to see drop output.
     """
     wrappers.expect_pager(context, "DROP TABLE\r\n", timeout=2)
+
+
+@then("we see transaction began")
+def step_see_transaction_began(context):
+    """
+    Wait to see transaction began.
+    """
+    wrappers.expect_pager(context, "BEGIN\r\n", timeout=2)
+
+
+@then("we see transaction rolled back")
+def step_see_transaction_rolled_back(context):
+    """
+    Wait to see transaction rollback.
+    """
+    wrappers.expect_pager(context, "ROLLBACK\r\n", timeout=2)