From d0469175e8bd5a8ac084282c8015a253550755de Mon Sep 17 00:00:00 2001 From: Roland Walker Date: Sat, 22 Aug 2026 09:08:10 -0400 Subject: [PATCH] exit on error for --execute containing /source If a /source'd file contains a special command, and the --special flag was not given to /source, then --execute should be able to print an error message to the stderr and exit with an exit code. The error message was previously visible only at the REPL. This requires adding an is_error property to SQLResult, which may be useful in other contexts. --- changelog.md | 5 +++ mycli/client_commands.py | 10 ++++-- mycli/client_query.py | 9 ++++++ mycli/main_modes/execute.py | 2 +- mycli/packages/sqlresult.py | 4 ++- test/pytests/test_client_commands.py | 7 +++-- test/pytests/test_client_query.py | 41 +++++++++++++++++++++++++ test/pytests/test_main.py | 4 +-- test/pytests/test_main_modes_execute.py | 12 ++++---- test/pytests/test_sqlresult.py | 2 ++ 10 files changed, 82 insertions(+), 14 deletions(-) diff --git a/changelog.md b/changelog.md index 084ad72b4..5b89d10c0 100644 --- a/changelog.md +++ b/changelog.md @@ -7,6 +7,11 @@ Features * Add `/source --special` to allow executing some special commands. +Bug Fixes +--------- +* Show CLI error on invalid `--execute` containing `/source`. + + 2.15.0 (2026/08/20) ============== diff --git a/mycli/client_commands.py b/mycli/client_commands.py index d12e476e4..a4ca46d62 100644 --- a/mycli/client_commands.py +++ b/mycli/client_commands.py @@ -367,11 +367,17 @@ def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]: special_query = query.rstrip(';') if special.is_special_command(special_query): if not allow_special: - yield SQLResult(status='Special commands are not supported without /source --special.') + yield SQLResult( + status='Special commands are not supported without /source --special.', + is_error=True, + ) return if not _source_special_command_is_safe(special_query): command, _verbosity, _arg = special.parse_special_command(special_query) - yield SQLResult(status=f'Special command is never permitted in source files: {command}.') + yield SQLResult( + status=f'Special command is never permitted in source files: {command}.', + is_error=True, + ) return yield from self.sqlexecute.run(special_query) continue diff --git a/mycli/client_query.py b/mycli/client_query.py index 1e563dd3b..5ae94b436 100644 --- a/mycli/client_query.py +++ b/mycli/client_query.py @@ -10,6 +10,10 @@ from mycli.sqlcompleter import SQLCompleter +class QueryError(Exception): + """An error reported while consuming query results.""" + + class ClientQueryMixin: if TYPE_CHECKING: schema_prefetcher: Any @@ -80,6 +84,7 @@ def run_query( query: str, checkpoint: str | None = None, new_line: bool = True, + raise_on_error: bool = False, ) -> None: """Runs *query*.""" assert self.sqlexecute is not None @@ -91,6 +96,10 @@ def run_query( self.main_formatter.query = query self.redirect_formatter.query = query self.explorer_formatter.query = query + if result.is_error and raise_on_error: + message = result.status_plain or 'Query failed.' + self.log_output(message) + raise QueryError(message) output = self.format_sqlresult( result, is_expanded=special.is_expanded_output(), diff --git a/mycli/main_modes/execute.py b/mycli/main_modes/execute.py index dcabf46ac..16f7cf604 100644 --- a/mycli/main_modes/execute.py +++ b/mycli/main_modes/execute.py @@ -46,7 +46,7 @@ def main_execute_from_cli(mycli: 'MyCli', cli_args: 'CliArgs') -> int: mycli.logger.warning('Unable to open TTY as stdin.') raise e if execution_confirmed: - mycli.run_query(execute_sql, checkpoint=cli_args.checkpoint) + mycli.run_query(execute_sql, checkpoint=cli_args.checkpoint, raise_on_error=True) return 0 else: return 1 diff --git a/mycli/packages/sqlresult.py b/mycli/packages/sqlresult.py index 0b9d4b9c7..d288a6aa6 100644 --- a/mycli/packages/sqlresult.py +++ b/mycli/packages/sqlresult.py @@ -17,11 +17,13 @@ class SQLResult: command: dict[str, str | float] | None = None image: bytes | None = None image_protocol: ImageProtocol = 'none' + is_error: bool = False def __str__(self): image = f'<{len(self.image)} bytes>' if self.image is not None else None return ( - f"{self.preamble}, {self.header}, {self.rows}, {self.postamble}, {self.status}, {self.command}, {image}, {self.image_protocol}" + f"{self.preamble}, {self.header}, {self.rows}, {self.postamble}, {self.status}, {self.command}, " + f"{image}, {self.image_protocol}, {self.is_error}" ) @cached_property diff --git a/test/pytests/test_client_commands.py b/test/pytests/test_client_commands.py index 3cd6301e3..449b09531 100644 --- a/test/pytests/test_client_commands.py +++ b/test/pytests/test_client_commands.py @@ -666,10 +666,13 @@ def test_execute_from_file_stops_at_disallowed_special_command(tmp_path: Path) - client.destructive_keywords = set() client.sqlexecute = FakeSQLExecute() - assert result_statuses(client.execute_from_file(f'--special {sql_file}')) == [ + results = list(client.execute_from_file(f'--special {sql_file}')) + + assert result_statuses(iter(results)) == [ 'ran select 1;', 'Special command is never permitted in source files: /pager.', ] + assert results[-1].is_error is True assert client.sqlexecute.runs == ['select 1;'] @@ -798,7 +801,7 @@ def test_execute_from_file_rejects_special_commands(command: str, tmp_path: Path client.sqlexecute = FakeSQLExecute() assert list(client.execute_from_file(str(sql_file))) == [ - SQLResult(status='Special commands are not supported without /source --special.') + SQLResult(status='Special commands are not supported without /source --special.', is_error=True) ] assert client.sqlexecute.runs == [] diff --git a/test/pytests/test_client_query.py b/test/pytests/test_client_query.py index f89e265f7..3e2108cfb 100644 --- a/test/pytests/test_client_query.py +++ b/test/pytests/test_client_query.py @@ -1,6 +1,8 @@ from types import SimpleNamespace from typing import Any, cast +import pytest + from mycli import client_query, main from mycli.packages.sqlresult import SQLResult from mycli.types import Query @@ -305,6 +307,45 @@ def test_run_query_writes_checkpoint(monkeypatch, tmp_path) -> None: assert state['checkpoint_path'].read_text(encoding='utf-8') == 'select 1;\n' +def test_run_query_raises_for_error_result_when_requested(tmp_path) -> None: + cli = make_bare_mycli() + logged_output: list[str] = [] + checkpoint_path = tmp_path / 'checkpoint.sql' + cli.sqlexecute = SimpleNamespace(run=lambda query: [SQLResult(status='source failed', is_error=True)]) + cli.log_query = lambda query: None + cli.log_output = logged_output.append + + with pytest.raises(client_query.QueryError, match='source failed'): + main.MyCli.run_query( + cli, + '/source test.sql', + checkpoint=str(checkpoint_path), + raise_on_error=True, + ) + + assert logged_output == ['source failed'] + assert checkpoint_path.read_text(encoding='utf-8') == '' + cli.checkpoint.close() + + +def test_run_query_displays_error_result_by_default(monkeypatch) -> None: + cli = make_bare_mycli() + result = SQLResult(status='source failed', is_error=True) + echoed: list[str] = [] + cli.sqlexecute = SimpleNamespace(run=lambda query: [result]) + cli.log_query = lambda query: None + cli.log_output = lambda line: None + cli.format_sqlresult = lambda result, **kwargs: [result.status_plain] + monkeypatch.setattr(client_query.special, 'is_expanded_output', lambda: False) + monkeypatch.setattr(client_query.special, 'is_redirected', lambda: False) + monkeypatch.setattr(client_query.special, 'is_show_warnings_enabled', lambda: False) + monkeypatch.setattr(client_query.click, 'echo', lambda line, nl=True: echoed.append(line)) + + main.MyCli.run_query(cli, '/source test.sql') + + assert echoed == ['source failed'] + + def test_run_query_displays_set_buffer_fallback_outside_repl(monkeypatch) -> None: cli = make_bare_mycli() status = 'Error: /favorite eval is only available in the interactive REPL.' diff --git a/test/pytests/test_main.py b/test/pytests/test_main.py index f6ad0ec37..25c3fc0b1 100644 --- a/test/pytests/test_main.py +++ b/test/pytests/test_main.py @@ -987,7 +987,7 @@ def __init__(self, **_args): def connect(self, **args): MockMyCli.connect_args = args - def run_query(self, query, checkpoint=None, new_line=True): + def run_query(self, query, checkpoint=None, new_line=True, raise_on_error=False): return [] def close(self, **args): @@ -2058,7 +2058,7 @@ def __init__(self, **_args): def connect(self, **_args): MockMyCli.connect_calls += 1 - def run_query(self, query, checkpoint=None, new_line=True): + def run_query(self, query, checkpoint=None, new_line=True, raise_on_error=False): MockMyCli.ran_queries.append(query) def run_cli(self): diff --git a/test/pytests/test_main_modes_execute.py b/test/pytests/test_main_modes_execute.py index 2dd9f82a3..01662be06 100644 --- a/test/pytests/test_main_modes_execute.py +++ b/test/pytests/test_main_modes_execute.py @@ -36,14 +36,14 @@ class DummyMyCli: def __init__(self, run_query_error: Exception | None = None) -> None: self.main_formatter = DummyFormatter() self.run_query_error = run_query_error - self.ran_queries: list[tuple[str, str | None]] = [] + self.ran_queries: list[tuple[str, str | None, bool]] = [] self.destructive_keywords = ['drop'] self.logger = DummyLogger() - def run_query(self, query: str, checkpoint: str | None = None) -> None: + def run_query(self, query: str, checkpoint: str | None = None, raise_on_error: bool = False) -> None: if self.run_query_error is not None: raise self.run_query_error - self.ran_queries.append((query, checkpoint)) + self.ran_queries.append((query, checkpoint, raise_on_error)) def main_execute_from_cli(mycli: DummyMyCli, cli_args: DummyCliArgs) -> int: @@ -94,7 +94,7 @@ def test_main_execute_from_cli_sets_format_and_runs_query( assert result == 0 assert mycli.main_formatter.format_name == expected_format - assert mycli.ran_queries == [(expected_sql, 'cp')] + assert mycli.ran_queries == [(expected_sql, 'cp', True)] assert secho_calls == [ ('Ignoring STDIN since --execute was also given.', True, 'red'), ('Ignoring --batch since --execute was also given.', True, 'red'), @@ -116,7 +116,7 @@ def test_main_execute_from_cli_does_not_warn_when_stdin_is_tty_and_batch_is_unse assert result == 0 assert mycli.main_formatter.format_name == 'csv' - assert mycli.ran_queries == [('select 1', None)] + assert mycli.ran_queries == [('select 1', None, True)] assert secho_calls == [] @@ -159,7 +159,7 @@ def confirm_destructive_query(keywords: list[str], query: str) -> bool: assert result == 0 assert execute_mode.sys.stdin is tty assert confirm_calls == [(['drop'], 'drop table t')] - assert mycli.ran_queries == [('drop table t', None)] + assert mycli.ran_queries == [('drop table t', None, True)] def test_main_execute_from_cli_returns_error_when_destructive_query_is_rejected(monkeypatch) -> None: diff --git a/test/pytests/test_sqlresult.py b/test/pytests/test_sqlresult.py index 9c19293ab..2ea776697 100644 --- a/test/pytests/test_sqlresult.py +++ b/test/pytests/test_sqlresult.py @@ -11,6 +11,7 @@ def test_sqlresult_str_includes_all_fields() -> None: postamble='after', status='ok', command={'name': 'watch', 'seconds': 1.0}, + is_error=True, ) assert 'before' in str(result) @@ -19,6 +20,7 @@ def test_sqlresult_str_includes_all_fields() -> None: assert 'after' in str(result) assert 'ok' in str(result) assert "{'name': 'watch', 'seconds': 1.0}" in str(result) + assert 'True' in str(result) def test_sqlresult_status_plain_handles_none_and_formatted_text() -> None: