diff --git a/changelog.md b/changelog.md index 7cca3de7..22008173 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,11 @@ Upcoming (TBD) ============== +Features +--------- +* Subcommand completions for the `/dsn` command. + + Bug Fixes --------- * Keep Vault username and password fields from being confused. diff --git a/mycli/packages/completion_engine.py b/mycli/packages/completion_engine.py index a1126887..0ff724d5 100644 --- a/mycli/packages/completion_engine.py +++ b/mycli/packages/completion_engine.py @@ -6,6 +6,7 @@ import sqlparse from sqlparse.sql import Comparison, Identifier, Token, Where +from mycli.packages.special.dsn_aliases import DSN_SUBCOMMANDS from mycli.packages.special.main import COMMANDS as SPECIAL_COMMANDS from mycli.packages.special.main import parse_special_command from mycli.packages.sql_utils import extract_tables, find_prev_keyword, last_word @@ -816,6 +817,17 @@ def suggest_special(text: str) -> list[dict[str, Any]]: if cmd in ["\\llm", "/llm", "\\ai", "/ai"]: return [{"type": "llm"}] + if cmd.lower() in (r'\dsn', '/dsn'): + dsn_arguments = _arg.split(maxsplit=1) + completing_delete_target = (len(dsn_arguments) == 1 and text[-1].isspace()) or (len(dsn_arguments) == 2 and not text[-1].isspace()) + if dsn_arguments and dsn_arguments[0].lower() == 'delete' and completing_delete_target: + return [{'type': 'dsn_alias'}] + if dsn_arguments and dsn_arguments[0].lower() == 'delete' and len(dsn_arguments) == 2: + return [] + if dsn_arguments and dsn_arguments[0].lower() in DSN_SUBCOMMANDS - {'delete'}: + return [] + return [{'type': 'special_subcommand', 'subcommands': list(DSN_SUBCOMMANDS)}] + return [{"type": "keyword"}, {"type": "special"}] diff --git a/mycli/packages/special/dsn_aliases.py b/mycli/packages/special/dsn_aliases.py index 23a47047..7c017e66 100644 --- a/mycli/packages/special/dsn_aliases.py +++ b/mycli/packages/special/dsn_aliases.py @@ -1,5 +1,7 @@ from __future__ import annotations +DSN_SUBCOMMANDS = {'help', 'list', 'show', 'save', 'delete'} + class DsnAliases: section_name: str = 'alias_dsn' diff --git a/mycli/sqlcompleter.py b/mycli/sqlcompleter.py index 025493ac..bf3cb98b 100644 --- a/mycli/sqlcompleter.py +++ b/mycli/sqlcompleter.py @@ -14,6 +14,7 @@ from mycli.packages.completion_engine import is_inside_quotes, suggest_type from mycli.packages.filepaths import complete_path, parse_path, suggest_path from mycli.packages.special import llm +from mycli.packages.special.dsn_aliases import DsnAliases from mycli.packages.special.favoritequeries import FavoriteQueries from mycli.packages.special.main import COMMANDS as SPECIAL_COMMANDS from mycli.packages.sql_utils import extract_columns_from_select, extract_tables, last_word @@ -1686,6 +1687,27 @@ def get_completions( ) completions.extend([(*x, rank) for x in subcommands_m]) + elif suggestion['type'] == 'special_subcommand': + subcommands_m = self.find_matches( + word_before_cursor, + suggestion['subcommands'], + start_only=True, + fuzzy=False, + text_before_cursor=document.text_before_cursor, + ) + completions.extend([(*x, rank) for x in subcommands_m]) + + elif suggestion['type'] == 'dsn_alias': + if hasattr(DsnAliases, 'instance'): + aliases_m = self.find_matches( + word_before_cursor, + DsnAliases.instance.list(), + start_only=True, + fuzzy=False, + text_before_cursor=document.text_before_cursor, + ) + completions.extend([(*x, rank) for x in aliases_m]) + elif suggestion["type"] == "enum_value": enum_values = self.populate_enum_values( suggestion["tables"], diff --git a/test/pytests/test_completion_engine.py b/test/pytests/test_completion_engine.py index e6b4bc89..2330dda5 100644 --- a/test/pytests/test_completion_engine.py +++ b/test/pytests/test_completion_engine.py @@ -7,6 +7,7 @@ from mycli.packages import completion_engine, special from mycli.packages.completion_engine import ( + DSN_SUBCOMMANDS, _aliases, _build_suggest_context, _charset_suggestion, @@ -841,6 +842,12 @@ def fake_parse(text: str): ('\\edit ', [{'type': 'file_name'}]), ('\\llm ', [{'type': 'llm'}]), ('\\ai ', [{'type': 'llm'}]), + ('/dsn ', [{'type': 'special_subcommand', 'subcommands': list(DSN_SUBCOMMANDS)}]), + ('/dsn delete ', [{'type': 'dsn_alias'}]), + ('/dsn delete pro', [{'type': 'dsn_alias'}]), + ('/dsn delete prod ', []), + ('/dsn show', []), + ('/dsn help ', []), ('pager ', [{'type': 'keyword'}, {'type': 'special'}]), ], ) diff --git a/test/pytests/test_smart_completion_public_schema_only.py b/test/pytests/test_smart_completion_public_schema_only.py index 5c4ab2df..7de30fa7 100644 --- a/test/pytests/test_smart_completion_public_schema_only.py +++ b/test/pytests/test_smart_completion_public_schema_only.py @@ -1,6 +1,7 @@ # type: ignore import os.path +from types import SimpleNamespace from unittest.mock import patch from prompt_toolkit.completion import Completion @@ -105,6 +106,35 @@ def test_special_name_completion(completer, complete_event): ] +def test_dsn_subcommand_completion(completer, complete_event): + text = '/dsn ' + result = completer.get_completions(Document(text=text, cursor_position=len(text)), complete_event) + + assert {completion.text for completion in result} == {'help', 'list', 'show', 'save', 'delete'} + + +def test_dsn_delete_alias_completion(completer, complete_event, monkeypatch): + import mycli.sqlcompleter as sqlcompleter + + monkeypatch.setattr( + sqlcompleter.DsnAliases, + 'instance', + SimpleNamespace(list=lambda: ['prod', 'staging']), + raising=False, + ) + text = '/dsn delete pro' + result = completer.get_completions(Document(text=text, cursor_position=len(text)), complete_event) + + assert list(result) == [Completion(text='prod', start_position=-3)] + + +@pytest.mark.parametrize('text', ['/dsn show', '/dsn show ', '/dsn help', '/dsn help ', '/dsn delete prod ']) +def test_dsn_show_and_help_do_not_offer_completions(completer, complete_event, text): + result = list(completer.get_completions(Document(text=text, cursor_position=len(text)), complete_event)) + + assert result == [] + + def test_empty_string_completion(completer, complete_event): text = "" position = 0 diff --git a/test/pytests/test_sqlcompleter.py b/test/pytests/test_sqlcompleter.py index 1b796eba..68bf7a97 100644 --- a/test/pytests/test_sqlcompleter.py +++ b/test/pytests/test_sqlcompleter.py @@ -522,6 +522,57 @@ def fake_get_completions(tokens: list[str]) -> list[str]: assert 'explain' not in partial_word +def test_get_completions_special_subcommand_branch(monkeypatch) -> None: + monkeypatch.setattr( + mycli.sqlcompleter, + 'suggest_type', + lambda full_text, before: [{'type': 'special_subcommand', 'subcommands': ['help', 'list', 'show', 'save', 'delete']}], + ) + completer = make_completer() + + result = [completion.text for completion in completer.get_completions(Document(text='/dsn s', cursor_position=6), None)] + + assert result == ['show', 'save'] + + +def test_get_completions_dsn_alias_branch(monkeypatch) -> None: + monkeypatch.setattr( + mycli.sqlcompleter, + 'suggest_type', + lambda full_text, before: [{'type': 'dsn_alias'}], + ) + monkeypatch.setattr( + mycli.sqlcompleter.DsnAliases, + 'instance', + SimpleNamespace(list=lambda: ['prod', 'staging']), + raising=False, + ) + completer = make_completer() + + result = [completion.text for completion in completer.get_completions(Document(text='/dsn delete pro', cursor_position=15), None)] + + assert result == ['prod'] + + +def test_get_completions_dsn_alias_branch_without_aliases(monkeypatch) -> None: + monkeypatch.setattr( + mycli.sqlcompleter, + 'suggest_type', + lambda full_text, before: [{'type': 'dsn_alias'}], + ) + monkeypatch.setattr( + mycli.sqlcompleter.DsnAliases, + 'instance', + SimpleNamespace(list=list), + raising=False, + ) + completer = make_completer() + + result = list(completer.get_completions(Document(text='/dsn delete ', cursor_position=12), None)) + + assert result == [] + + def test_find_files_populate_scoped_cols_and_enum_helpers(monkeypatch) -> None: completer = make_completer() completer.extend_schemata('test')