diff --git a/changelog.rst b/changelog.rst index 25b4a7a85..7f997b496 100644 --- a/changelog.rst +++ b/changelog.rst @@ -34,6 +34,10 @@ Features: * Support dsn specific init-command in the config file * Add suggestion when setting the search_path * Allow per dsn_alias ssh tunnel selection +* Add support for forcing destructive commands without confirmation. + * Command line option `-y` or `--yes`. + * Skips the destructive command confirmation prompt when enabled. + * Useful for automated scripts and CI/CD pipelines. Internal: --------- diff --git a/pgcli/main.py b/pgcli/main.py index 94ac9af10..1c2e60088 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -187,12 +187,14 @@ def __init__( warn=None, ssh_tunnel_url: str | None = None, log_file: str | None = None, + force_destructive: bool = False, ): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None self.watch_command = None + self.force_destructive = force_destructive # Load config. c = self.config = get_config(pgclirc_file) @@ -510,7 +512,7 @@ def execute_from_file(self, pattern, **_): ): message = "Destructive statements must be run within a transaction. Command execution stopped." return [(None, None, None, message)] - destroy = confirm_destructive_query(query, self.destructive_warning, self.dsn_alias) + destroy = confirm_destructive_query(query, self.destructive_warning, self.dsn_alias, self.force_destructive) if destroy is False: message = "Wise choice. Command execution stopped." return [(None, None, None, message)] @@ -824,11 +826,11 @@ def execute_command(self, text, handle_closed_connection=True): ): click.secho("Destructive statements must be run within a transaction.") raise KeyboardInterrupt - destroy = confirm_destructive_query(text, self.destructive_warning, self.dsn_alias) + destroy = confirm_destructive_query(text, self.destructive_warning, self.dsn_alias, self.force_destructive) if destroy is False: click.secho("Wise choice!") raise KeyboardInterrupt - elif destroy: + elif destroy and not self.force_destructive: click.secho("Your call!") output, query = self._evaluate_command(text) @@ -1488,6 +1490,14 @@ def echo_via_pager(self, text, color=None): type=str, help="SQL statement to execute after connecting.", ) +@click.option( + "-y", + "--yes", + "force_destructive", + is_flag=True, + default=False, + help="Force destructive commands without confirmation prompt.", +) @click.argument("dbname", default=lambda: None, envvar="PGDATABASE", nargs=1) @click.argument("username", default=lambda: None, envvar="PGUSER", nargs=1) def cli( @@ -1516,6 +1526,7 @@ def cli( ssh_tunnel: str, init_command: str, log_file: str, + force_destructive: bool, ): if version: print("Version:", __version__) @@ -1574,6 +1585,7 @@ def cli( warn=warn, ssh_tunnel_url=ssh_tunnel, log_file=log_file, + force_destructive=force_destructive, ) # Choose which ever one has a valid value. diff --git a/pgcli/packages/prompt_utils.py b/pgcli/packages/prompt_utils.py index 997b86e7f..4d21c09d4 100644 --- a/pgcli/packages/prompt_utils.py +++ b/pgcli/packages/prompt_utils.py @@ -3,22 +3,31 @@ from .parseutils import is_destructive -def confirm_destructive_query(queries, keywords, alias): +def confirm_destructive_query(queries, keywords, alias, force=False): """Check if the query is destructive and prompts the user to confirm. Returns: * None if the query is non-destructive or we can't prompt the user. - * True if the query is destructive and the user wants to proceed. + * True if the query is destructive and the user wants to proceed, or if the + confirmation was bypassed with `force`. * False if the query is destructive and the user doesn't want to proceed. """ + if not is_destructive(queries, keywords): + return None + + if force: + return True + + if not sys.stdin.isatty(): + return None + info = "You're about to run a destructive command" if alias: info += f" in {click.style(alias, fg='red')}" prompt_text = f"{info}.\nDo you want to proceed?" - if is_destructive(queries, keywords) and sys.stdin.isatty(): - return confirm(prompt_text) + return confirm(prompt_text) def confirm(*args, **kwargs): diff --git a/tests/test_main.py b/tests/test_main.py index 52415a008..f33b06c51 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -636,3 +636,52 @@ def test_notifications(executor): with mock.patch("pgcli.main.click.secho") as mock_secho: run(executor, "notify chan1, 'testing2'") mock_secho.assert_not_called() + + +def test_force_destructive_flag(): + """Test that PGCli can be initialized with force_destructive flag.""" + cli = PGCli(force_destructive=True) + assert cli.force_destructive is True + + cli = PGCli(force_destructive=False) + assert cli.force_destructive is False + + cli = PGCli() + assert cli.force_destructive is False + + +@dbtest +def test_force_destructive_skips_confirmation(executor): + """Test that force_destructive=True skips confirmation for destructive commands.""" + cli = PGCli(pgexecute=executor, force_destructive=True) + cli.destructive_warning = ["drop", "alter"] + + # The proceed/abort decision lives inside confirm_destructive_query, which is + # told to force; what must not happen is the user being prompted. + with mock.patch("pgcli.packages.prompt_utils.confirm") as mock_prompt: + # Execute a destructive command + result = cli.execute_command("ALTER TABLE test_table ADD COLUMN test_col TEXT;") + + # Verify that the user was never prompted + mock_prompt.assert_not_called() + + # Verify that the command was attempted (even if it fails due to missing table) + assert result is not None + + +@dbtest +def test_without_force_destructive_calls_confirmation(executor): + """Test that without force_destructive, confirmation is called for destructive commands.""" + cli = PGCli(pgexecute=executor, force_destructive=False) + cli.destructive_warning = ["drop", "alter"] + + # Mock confirm_destructive_query to return True (user confirms) + with mock.patch("pgcli.main.confirm_destructive_query", return_value=True) as mock_confirm: + # Execute a destructive command + result = cli.execute_command("ALTER TABLE test_table ADD COLUMN test_col TEXT;") + + # Verify that confirm_destructive_query WAS called + mock_confirm.assert_called_once() + + # Verify that the command was attempted + assert result is not None