Let /dsn offer its subcommands as completions

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.

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"}]
 
 

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'

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"],

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'}]),
     ],
 )

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

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