/source --special flag, allowing special /commands from an allowlist

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)

mycli/client_commands.py

@@ -7,13 +7,16 @@
 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
 from mycli.packages import special
 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 <filename>",
+            "/source [--special] <filename>",
             "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

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',

test/features/fixture_data/help_commands.txt

@@ -29,7 +29,7 @@
 | /quit           | /q       | /quit                                  | Quit.                                                       |
 | /redirectformat | /Tr      | /redirectformat <format>               | Change the table format used to output redirected results.  |
 | /rehash         | /#       | /rehash                                | Refresh auto-completions.                                   |
-| /source         | /.       | /source <filename>                     | Execute queries from a file.                                |
+| /source         | /.       | /source [--special] <filename>         | Execute queries from a file.                                |
 | /status         | /s       | /status                                | Get status information from the server.                     |
 | /system         | <null>   | /system [-r] <command>                 | Execute a system shell command (raw mode with -r).          |
 | /tableformat    | /T       | /tableformat <format>                  | Change the table format used to output interactive results. |

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] <filename>', '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 == []