Skip to content
Open
Show file tree
Hide file tree
Changes from all 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:
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