diff --git a/changelog.rst b/changelog.rst index 96eefd747..5b6cffa78 100644 --- a/changelog.rst +++ b/changelog.rst @@ -9,6 +9,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 0b4b64f59..8606302a3 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -185,12 +185,14 @@ def __init__( warn=None, ssh_tunnel_url: Optional[str] = None, log_file: Optional[str] = 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) @@ -484,7 +486,10 @@ 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) + if self.force_destructive: + destroy = True + else: + destroy = confirm_destructive_query(query, self.destructive_warning, self.dsn_alias) if destroy is False: message = "Wise choice. Command execution stopped." return [(None, None, None, message)] @@ -792,11 +797,14 @@ 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) + if self.force_destructive: + destroy = True + else: + destroy = confirm_destructive_query(text, self.destructive_warning, self.dsn_alias) 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) @@ -1426,6 +1434,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( @@ -1454,6 +1470,7 @@ def cli( ssh_tunnel: str, init_command: str, log_file: str, + force_destructive: bool, ): if version: print("Version:", __version__) @@ -1512,6 +1529,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/tests/test_main.py b/tests/test_main.py index 5cf1d09f8..defcb206c 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -595,3 +595,51 @@ 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"] + + # Mock confirm_destructive_query to ensure it's not called + with mock.patch("pgcli.main.confirm_destructive_query") 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 NOT called + mock_confirm.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