Highlight indexed columns in completion menus

changelog.md

@@ -6,6 +6,7 @@ Features
 * Show purpose in tabular `--checkup` output.
 * Add optional dependencies to `--checkup`.
 * Improve completions for `/command`s.
+* Highlight indexed columns in completions with a suffix and/or a text style.
 
 
 Documentation

mycli/TIPS

@@ -144,6 +144,8 @@ display output in an interactive explorer \x at the end of a query!
 
 run SQL scripts in batch mode using the standard input!
 
+indexed columns are tagged with a "*" in completion menus!
+
 ###
 ### keystrokes
 ###

mycli/client.py

@@ -183,6 +183,7 @@ def __init__(
         self.initialize_logging()
 
         keyword_casing = c["main"].get("keyword_casing", "auto")
+        indexed_column_suffix = c['main'].get('indexed_column_suffix', '')
 
         self.highlight_preview = c['search'].as_bool('highlight_preview')
 
@@ -191,7 +192,10 @@ def __init__(
         # Initialize completer.
         self.smart_completion = c["main"].as_bool("smart_completion")
         self.completer = SQLCompleter(
-            self.smart_completion, supported_formats=self.main_formatter.supported_formats, keyword_casing=keyword_casing
+            self.smart_completion,
+            supported_formats=self.main_formatter.supported_formats,
+            keyword_casing=keyword_casing,
+            indexed_column_suffix=indexed_column_suffix,
         )
         self._completer_lock = threading.Lock()
 

mycli/client_query.py

@@ -52,6 +52,7 @@ def refresh_completions(self, reset: bool = False) -> list[SQLResult]:
                 "smart_completion": self.smart_completion,
                 "supported_formats": self.main_formatter.supported_formats,
                 "keyword_casing": self.completer.keyword_casing,
+                "indexed_column_suffix": self.completer.indexed_column_suffix,
             },
         )
 

mycli/clistyle.py

@@ -15,6 +15,7 @@
 TOKEN_TO_PROMPT_STYLE: dict[Token, str] = {
     Token.Menu.Completions.Completion.Current: "completion-menu.completion.current",
     Token.Menu.Completions.Completion: "completion-menu.completion",
+    Token.Menu.Completions.Completion.Indexed: "completion-menu.completion.indexed",
     Token.Menu.Completions.Meta.Current: "completion-menu.meta.completion.current",
     Token.Menu.Completions.Meta: "completion-menu.meta.completion",
     Token.Menu.Completions.MultiColumnMeta: "completion-menu.multi-column-meta",

mycli/completion_refresher.py

@@ -167,6 +167,11 @@ def refresh_tables(completer: SQLCompleter, executor: SQLExecute) -> None:
     completer.extend_columns(table_columns_dbresult, kind="tables")
 
 
+@refresher("indexed_columns")
+def refresh_indexed_columns(completer: SQLCompleter, executor: SQLExecute) -> None:
+    completer.extend_indexed_columns(executor.indexed_columns())
+
+
 @refresher("foreign_keys")
 def refresh_foreign_keys(completer: SQLCompleter, executor: SQLExecute) -> None:
     completer.extend_foreign_keys(executor.foreign_keys())

mycli/myclirc

@@ -9,6 +9,12 @@ show_warnings = False
 # possible completions will be listed.
 smart_completion = True
 
+# Text appended to indexed column names in the completion menu.  This text is
+# not inserted into the query. Leave empty to disable the marker.  Quote values
+# containing spaces, commas, or comment characters.
+# Alternative: ' [indexed]'
+indexed_column_suffix = *
+
 # Minimum characters typed before offering completion suggestions.  Forward
 # slash for a command is an exception which always offers completions.
 # Suggestion: 3.
@@ -411,6 +417,7 @@ default_username_field = username
 # Completion menus
 completion-menu.completion.current = 'bg:#ffffff #000000'
 completion-menu.completion = 'bg:#008888 #ffffff'
+completion-menu.completion.indexed = 'bg:#008888 #ffffff'
 completion-menu.meta.completion.current = 'bg:#44aaaa #000000'
 completion-menu.meta.completion = 'bg:#448888 #ffffff'
 completion-menu.multi-column-meta = 'bg:#aaffff #000000'

mycli/schema_prefetcher.py

@@ -174,6 +174,7 @@ def _invalidate_after_visibility_deadline(self) -> None:
     def _prefetch_one(self, executor: SQLExecute, schema: str) -> None:
         _logger.debug('prefetching schema %r', schema)
         table_rows = list(executor.table_columns(schema=schema))
+        indexed_rows = list(executor.indexed_columns(schema=schema))
         fk_rows = list(executor.foreign_keys(schema=schema))
         enum_rows = list(executor.enum_values(schema=schema))
         func_rows = list(executor.functions(schema=schema))
@@ -189,6 +190,12 @@ def _prefetch_one(self, executor: SQLExecute, schema: str) -> None:
             cols = table_columns.setdefault(esc_table, ['*'])
             cols.append(esc_col)
 
+        indexed_columns: dict[str, set[str]] = {}
+        for table, column in indexed_rows:
+            esc_table = completer.escape_name(table)
+            esc_col = completer.escape_name(column)
+            indexed_columns.setdefault(esc_table, set()).add(esc_col)
+
         fk_tables: dict[str, set[str]] = {}
         fk_relations: list[tuple[str, str, str, str]] = []
         for table, col, ref_table, ref_col in fk_rows:
@@ -224,6 +231,7 @@ def _prefetch_one(self, executor: SQLExecute, schema: str) -> None:
             live_completer.load_schema_metadata(
                 schema=schema,
                 table_columns=table_columns,
+                indexed_columns=indexed_columns,
                 foreign_keys=fk_payload,
                 enum_values=enum_values,
                 functions=functions,

mycli/sqlcompleter.py

@@ -21,6 +21,7 @@
 
 _logger = logging.getLogger(__name__)
 _CASE_CHANGE_PAT = re.compile('(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])')
+_INDEXED_COLUMN_STYLE = 'class:completion-menu.completion.indexed'
 
 
 class Fuzziness(IntEnum):
@@ -938,9 +939,11 @@ def __init__(
         smart_completion: bool = True,
         supported_formats: tuple = (),
         keyword_casing: str = "auto",
+        indexed_column_suffix: str = '*',
     ) -> None:
         super(self.__class__, self).__init__()
         self.smart_completion = smart_completion
+        self.indexed_column_suffix = indexed_column_suffix
         self.reserved_words = set()
         for x in self.keywords:
             self.reserved_words.update(x.split())
@@ -1043,6 +1046,15 @@ def extend_columns(self, column_data: list[tuple[str, str]], kind: Literal['tabl
             metadata[self.dbname][relname].append(column)
             self.all_completions.add(column)
 
+    def extend_indexed_columns(self, index_data: Iterable[tuple[str, str]]) -> None:
+        """Extend metadata for columns that lead an index."""
+        metadata = self.dbmetadata["indexed_columns"]
+        schema_meta = metadata.setdefault(self.dbname, {})
+        for table, column in index_data:
+            table = self.escape_name(table)
+            column = self.escape_name(column)
+            schema_meta.setdefault(table, set()).add(column)
+
     def extend_enum_values(self, enum_data: Iterable[tuple[str, str, list[str]]]) -> None:
         metadata = self.dbmetadata["enum_values"]
         if self.dbname not in metadata:
@@ -1162,6 +1174,7 @@ def load_schema_metadata(
         self,
         schema: str,
         table_columns: dict[str, list[str]],
+        indexed_columns: dict[str, set[str]],
         foreign_keys: dict[str, Any],
         enum_values: dict[str, dict[str, list[str]]],
         functions: dict[str, None],
@@ -1177,6 +1190,7 @@ def load_schema_metadata(
         if not schema:
             return
         self.dbmetadata["tables"][schema] = table_columns
+        self.dbmetadata["indexed_columns"][schema] = indexed_columns
         self.dbmetadata["views"].setdefault(schema, {})
         self.dbmetadata["functions"][schema] = functions
         self.dbmetadata["procedures"][schema] = procedures
@@ -1193,7 +1207,7 @@ def copy_other_schemas_from(self, source: "SQLCompleter", exclude: str | None) -
         using qualified completions (``OtherSchema.table``) without a
         re-fetch.
         """
-        kinds = ("tables", "views", "functions", "procedures", "enum_values", "foreign_keys")
+        kinds = ("tables", "views", "functions", "procedures", "enum_values", "foreign_keys", "indexed_columns")
         for kind in kinds:
             src_map = source.dbmetadata.get(kind, {})
             dest_map = self.dbmetadata.setdefault(kind, {})
@@ -1238,6 +1252,7 @@ def reset_completions(self) -> None:
             "procedures": {},
             "enum_values": {},
             "foreign_keys": {},
+            "indexed_columns": {},
         }
         self.all_completions = set(self.keywords + self.functions)
 
@@ -1428,6 +1443,7 @@ def get_completions(
             return (Completion(x[0], -len(text_for_len)) for x in matches)
 
         completions: list[tuple[str, int, int]] = []
+        indexed_column_candidates: set[str] = set()
         suggestions = suggest_type(document.text, document.text_before_cursor)
         rigid_sort = False
         length_based_on_path = False
@@ -1451,10 +1467,16 @@ def get_completions(
                     # showing all columns. So make them unique and sort them.
                     scoped_cols = sorted(set(scoped_cols), key=lambda s: s.strip('`'))
 
-                cols = self.find_matches(
-                    word_before_cursor,
-                    scoped_cols,
-                    text_before_cursor=document.text_before_cursor,
+                cols = list(
+                    self.find_matches(
+                        word_before_cursor,
+                        scoped_cols,
+                        text_before_cursor=document.text_before_cursor,
+                    )
+                )
+                indexed_columns = {self._strip_backticks(column).casefold() for column in self.populate_scoped_indexed_columns(tables)}
+                indexed_column_candidates.update(
+                    candidate for candidate, _fuzziness in cols if self._strip_backticks(candidate).casefold() in indexed_columns
                 )
                 completions.extend([(*x, rank) for x in cols])
 
@@ -1745,9 +1767,25 @@ def completion_sort_key(item: tuple[str, int, int], text_for_len: str):
             uniq_completions_str = dict.fromkeys(x[0] for x in sorted_completions)
 
         if length_based_on_path:
-            return (Completion(x, -len(last_for_len_paths)) for x in uniq_completions_str)
+            return (
+                Completion(
+                    x,
+                    -len(last_for_len_paths),
+                    display=f'{x}{self.indexed_column_suffix}' if x in indexed_column_candidates else None,
+                    style=_INDEXED_COLUMN_STYLE if x in indexed_column_candidates else '',
+                )
+                for x in uniq_completions_str
+            )
         else:
-            return (Completion(x, -len(text_for_len)) for x in uniq_completions_str)
+            return (
+                Completion(
+                    x,
+                    -len(text_for_len),
+                    display=f'{x}{self.indexed_column_suffix}' if x in indexed_column_candidates else None,
+                    style=_INDEXED_COLUMN_STYLE if x in indexed_column_candidates else '',
+                )
+                for x in uniq_completions_str
+            )
 
     def find_files(self, word: str) -> Generator[tuple[str, int], None, None]:
         """Yield matching directory or file names.
@@ -1819,6 +1857,24 @@ def populate_scoped_cols(self, scoped_tbls: list[tuple[str | None, str, str | No
 
         return columns
 
+    def populate_scoped_indexed_columns(self, scoped_tbls: list[tuple[str | None, str, str | None]]) -> set[str]:
+        """Find leading indexed columns in a set of scoped tables."""
+        metadata = self.dbmetadata["indexed_columns"]
+        indexed_columns: set[str] = set()
+
+        if not scoped_tbls:
+            for columns in metadata.get(self.dbname, {}).values():
+                indexed_columns.update(columns)
+            return indexed_columns
+
+        for schema, relname, _alias in scoped_tbls:
+            schema_meta = metadata.get(schema or self.dbname, {})
+            escaped_relname = self.escape_name(relname)
+            indexed_columns.update(schema_meta.get(relname, set()))
+            indexed_columns.update(schema_meta.get(escaped_relname, set()))
+
+        return indexed_columns
+
     def populate_enum_values(
         self,
         scoped_tbls: list[tuple[str | None, str, str | None]],

mycli/sqlexecute.py

@@ -108,6 +108,13 @@ class SQLExecute:
                                     where table_schema = %s
                                     order by table_name,ordinal_position"""
 
+    indexed_columns_query = """SELECT DISTINCT TABLE_NAME, COLUMN_NAME
+                                    FROM information_schema.STATISTICS
+                                    WHERE TABLE_SCHEMA = %s
+                                      AND SEQ_IN_INDEX = 1
+                                      AND COLUMN_NAME IS NOT NULL
+                                    ORDER BY TABLE_NAME, COLUMN_NAME"""
+
     enum_values_query = """select TABLE_NAME, COLUMN_NAME, COLUMN_TYPE from information_schema.columns
                                     where table_schema = %s and data_type = 'enum'
                                     order by table_name,ordinal_position"""
@@ -439,6 +446,18 @@ def table_columns(self, schema: str | None = None) -> Generator[tuple[str, str],
             cur.execute(self.table_columns_query, (target,))
             yield from cur
 
+    def indexed_columns(self, schema: str | None = None) -> Generator[tuple[str, str], None, None]:
+        """Yields leading indexed (table name, column name) pairs for *schema*."""
+        target = schema if schema is not None else self.dbname
+        assert isinstance(self.conn, Connection)
+        with self.conn.cursor() as cur:
+            _logger.debug("Indexed Columns Query. sql: %r schema: %r", self.indexed_columns_query, target)
+            try:
+                cur.execute(self.indexed_columns_query, (target,))
+                yield from cur
+            except Exception as e:
+                _logger.error('No indexed-column metadata due to %r', e)
+
     def enum_values(self, schema: str | None = None) -> Generator[tuple[str, str, list[str]], None, None]:
         """Yields (table name, column name, enum values) tuples for *schema*."""
         target = schema if schema is not None else self.dbname

test/myclirc

@@ -9,6 +9,12 @@ show_warnings = False
 # possible completions will be listed.
 smart_completion = True
 
+# Text appended to indexed column names in the completion menu.  This text is
+# not inserted into the query. Leave empty to disable the marker.  Quote values
+# containing spaces, commas, or comment characters.
+# Alternative: ' [indexed]'
+indexed_column_suffix = *
+
 # Minimum characters typed before offering completion suggestions.  Forward
 # slash for a command is an exception which always offers completions.
 # Suggestion: 3.
@@ -411,6 +417,7 @@ default_username_field = username
 # Completion menus
 completion-menu.completion.current = 'bg:#ffffff #000000'
 completion-menu.completion = 'bg:#008888 #ffffff'
+completion-menu.completion.indexed = 'bg:#008888 #ffffff'
 completion-menu.meta.completion.current = 'bg:#44aaaa #000000'
 completion-menu.meta.completion = 'bg:#448888 #ffffff'
 completion-menu.multi-column-meta = 'bg:#aaffff #000000'

test/pytests/test_client.py

@@ -79,6 +79,36 @@ def test_init_uses_default_plot_theme_for_empty_value(monkeypatch: pytest.Monkey
     assert cli.plot_theme == 'carbong90'
 
 
+def test_init_configures_indexed_column_suffix(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    patch_constructor_side_effects(monkeypatch)
+    myclirc = write_myclirc(
+        tmp_path,
+        """
+        [main]
+        indexed_column_suffix = " [indexed]"
+        """,
+    )
+
+    cli = MyCli(myclirc=myclirc)
+
+    assert cli.completer.indexed_column_suffix == ' [indexed]'
+
+
+def test_init_allows_empty_indexed_column_suffix(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    patch_constructor_side_effects(monkeypatch)
+    myclirc = write_myclirc(
+        tmp_path,
+        """
+        [main]
+        indexed_column_suffix =
+        """,
+    )
+
+    cli = MyCli(myclirc=myclirc)
+
+    assert cli.completer.indexed_column_suffix == ''
+
+
 def test_init_configures_kitty_image_protocol(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
     patch_constructor_side_effects(monkeypatch)
     myclirc = write_myclirc(

test/pytests/test_client_query.py

@@ -24,6 +24,7 @@ def make_refresh_cli() -> tuple[Any, dict[str, Any]]:
     cli._on_completions_refreshed = callback
     cli.completer = SimpleNamespace(
         keyword_casing='upper',
+        indexed_column_suffix=' [indexed]',
         set_dbname=lambda dbname: state['set_dbname_calls'].append(dbname),
     )
     cli.main_formatter = SimpleNamespace(supported_formats=['ascii', 'csv'])
@@ -64,6 +65,7 @@ def test_refresh_completions_passes_options_to_refresher() -> None:
                 'smart_completion': True,
                 'supported_formats': ['ascii', 'csv'],
                 'keyword_casing': 'upper',
+                'indexed_column_suffix': ' [indexed]',
             },
         )
     ]
@@ -82,7 +84,11 @@ def test_refresh_completions_updates_dbname_when_reset() -> None:
     set_dbname_calls: list[str] = []
     cli.schema_prefetcher = SimpleNamespace(stop=lambda: None)
     cli.sqlexecute = SimpleNamespace(dbname='next_db')
-    cli.completer = SimpleNamespace(keyword_casing='lower', set_dbname=lambda dbname: set_dbname_calls.append(dbname))
+    cli.completer = SimpleNamespace(
+        keyword_casing='lower',
+        indexed_column_suffix='*',
+        set_dbname=lambda dbname: set_dbname_calls.append(dbname),
+    )
     cli.main_formatter = SimpleNamespace(supported_formats=['table'])
     cli.completion_refresher = SimpleNamespace(refresh=lambda executor, callbacks, options: None)
 
@@ -97,7 +103,11 @@ def test_refresh_completions_uses_lock_when_reset() -> None:
     cli.schema_prefetcher = SimpleNamespace(stop=lambda: None)
     cli.sqlexecute = SimpleNamespace(dbname='next_db')
     cli._completer_lock = cast(Any, ReusableLock(lambda: entered_lock.__setitem__('count', entered_lock['count'] + 1)))
-    cli.completer = SimpleNamespace(keyword_casing='lower', set_dbname=lambda dbname: None)
+    cli.completer = SimpleNamespace(
+        keyword_casing='lower',
+        indexed_column_suffix='*',
+        set_dbname=lambda dbname: None,
+    )
     cli.main_formatter = SimpleNamespace(supported_formats=['table'])
     cli.completion_refresher = SimpleNamespace(refresh=lambda executor, callbacks, options: None)