Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions changelog.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:
---------
Expand Down
24 changes: 21 additions & 3 deletions pgcli/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than check self.force_destructive in two places, I'd change the signature of confirm_destructive_query to include it as a parameter.

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)]
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -1454,6 +1470,7 @@ def cli(
ssh_tunnel: str,
init_command: str,
log_file: str,
force_destructive: bool,
):
if version:
print("Version:", __version__)
Expand Down Expand Up @@ -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.
Expand Down
48 changes: 48 additions & 0 deletions tests/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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