diff --git a/changelog.md b/changelog.md index e6ace488c..084ad72b4 100644 --- a/changelog.md +++ b/changelog.md @@ -4,6 +4,7 @@ Upcoming (TBD) Features --------- * Advertise beta `--boundary-id` option in helpdoc and completions. +* Add `/source --special` to allow executing some special commands. 2.15.0 (2026/08/20) diff --git a/mycli/client_commands.py b/mycli/client_commands.py index 26a92df03..d12e476e4 100644 --- a/mycli/client_commands.py +++ b/mycli/client_commands.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, cast import click +import sqlparse from mycli.config import write_default_config from mycli.main_modes.repl import set_all_external_titles @@ -14,6 +15,8 @@ from mycli.packages.batch_utils import statements_from_filehandle from mycli.packages.filepaths import dir_path_exists from mycli.packages.interactive_utils import confirm_destructive_query +from mycli.packages.special import main as special_main +from mycli.packages.special.iocommands import expand_favorite_query from mycli.packages.special.main import ArgType, SpecialCommandAlias from mycli.packages.sqlresult import SQLResult from mycli.sqlexecute import SQLExecute @@ -30,6 +33,28 @@ DSN_CONFIG_VALUE = object() FAVORITES_CONFIG_VALUE = object() HIDDEN_CONFIG_SECTIONS = frozenset({'alias_dsn', 'favorite_queries'}) +SOURCE_SAFE_SPECIAL_COMMANDS = frozenset({ + 'connect', + 'fd', + 'fs', + 'help', + 'l', + 'nowarnings', + 'prompt', + 'redirectformat', + 'rehash', + 'status', + 'tableformat', + 'timing', + 'use', + 'warnings', + 'dt', +}) +SOURCE_SAFE_SUBCOMMANDS = { + 'config': frozenset({'help', 'get', 'search'}), + 'dsn': frozenset({'help', 'list', 'show', 'save', 'delete'}), + 'favorite': frozenset({'help', 'list', 'reload', 'run', 'save', 'delete'}), +} def _render_config_value(value: Any) -> str: @@ -38,6 +63,56 @@ def _render_config_value(value: Any) -> str: return str(value) +def _parse_source_arguments(arg: str) -> tuple[str, bool]: + arguments = arg.split(maxsplit=1) + if arguments and arguments[0] == '--special': + return (arguments[1] if len(arguments) == 2 else '', True) + return (arg, False) + + +def _registered_special_command(query: str) -> tuple[str, str] | None: + command, _verbosity, arg = special.parse_special_command(query) + registered = special_main.COMMANDS.get(command) + if registered is None: + registered = special_main.COMMANDS.get(command.lower()) + if registered is None: + return None + return registered.command.removeprefix('\\').removeprefix('/').lower(), arg + + +def _favorite_source_command_is_safe(arg: str) -> bool: + query, _error = expand_favorite_query(arg) + if query is None: + return True + return not any(special.is_special_command(statement.rstrip(';')) for statement in sqlparse.split(query)) + + +def _source_special_command_is_safe(query: str) -> bool: + parsed = _registered_special_command(query) + if parsed is None: + return False + + command, arg = parsed + if command == 'f': + return not arg or _favorite_source_command_is_safe(arg) + if command in ('fd', 'fs'): + return True + if command in SOURCE_SAFE_SPECIAL_COMMANDS: + return True + + subcommands = SOURCE_SAFE_SUBCOMMANDS.get(command) + if subcommands is None: + return False + arguments = arg.split(maxsplit=1) + subcommand = arguments[0].lower() if arguments else 'help' + if subcommand not in subcommands: + return False + if command == 'favorite' and subcommand == 'run': + run_arg = arguments[1] if len(arguments) == 2 else '' + return not run_arg or _favorite_source_command_is_safe(run_arg) + return True + + def _iter_config_values( config: Mapping[str, Any], prefix: str = '', @@ -143,7 +218,7 @@ def register_special_commands(self) -> None: special.register_special_command( self.execute_from_file, "source", - "/source ", + "/source [--special] ", "Execute queries from a file.", aliases=[SpecialCommandAlias("\\.", case_sensitive=False)], ) @@ -266,12 +341,13 @@ def change_db(self, arg: str, **_) -> Generator[SQLResult, None, None]: yield SQLResult(status=msg) def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]: - if not arg: + filename, allow_special = _parse_source_arguments(arg) + if not filename: yield SQLResult(status="Missing required argument: filename.") return try: - file_h = open(os.path.expanduser(arg)) + file_h = open(os.path.expanduser(filename)) except OSError as error: yield SQLResult(status=str(error)) return @@ -288,9 +364,17 @@ def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]: yield SQLResult(status=str(error)) return - if special.is_special_command(query.rstrip(';')): - yield SQLResult(status='Special commands are not supported in source files.') - return + 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.') + 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}.') + return + yield from self.sqlexecute.run(special_query) + continue if self.destructive_warning and confirm_destructive_query(self.destructive_keywords, query) is False: continue diff --git a/mycli/packages/completion_engine.py b/mycli/packages/completion_engine.py index 716ce959d..534e67737 100644 --- a/mycli/packages/completion_engine.py +++ b/mycli/packages/completion_engine.py @@ -811,6 +811,20 @@ def suggest_special(text: str) -> list[dict[str, Any]]: r'/.', 'source', '/source', + ]: + source_arguments = _arg.split(maxsplit=1) + if not source_arguments: + return [ + {'type': 'special_subcommand', 'subcommands': ['--special']}, + {'type': 'file_name'}, + ] + if source_arguments[0].startswith('-') and source_arguments[0] != '--special': + return [{'type': 'special_subcommand', 'subcommands': ['--special']}] + if source_arguments[0] == '--special' and len(source_arguments) == 1 and not text[-1].isspace(): + return [] + return [{'type': 'file_name'}] + + if cmd.lower() in [ r'\o', '/o', r'\once', diff --git a/test/features/fixture_data/help_commands.txt b/test/features/fixture_data/help_commands.txt index b17340992..44d68a22a 100644 --- a/test/features/fixture_data/help_commands.txt +++ b/test/features/fixture_data/help_commands.txt @@ -29,7 +29,7 @@ | /quit | /q | /quit | Quit. | | /redirectformat | /Tr | /redirectformat | Change the table format used to output redirected results. | | /rehash | /# | /rehash | Refresh auto-completions. | -| /source | /. | /source | Execute queries from a file. | +| /source | /. | /source [--special] | Execute queries from a file. | | /status | /s | /status | Get status information from the server. | | /system | | /system [-r] | Execute a system shell command (raw mode with -r). | | /tableformat | /T | /tableformat | Change the table format used to output interactive results. | diff --git a/test/pytests/test_client_commands.py b/test/pytests/test_client_commands.py index c3cca4413..3cd6301e3 100644 --- a/test/pytests/test_client_commands.py +++ b/test/pytests/test_client_commands.py @@ -118,6 +118,7 @@ def test_register_special_commands_registers_expected_commands(monkeypatch: pyte assert calls[3][0] == client.change_table_format assert calls[4][0] == client.change_redirect_format assert calls[5][0] == client.execute_from_file + assert calls[5][2:4] == ('/source [--special] ', 'Execute queries from a file.') assert calls[6][0] == client.change_prompt_format assert calls[6][2:4] == ('/prompt [string]', 'Show or change prompt format.') assert calls[7][0] == client.config_command @@ -616,6 +617,168 @@ def test_execute_from_file_runs_file_query(tmp_path: Path) -> None: assert client.sqlexecute.runs == ['select 1;'] +def test_execute_from_file_parses_special_option_and_preserves_filename( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = DummyClient() + client.destructive_warning = False + client.sqlexecute = FakeSQLExecute() + file_h = IteratedFile('select 1;') + opened_paths: list[str] = [] + + def open_file(path: str) -> IteratedFile: + opened_paths.append(path) + return file_h + + monkeypatch.setattr(client_commands, 'open', open_file, raising=False) + + assert result_statuses(client.execute_from_file('--special query file.sql')) == ['ran select 1;'] + assert opened_paths == ['query file.sql'] + + +def test_execute_from_file_reports_missing_filename_after_special_option() -> None: + client = DummyClient() + + assert list(client.execute_from_file('--special')) == [SQLResult(status='Missing required argument: filename.')] + + +def test_execute_from_file_runs_permitted_special_commands(tmp_path: Path) -> None: + client = DummyClient() + sql_file = tmp_path / 'query.sql' + sql_file.write_text('select 1; /status; select 2;', encoding='utf-8') + client.destructive_warning = False + client.destructive_keywords = set() + client.sqlexecute = FakeSQLExecute() + + assert result_statuses(client.execute_from_file(f'--special {sql_file}')) == [ + 'ran select 1;', + 'ran /status', + 'ran select 2;', + ] + assert client.sqlexecute.runs == ['select 1;', '/status', 'select 2;'] + + +def test_execute_from_file_stops_at_disallowed_special_command(tmp_path: Path) -> None: + client = DummyClient() + sql_file = tmp_path / 'query.sql' + sql_file.write_text('select 1; /pager; select 2;', encoding='utf-8') + client.destructive_warning = False + client.destructive_keywords = set() + client.sqlexecute = FakeSQLExecute() + + assert result_statuses(client.execute_from_file(f'--special {sql_file}')) == [ + 'ran select 1;', + 'Special command is never permitted in source files: /pager.', + ] + assert client.sqlexecute.runs == ['select 1;'] + + +def test_execute_from_file_requires_semicolon_for_special_commands(tmp_path: Path) -> None: + client = DummyClient() + sql_file = tmp_path / 'query.sql' + sql_file.write_text('/status\nselect 1;', encoding='utf-8') + client.destructive_warning = False + client.destructive_keywords = set() + client.sqlexecute = FakeSQLExecute() + + assert result_statuses(client.execute_from_file(f'--special {sql_file}')) == ['ran /status\nselect 1;'] + assert client.sqlexecute.runs == ['/status\nselect 1;'] + + +@pytest.mark.parametrize( + ('command', 'arg', 'expected'), + [ + ('status', '', True), + ('connect', 'db', True), + ('config', 'get main.prompt', True), + ('config', 'edit', False), + ('dsn', 'list', True), + ('dsn', 'edit prod', False), + ('favorite', 'list', True), + ('favorite', 'eval report', False), + ('pager', '', False), + ('delimiter', '$$', False), + ('plugin_command', '', False), + ], +) +def test_source_special_command_policy( + monkeypatch: pytest.MonkeyPatch, + command: str, + arg: str, + expected: bool, +) -> None: + monkeypatch.setattr(client_commands, '_registered_special_command', lambda query: (command, arg)) + + assert client_commands._source_special_command_is_safe('/command') is expected + + +def test_registered_source_special_command_uses_case_insensitive_registry_lookup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + registered = special_main.SpecialCommand( + handler=lambda: None, + command='status', + usage='/status', + description='Show status.', + arg_type=special_main.ArgType.NO_ARGUMENT, + hidden=False, + case_sensitive=False, + aliases=None, + backslash_only=False, + ) + monkeypatch.setattr(special_main, 'COMMANDS', {'/status': registered}) + + assert client_commands._registered_special_command('/STATUS verbose') == ('status', 'verbose') + + +def test_registered_source_special_command_rejects_unknown_command(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(special_main, 'COMMANDS', {}) + + assert client_commands._registered_special_command('/unknown') is None + assert client_commands._source_special_command_is_safe('/unknown') is False + + +@pytest.mark.parametrize('command', ['fd', 'fs']) +def test_source_special_command_policy_allows_favorite_aliases( + monkeypatch: pytest.MonkeyPatch, + command: str, +) -> None: + monkeypatch.setattr(client_commands, '_registered_special_command', lambda query: (command, 'report')) + + assert client_commands._source_special_command_is_safe('/command') is True + + +@pytest.mark.parametrize( + ('expanded_query', 'expected'), + [ + ('select 1; select 2;', True), + ('select 1; /system echo unsafe;', False), + (None, True), + ], +) +def test_favorite_source_command_requires_sql_only_expansion( + monkeypatch: pytest.MonkeyPatch, + expanded_query: str | None, + expected: bool, +) -> None: + monkeypatch.setattr( + client_commands, + 'expand_favorite_query', + lambda arg: (expanded_query, None if expanded_query is not None else 'invalid arguments'), + ) + + assert client_commands._favorite_source_command_is_safe('report') is expected + + +@pytest.mark.parametrize('command', ['f', 'favorite']) +def test_source_favorite_run_uses_expansion_policy(monkeypatch: pytest.MonkeyPatch, command: str) -> None: + arg = 'report' if command == 'f' else 'run report' + monkeypatch.setattr(client_commands, '_registered_special_command', lambda query: (command, arg)) + monkeypatch.setattr(client_commands, '_favorite_source_command_is_safe', lambda favorite_arg: False) + + assert client_commands._source_special_command_is_safe('/favorite') is False + + @pytest.mark.parametrize( 'command', [ @@ -634,7 +797,9 @@ def test_execute_from_file_rejects_special_commands(command: str, tmp_path: Path client.destructive_keywords = set() client.sqlexecute = FakeSQLExecute() - assert list(client.execute_from_file(str(sql_file))) == [SQLResult(status='Special commands are not supported in source files.')] + assert list(client.execute_from_file(str(sql_file))) == [ + SQLResult(status='Special commands are not supported without /source --special.') + ] assert client.sqlexecute.runs == [] diff --git a/test/pytests/test_completion_engine.py b/test/pytests/test_completion_engine.py index 8bff2a59d..25ac19beb 100644 --- a/test/pytests/test_completion_engine.py +++ b/test/pytests/test_completion_engine.py @@ -887,8 +887,19 @@ def test_suggest_type_handles_parser_results_shorter_than_cursor(monkeypatch): ('/f report --user="henry', []), ('\\dt ', [{'type': 'table', 'schema': []}, {'type': 'view', 'schema': []}, {'type': 'schema'}]), ('\\dt+ ', [{'type': 'table', 'schema': []}, {'type': 'view', 'schema': []}, {'type': 'schema'}]), - ('\\. ', [{'type': 'file_name'}]), - ('source ', [{'type': 'file_name'}]), + ( + '\\. ', + [{'type': 'special_subcommand', 'subcommands': ['--special']}, {'type': 'file_name'}], + ), + ( + 'source ', + [{'type': 'special_subcommand', 'subcommands': ['--special']}, {'type': 'file_name'}], + ), + ('source --s', [{'type': 'special_subcommand', 'subcommands': ['--special']}]), + ('source --special', []), + ('source --special ', [{'type': 'file_name'}]), + ('source --special query.sql', [{'type': 'file_name'}]), + ('source query.sql', [{'type': 'file_name'}]), ('\\o ', [{'type': 'file_name'}]), ('\\once ', [{'type': 'file_name'}]), ('tee ', [{'type': 'file_name'}]), @@ -1842,7 +1853,10 @@ def test_source_is_file(expression): aliases=[special.SpecialCommandAlias('\\.', case_sensitive=False)], ) suggestions = suggest_type(expression, expression) - assert suggestions == [{"type": "file_name"}] + assert suggestions == [ + {'type': 'special_subcommand', 'subcommands': ['--special']}, + {'type': 'file_name'}, + ] @pytest.mark.parametrize( diff --git a/test/pytests/test_smart_completion_public_schema_only.py b/test/pytests/test_smart_completion_public_schema_only.py index c03ce8d5b..8a6e475a2 100644 --- a/test/pytests/test_smart_completion_public_schema_only.py +++ b/test/pytests/test_smart_completion_public_schema_only.py @@ -720,8 +720,11 @@ def dummy_list_path(dir_name): @pytest.mark.parametrize( "text,expected", [ - ('source ', [('/', 0), ('~', 0), ('.', 0), ('..', 0)]), + ('source ', [('--special', 0), ('/', 0), ('~', 0), ('.', 0), ('..', 0)]), + ('source --s', [('--special', -3)]), + ('source --special ', [('/', 0), ('~', 0), ('.', 0), ('..', 0)]), ("source /", [("dir1", 0), ("file1.sql", 0), ("file2.sql", 0)]), + ('source --special /', [('dir1', 0), ('file1.sql', 0), ('file2.sql', 0)]), ("source /dir1/", [("subdir1", 0), ("subfile1.sql", 0), ("subfile2.sql", 0)]), ("source /dir1/subdir1/", [("lastfile.sql", 0)]), ],