diff options
author | Andy Schoenberger <akschoenberger@gmail.com> | 2022-11-17 23:13:05 -0500 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-11-17 20:13:05 -0800 |
commit | d6ca4c346462844f3209cedfdc892def2fddf38d (patch) | |
tree | bd5da927a2d1b4ce1506d461ab58c18a318f8c0e | |
parent | fa054a5546d33ed34586513f71f371f81ef6ecfd (diff) |
Add config option to not automatically restart connection on destructive warning abort; defaults to not restarting. (#1379)
-rw-r--r-- | changelog.rst | 2 | ||||
-rw-r--r-- | pgcli/main.py | 17 | ||||
-rw-r--r-- | pgcli/pgclirc | 6 | ||||
-rw-r--r-- | tests/features/basic_commands.feature | 7 | ||||
-rw-r--r-- | tests/features/crud_database.feature | 2 | ||||
-rw-r--r-- | tests/features/crud_table.feature | 29 | ||||
-rw-r--r-- | tests/features/expanded.feature | 6 | ||||
-rw-r--r-- | tests/features/steps/basic_commands.py | 61 | ||||
-rw-r--r-- | tests/features/steps/crud_table.py | 83 |
9 files changed, 190 insertions, 23 deletions
diff --git a/changelog.rst b/changelog.rst index 6fdd2e37..19b1c3ea 100644 --- a/changelog.rst +++ b/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): =================== diff --git a/pgcli/main.py b/pgcli/main.py index 9353b6c1..8c7d59dc 100644 --- a/pgcli/main.py +++ b/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: diff --git a/pgcli/pgclirc b/pgcli/pgclirc index 7626db5e..b849adfe 100644 --- a/pgcli/pgclirc +++ b/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 diff --git a/tests/features/basic_commands.feature b/tests/features/basic_commands.feature index cd15306b..da27bb96 100644 --- a/tests/features/basic_commands.feature +++ b/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 diff --git a/tests/features/crud_database.feature b/tests/features/crud_database.feature index ed13bbe0..87da4e39 100644 --- a/tests/features/crud_database.feature +++ b/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 diff --git a/tests/features/crud_table.feature b/tests/features/crud_table.feature index 1f9db4a0..8a43c5c0 100644 --- a/tests/features/crud_table.feature +++ b/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 diff --git a/tests/features/expanded.feature b/tests/features/expanded.feature index 4f381f81..e4860486 100644 --- a/tests/features/expanded.feature +++ b/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 diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py index 7c878143..3b5c82b3 100644 --- a/tests/features/steps/basic_commands.py +++ b/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") diff --git a/tests/features/steps/crud_table.py b/tests/features/steps/crud_table.py index 0375883a..27d543e7 100644 --- a/tests/features/steps/crud_table.py +++ b/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,19 +111,20 @@ 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 - | yyy |\r + | {x} |\r +-----+\r SELECT 1\r """ @@ -102,6 +133,26 @@ def step_see_data_selected(context): ) +@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 + +---+\r + SELECT 0\r + """ + ), + timeout=1, + ) + + @then("we see record deleted") def step_see_data_deleted(context): """ @@ -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) |