diff --git a/changelog.md b/changelog.md index 249659d8..2364af42 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,11 @@ Upcoming (TBD) ============== +Bug Fixes +-------- +* Allow shell-style redirects with `/source` when the filename is unquoted. + + Documentation -------- * Badge color nit in `README.md`. diff --git a/mycli/client_commands.py b/mycli/client_commands.py index 938a0336..de90a0a2 100644 --- a/mycli/client_commands.py +++ b/mycli/client_commands.py @@ -4,13 +4,10 @@ import logging import os import re -import shlex from typing import TYPE_CHECKING, Any, cast import click -import sqlparse -from mycli.compat import WIN from mycli.config import write_default_config from mycli.main_modes.repl import set_all_external_titles from mycli.packages import special @@ -18,9 +15,12 @@ from mycli.packages.filepaths import dir_path_exists from mycli.packages.interactive_utils import confirm_destructive_query from mycli.packages.ptoolkit.history import FileHistoryWithTimestamp -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.special.source import ( + parse_source_arguments, + parse_source_filename, + source_special_command_is_safe, +) from mycli.packages.sqlresult import SQLResult from mycli.sqlexecute import SQLExecute @@ -36,30 +36,6 @@ DSN_CONFIG_VALUE = object() FAVORITES_CONFIG_VALUE = object() HIDDEN_CONFIG_SECTIONS = frozenset({'alias_dsn', 'favorite_queries'}) -INVALID_SOURCE_FILENAME = 'Source accepts exactly one filename; filenames containing spaces must be quoted.' -SOURCE_SAFE_SPECIAL_COMMANDS = frozenset({ - 'connect', - 'fd', - 'fs', - 'help', - 'l', - 'nowarnings', - 'ping', - '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: @@ -68,106 +44,6 @@ def _render_config_value(value: Any) -> str: return str(value) -def _parse_source_arguments(arg: str) -> tuple[str, bool, bool, bool]: - allow_special = False - show_queries = False - page_output = False - filename = arg - while arguments := filename.split(maxsplit=1): - if arguments[0] == '--special': - allow_special = True - elif arguments[0] == '--show': - show_queries = True - elif arguments[0] == '--page': - page_output = True - else: - break - filename = arguments[1] if len(arguments) == 2 else '' - return filename, allow_special, show_queries, page_output - - -def _has_unquoted_whitespace(value: str) -> bool: - quote: str | None = None - escaped = False - for character in value: - if escaped: - if quote is None and character.isspace(): - return True - escaped = False - continue - if not WIN and character == '\\' and quote != "'": - escaped = True - continue - if character in ("'", '"'): - if quote is None: - quote = character - elif quote == character: - quote = None - elif quote is None and character.isspace(): - return True - return False - - -def _parse_source_filename(filename: str) -> str: - if not filename: - return '' - if _has_unquoted_whitespace(filename): - raise ValueError(INVALID_SOURCE_FILENAME) - try: - arguments = shlex.split(filename, posix=not WIN) - except ValueError as error: - raise ValueError(f'Invalid source filename: {error}.') from None - if len(arguments) != 1: - raise ValueError(INVALID_SOURCE_FILENAME) - parsed_filename = arguments[0] - if WIN and len(parsed_filename) >= 2 and parsed_filename[0] == parsed_filename[-1] and parsed_filename[0] in ("'", '"'): - parsed_filename = parsed_filename[1:-1] - return parsed_filename - - -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 = '', @@ -412,11 +288,11 @@ 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]: - filename, allow_special, show_queries, page_output = _parse_source_arguments(arg) + filename, allow_special, show_queries, page_output = parse_source_arguments(arg) if page_output: yield SQLResult(command={'name': 'source_page'}) try: - filename = _parse_source_filename(filename) + filename = parse_source_filename(filename) except ValueError as error: yield SQLResult(status=str(error), is_error=True) return @@ -450,7 +326,7 @@ def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]: is_error=True, ) return - if not _source_special_command_is_safe(special_query): + 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}.', diff --git a/mycli/packages/hybrid_redirection.py b/mycli/packages/hybrid_redirection.py index 7840b4da..14dd310b 100644 --- a/mycli/packages/hybrid_redirection.py +++ b/mycli/packages/hybrid_redirection.py @@ -1,14 +1,34 @@ import functools import logging +import re import shlex import sqlglot from mycli.compat import WIN from mycli.packages.special.delimitercommand import DelimiterCommand +from mycli.packages.special.source import ( + parse_source_arguments, + parse_source_filename, +) logger = logging.getLogger(__name__) delimiter_command = DelimiterCommand() +SOURCE_COMMAND_PATTERN = re.compile(r'^([/]?source|[/\\]\.)\s+', re.IGNORECASE) +SOURCE_OPTIONS_PATTERN = re.compile( + r'^([/]?source|[/\\]\.)\s+(?P(?:(?:--special|--show|--page)\s+)*)', + re.IGNORECASE, +) + + +def tokenize_command(command: str) -> list[sqlglot.Token]: + """Tokenize a command without treating source options as SQL comments.""" + source_match = SOURCE_OPTIONS_PATTERN.match(command) + if source_match: + options_start, options_end = source_match.span('options') + options = command[options_start:options_end].replace('-', '_') + command = command[:options_start] + options + command[options_end:] + return sqlglot.tokenize(command) def find_token_indices(tokens: list[sqlglot.Token]) -> dict[str, list[int]]: @@ -44,6 +64,16 @@ def find_sql_part( ): leftmost_dollar_pos = tokens[true_dollar_indices[0]].start sql_part = command[0:leftmost_dollar_pos].strip().removesuffix(delimiter_command.current).rstrip() + if SOURCE_COMMAND_PATTERN.match(sql_part): + source_arg_str = SOURCE_COMMAND_PATTERN.sub('', sql_part) + try: + filename, _allow_special, _show_queries, _page_output = parse_source_arguments(source_arg_str) + filename = parse_source_filename(filename) + except ValueError: + return '' + if not filename: + return '' + return sql_part try: statements = sqlglot.parse(sql_part, read='mysql') except sqlglot.errors.ParseError: @@ -142,7 +172,7 @@ def get_redirect_components(command: str) -> tuple[str | None, str | None, str | """Get the parts of a hybrid shell-style redirect command.""" try: - tokens = sqlglot.tokenize(command) + tokens = tokenize_command(command) except sqlglot.errors.TokenError: return None, None, None, None diff --git a/mycli/packages/special/__init__.py b/mycli/packages/special/__init__.py index 361b3752..fb8afca9 100644 --- a/mycli/packages/special/__init__.py +++ b/mycli/packages/special/__init__.py @@ -3,6 +3,7 @@ 'mycli.packages.special.iocommands', 'mycli.packages.special.llm', 'mycli.packages.special.main', + 'mycli.packages.special.source', ] import os @@ -54,6 +55,10 @@ write_pipe_once, write_tee, ) +from mycli.packages.special.source import ( + parse_source_arguments, + parse_source_filename, +) if not os.environ.get('MYCLI_LLM_OFF'): from mycli.packages.special.llm import ( @@ -112,12 +117,15 @@ def sql_using_llm(*args, **kwargs): # type: ignore[no-redef, misc] 'is_llm_command', 'is_pager_enabled', 'is_redirected', - 'is_special_command', + 'is_show_favorite_query', 'is_show_warnings_enabled', + 'is_special_command', 'is_timing_enabled', 'list_databases', 'list_tables', 'open_external_editor', + 'parse_source_arguments', + 'parse_source_filename', 'parse_special_command', 'ping', 'register_special_command', @@ -131,10 +139,9 @@ def sql_using_llm(*args, **kwargs): # type: ignore[no-redef, misc] 'set_pager', 'set_pager_enabled', 'set_redirect', + 'set_show_favorite_query', 'set_show_warnings_enabled', 'set_timing_enabled', - 'set_show_favorite_query', - 'is_show_favorite_query', 'special_command', 'split_queries', 'sql_using_llm', diff --git a/mycli/packages/special/source.py b/mycli/packages/special/source.py new file mode 100644 index 00000000..f45d4d65 --- /dev/null +++ b/mycli/packages/special/source.py @@ -0,0 +1,133 @@ +import shlex + +import sqlparse + +from mycli.compat import WIN +from mycli.packages import special +from mycli.packages.special import main as special_main +from mycli.packages.special.iocommands import expand_favorite_query + +INVALID_SOURCE_FILENAME = 'Source accepts exactly one filename; filenames containing spaces must be quoted.' +SOURCE_SAFE_SPECIAL_COMMANDS = frozenset({ + 'connect', + 'fd', + 'fs', + 'help', + 'l', + 'nowarnings', + 'ping', + '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 _has_unquoted_whitespace(value: str) -> bool: + quote: str | None = None + escaped = False + for character in value: + if escaped: + if quote is None and character.isspace(): + return True + escaped = False + continue + if not WIN and character == '\\' and quote != "'": + escaped = True + continue + if character in ("'", '"'): + if quote is None: + quote = character + elif quote == character: + quote = None + elif quote is None and character.isspace(): + return True + return 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 parse_source_arguments(arg: str) -> tuple[str, bool, bool, bool]: + allow_special = False + show_queries = False + page_output = False + filename = arg + while arguments := filename.split(maxsplit=1): + if arguments[0] == '--special': + allow_special = True + elif arguments[0] == '--show': + show_queries = True + elif arguments[0] == '--page': + page_output = True + else: + break + filename = arguments[1] if len(arguments) == 2 else '' + return filename, allow_special, show_queries, page_output + + +def parse_source_filename(filename: str) -> str: + if not filename: + return '' + if _has_unquoted_whitespace(filename): + raise ValueError(INVALID_SOURCE_FILENAME) + try: + arguments = shlex.split(filename, posix=not WIN) + except ValueError as error: + raise ValueError(f'Invalid source filename: {error}.') from None + if len(arguments) != 1: + raise ValueError(INVALID_SOURCE_FILENAME) + parsed_filename = arguments[0] + if WIN and len(parsed_filename) >= 2 and parsed_filename[0] == parsed_filename[-1] and parsed_filename[0] in ("'", '"'): + parsed_filename = parsed_filename[1:-1] + return parsed_filename + + +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 diff --git a/test/pytests/test_client_commands.py b/test/pytests/test_client_commands.py index ca5bed49..2cbd3491 100644 --- a/test/pytests/test_client_commands.py +++ b/test/pytests/test_client_commands.py @@ -14,6 +14,7 @@ from mycli.client_commands import ClientCommandsMixin from mycli.packages import special from mycli.packages.special import main as special_main +from mycli.packages.special import source as source_commands from mycli.packages.sqlresult import SQLResult @@ -96,73 +97,6 @@ def result_statuses(results: Any) -> list[str | None]: return [result.status for result in list(results)] -@pytest.mark.parametrize( - ('arg', 'expected'), - [ - ('query.sql', ('query.sql', False, False, False)), - ('--special query.sql', ('query.sql', True, False, False)), - ('--show query.sql', ('query.sql', False, True, False)), - ('--page query.sql', ('query.sql', False, False, True)), - ('--special --show --page query file.sql', ('query file.sql', True, True, True)), - ('--page --show --special query file.sql', ('query file.sql', True, True, True)), - ('--show --show query.sql', ('query.sql', False, True, False)), - ('--page --page query.sql', ('query.sql', False, False, True)), - ('--show', ('', False, True, False)), - ], -) -def test_parse_source_arguments(arg: str, expected: tuple[str, bool, bool, bool]) -> None: - assert client_commands._parse_source_arguments(arg) == expected - - -@pytest.mark.parametrize( - ('filename', 'expected'), - [ - ('query.sql', 'query.sql'), - ('"query file.sql"', 'query file.sql'), - ("'query file.sql'", 'query file.sql'), - ('prefix" query".sql', 'prefix query.sql'), - ], -) -def test_parse_source_filename(filename: str, expected: str) -> None: - assert client_commands._parse_source_filename(filename) == expected - - -@pytest.mark.parametrize( - 'filename', - [ - 'query file.sql', - r'query\ file.sql', - '"first file.sql" second.sql', - ], -) -def test_parse_source_filename_rejects_multiple_unquoted_arguments(filename: str) -> None: - with pytest.raises(ValueError, match='filenames containing spaces must be quoted'): - client_commands._parse_source_filename(filename) - - -def test_parse_source_filename_rejects_unclosed_quote() -> None: - with pytest.raises(ValueError, match='No closing quotation'): - client_commands._parse_source_filename('"query file.sql') - - -def test_parse_source_filename_rejects_missing_parsed_argument(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(client_commands.shlex, 'split', lambda *_args, **_kwargs: []) - - with pytest.raises(ValueError, match='accepts exactly one filename'): - client_commands._parse_source_filename('query.sql') - - -def test_source_filename_whitespace_scanner_allows_escaped_non_whitespace() -> None: - assert not client_commands._has_unquoted_whitespace(r'query\name.sql') - - -def test_parse_source_filename_preserves_windows_backslashes(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(client_commands, 'WIN', True) - - assert client_commands._parse_source_filename(r'C:\queries\query.sql') == r'C:\queries\query.sql' - assert client_commands._parse_source_filename(r'"C:\my queries\query.sql"') == r'C:\my queries\query.sql' - - def test_register_special_commands_registers_expected_commands(monkeypatch: pytest.MonkeyPatch) -> None: client = DummyClient() calls: list[tuple[Any, ...]] = [] @@ -801,7 +735,7 @@ def test_execute_from_file_rejects_unquoted_filename_with_spaces(monkeypatch: py opened_paths: list[str] = [] monkeypatch.setattr(client_commands, 'open', lambda path: opened_paths.append(path), raising=False) - assert list(client.execute_from_file('query file.sql')) == [SQLResult(status=client_commands.INVALID_SOURCE_FILENAME, is_error=True)] + assert list(client.execute_from_file('query file.sql')) == [SQLResult(status=source_commands.INVALID_SOURCE_FILENAME, is_error=True)] assert opened_paths == [] @@ -810,7 +744,7 @@ def test_execute_from_file_pages_invalid_filename_error() -> None: assert list(client.execute_from_file('--page query file.sql')) == [ SQLResult(command={'name': 'source_page'}), - SQLResult(status=client_commands.INVALID_SOURCE_FILENAME, is_error=True), + SQLResult(status=source_commands.INVALID_SOURCE_FILENAME, is_error=True), ] @@ -887,101 +821,6 @@ def test_execute_from_file_requires_semicolon_for_special_commands(tmp_path: Pat assert client.sqlexecute.runs == ['/status\nselect 1;'] -@pytest.mark.parametrize( - ('command', 'arg', 'expected'), - [ - ('status', '', True), - ('ping', '', 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', [ diff --git a/test/pytests/test_hybrid_redirection.py b/test/pytests/test_hybrid_redirection.py index 608217a6..6922d3cf 100644 --- a/test/pytests/test_hybrid_redirection.py +++ b/test/pytests/test_hybrid_redirection.py @@ -54,6 +54,22 @@ def test_find_sql_part_handles_valid_parse_custom_delimiter_and_invalid_sql(rese assert hybrid_redirection.find_sql_part('select 1; select 2 $> out.txt', multiple_tokens, [5]) == '' +@pytest.mark.parametrize( + ('command', 'expected'), + [ + ('/source query.sql $> out.txt', '/source query.sql'), + ('source "query file.sql" $| cat', 'source "query file.sql"'), + ('/source query file.sql $> out.txt', ''), + ('/source "" $> out.txt', ''), + ], +) +def test_find_sql_part_handles_source_commands(command: str, expected: str) -> None: + tokens = tokenize(command) + indices = hybrid_redirection.find_token_indices(tokens) + + assert hybrid_redirection.find_sql_part(command, tokens, indices['true_dollar']) == expected + + def test_find_command_and_file_tokens_extract_expected_parts() -> None: tokens = tokenize('select 1 $| cat $>> out.txt') indices = hybrid_redirection.find_token_indices(tokens) @@ -149,6 +165,29 @@ def test_get_redirect_components_valid_paths_and_logging() -> None: ) +@pytest.mark.parametrize('option', ['--special', '--show', '--page']) +def test_get_redirect_components_preserves_source_options(option: str) -> None: + command = f'/source {option} query.sql $> out.txt' + + assert hybrid_redirection.get_redirect_components(command) == ( + f'/source {option} query.sql', + None, + '>', + 'out.txt', + ) + + +def test_get_redirect_components_handles_combined_source_options_and_redirects() -> None: + command = '/source --page --show --special query.sql $| cat $>> out.txt' + + assert hybrid_redirection.get_redirect_components(command) == ( + '/source --page --show --special query.sql', + 'cat', + '>>', + 'out.txt', + ) + + def test_get_redirect_components_returns_none_on_token_error(monkeypatch) -> None: monkeypatch.setattr( hybrid_redirection.sqlglot, 'tokenize', lambda command: (_ for _ in ()).throw(sqlglot.errors.TokenError('bad token')) diff --git a/test/pytests/test_special_source.py b/test/pytests/test_special_source.py new file mode 100644 index 00000000..fe2c8d25 --- /dev/null +++ b/test/pytests/test_special_source.py @@ -0,0 +1,167 @@ +import pytest + +from mycli.packages.special import main as special_main +from mycli.packages.special import source + + +@pytest.mark.parametrize( + ('arg', 'expected'), + [ + ('query.sql', ('query.sql', False, False, False)), + ('--special query.sql', ('query.sql', True, False, False)), + ('--show query.sql', ('query.sql', False, True, False)), + ('--page query.sql', ('query.sql', False, False, True)), + ('--special --show --page query file.sql', ('query file.sql', True, True, True)), + ('--page --show --special query file.sql', ('query file.sql', True, True, True)), + ('--show --show query.sql', ('query.sql', False, True, False)), + ('--page --page query.sql', ('query.sql', False, False, True)), + ('--show', ('', False, True, False)), + ], +) +def test_parse_source_arguments(arg: str, expected: tuple[str, bool, bool, bool]) -> None: + assert source.parse_source_arguments(arg) == expected + + +@pytest.mark.parametrize( + ('filename', 'expected'), + [ + ('', ''), + ('query.sql', 'query.sql'), + ('"query file.sql"', 'query file.sql'), + ("'query file.sql'", 'query file.sql'), + ('prefix" query".sql', 'prefix query.sql'), + ], +) +def test_parse_source_filename(filename: str, expected: str) -> None: + assert source.parse_source_filename(filename) == expected + + +@pytest.mark.parametrize( + 'filename', + [ + 'query file.sql', + r'query\ file.sql', + '"first file.sql" second.sql', + ], +) +def test_parse_source_filename_rejects_multiple_unquoted_arguments(filename: str) -> None: + with pytest.raises(ValueError, match='filenames containing spaces must be quoted'): + source.parse_source_filename(filename) + + +def test_parse_source_filename_rejects_unclosed_quote() -> None: + with pytest.raises(ValueError, match='No closing quotation'): + source.parse_source_filename('"query file.sql') + + +def test_parse_source_filename_rejects_missing_parsed_argument(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(source.shlex, 'split', lambda *_args, **_kwargs: []) + + with pytest.raises(ValueError, match='accepts exactly one filename'): + source.parse_source_filename('query.sql') + + +def test_source_filename_whitespace_scanner_allows_escaped_non_whitespace() -> None: + assert not source._has_unquoted_whitespace(r'query\name.sql') + + +def test_parse_source_filename_preserves_windows_backslashes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(source, 'WIN', True) + + assert source.parse_source_filename(r'C:\queries\query.sql') == r'C:\queries\query.sql' + assert source.parse_source_filename(r'"C:\my queries\query.sql"') == r'C:\my queries\query.sql' + + +@pytest.mark.parametrize( + ('command', 'arg', 'expected'), + [ + ('status', '', True), + ('ping', '', 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(source, '_registered_special_command', lambda query: (command, arg)) + + assert source.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 source._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 source._registered_special_command('/unknown') is None + assert source.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(source, '_registered_special_command', lambda query: (command, 'report')) + + assert source.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( + source, + 'expand_favorite_query', + lambda arg: (expanded_query, None if expanded_query is not None else 'invalid arguments'), + ) + + assert source._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(source, '_registered_special_command', lambda query: (command, arg)) + monkeypatch.setattr(source, '_favorite_source_command_is_safe', lambda favorite_arg: False) + + assert source.source_special_command_is_safe('/favorite') is False