diff --git a/changelog.md b/changelog.md index 35a4ea39..58a59c7d 100644 --- a/changelog.md +++ b/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 diff --git a/mycli/TIPS b/mycli/TIPS index c385b22d..85398499 100644 --- a/mycli/TIPS +++ b/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 ### diff --git a/mycli/client.py b/mycli/client.py index e961f15a..88ca5b79 100644 --- a/mycli/client.py +++ b/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() diff --git a/mycli/client_query.py b/mycli/client_query.py index da4834b0..486cf884 100644 --- a/mycli/client_query.py +++ b/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, }, ) diff --git a/mycli/clistyle.py b/mycli/clistyle.py index c86694e8..ef02b6aa 100644 --- a/mycli/clistyle.py +++ b/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", diff --git a/mycli/completion_refresher.py b/mycli/completion_refresher.py index 914a8709..bdb00472 100644 --- a/mycli/completion_refresher.py +++ b/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()) diff --git a/mycli/myclirc b/mycli/myclirc index a7882745..038dbf4d 100644 --- a/mycli/myclirc +++ b/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' diff --git a/mycli/schema_prefetcher.py b/mycli/schema_prefetcher.py index aedcc69e..e099af4b 100644 --- a/mycli/schema_prefetcher.py +++ b/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, diff --git a/mycli/sqlcompleter.py b/mycli/sqlcompleter.py index bf3cb98b..362ea733 100644 --- a/mycli/sqlcompleter.py +++ b/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]], diff --git a/mycli/sqlexecute.py b/mycli/sqlexecute.py index 5c5030be..9b780557 100644 --- a/mycli/sqlexecute.py +++ b/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 diff --git a/test/myclirc b/test/myclirc index fca7a1a5..05632777 100644 --- a/test/myclirc +++ b/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' diff --git a/test/pytests/test_client.py b/test/pytests/test_client.py index 733b112d..471c34de 100644 --- a/test/pytests/test_client.py +++ b/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( diff --git a/test/pytests/test_client_query.py b/test/pytests/test_client_query.py index 8d8fd9ff..7aebaeff 100644 --- a/test/pytests/test_client_query.py +++ b/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) diff --git a/test/pytests/test_clistyle.py b/test/pytests/test_clistyle.py index 3e152c9f..c216f80f 100644 --- a/test/pytests/test_clistyle.py +++ b/test/pytests/test_clistyle.py @@ -141,6 +141,7 @@ def test_style_factory_helpers_updates_known_tokens(monkeypatch, caplog) -> None cli_style = { 'Token.Prompt': 'Token.Name', 'Token.Toolbar': 'Token.Name', + 'completion-menu.completion.indexed': 'ansired', 'search': 'ansigreen', 'search.current': 'skip-me', 'sql.keyword': 'ansired', @@ -152,6 +153,7 @@ def test_style_factory_helpers_updates_known_tokens(monkeypatch, caplog) -> None output_style = clistyle.style_factory_helpers('native', cli_style) assert output_style.styles[Token.Prompt] == 'ansiblue' + assert output_style.styles[Token.Menu.Completions.Completion.Indexed] == 'ansired' assert output_style.styles[Token.SearchMatch] == 'ansigreen' assert Token.SearchMatch.Current not in output_style.styles assert output_style.styles[Token.Keyword] == 'ansired' @@ -186,6 +188,15 @@ def fake_get_style_by_name(name: str): def test_style_factory_ptoolkit_returns_merged_style_object() -> None: - style = clistyle.style_factory_ptoolkit('native', {'prompt': 'bold'}) + style = clistyle.style_factory_ptoolkit( + 'native', + { + 'prompt': 'bold', + 'completion-menu.completion.indexed': '#ff0000', + }, + ) assert style.get_attrs_for_style_str('class:prompt') == PromptStyle([('prompt', 'bold')]).get_attrs_for_style_str('class:prompt') + assert style.get_attrs_for_style_str('class:completion-menu.completion.indexed') == PromptStyle([ + ('completion-menu.completion.indexed', '#ff0000') + ]).get_attrs_for_style_str('class:completion-menu.completion.indexed') diff --git a/test/pytests/test_completer_use_switch.py b/test/pytests/test_completer_use_switch.py index 17e1d496..dd121dcd 100644 --- a/test/pytests/test_completer_use_switch.py +++ b/test/pytests/test_completer_use_switch.py @@ -26,6 +26,7 @@ def _make_completer() -> SQLCompleter: completer.load_schema_metadata( schema="old_db", table_columns={"orders": ["*", "id", "total"]}, + indexed_columns={"orders": {"id"}}, foreign_keys={}, enum_values={}, functions={}, @@ -63,6 +64,7 @@ def test_columns_available_once_schema_loads() -> None: completer.load_schema_metadata( schema="new_db", table_columns={"customers": ["*", "name", "email"]}, + indexed_columns={"customers": {"email"}}, foreign_keys={}, enum_values={}, functions={}, @@ -105,6 +107,7 @@ def writer() -> None: completer.load_schema_metadata( schema=schema, table_columns={"t": ["*", "c1", "c2"]}, + indexed_columns={"t": {"c1"}}, foreign_keys={}, enum_values={}, functions={}, diff --git a/test/pytests/test_completion_refresher.py b/test/pytests/test_completion_refresher.py index 53b46179..063005fe 100644 --- a/test/pytests/test_completion_refresher.py +++ b/test/pytests/test_completion_refresher.py @@ -57,6 +57,7 @@ def test_ctor(refresher) -> None: "databases", "schemata", "tables", + "indexed_columns", "foreign_keys", "enum_values", "users", @@ -489,6 +490,7 @@ def test_refresh_helpers_delegate_to_completer_and_executor(monkeypatch) -> None executor.dbname = 'current_db' executor.databases.return_value = ['db1', 'db2'] executor.table_columns.return_value = iter([('tbl', 'col')]) + executor.indexed_columns.return_value = iter([('tbl', 'col')]) executor.foreign_keys.return_value = iter([('tbl', 'col', 'other', 'id')]) executor.enum_values.return_value = iter([('tbl', 'status', ['open'])]) executor.users.return_value = iter([('app',)]) @@ -502,6 +504,7 @@ def test_refresh_helpers_delegate_to_completer_and_executor(monkeypatch) -> None completion_refresher.refresh_databases(completer, executor) completion_refresher.refresh_schemata(completer, executor) completion_refresher.refresh_tables(completer, executor) + completion_refresher.refresh_indexed_columns(completer, executor) completion_refresher.refresh_foreign_keys(completer, executor) completion_refresher.refresh_enum_values(completer, executor) completion_refresher.refresh_users(completer, executor) @@ -516,6 +519,7 @@ def test_refresh_helpers_delegate_to_completer_and_executor(monkeypatch) -> None completer.set_dbname.assert_called_once_with('current_db') completer.extend_relations.assert_called_once_with([('tbl', 'col')], kind='tables') completer.extend_columns.assert_called_once_with([('tbl', 'col')], kind='tables') + completer.extend_indexed_columns.assert_called_once_with(executor.indexed_columns.return_value) completer.extend_foreign_keys.assert_called_once_with(executor.foreign_keys.return_value) completer.extend_enum_values.assert_called_once_with(executor.enum_values.return_value) completer.extend_users.assert_called_once_with(executor.users.return_value) diff --git a/test/pytests/test_naive_completion.py b/test/pytests/test_naive_completion.py index fb7556d7..26ec111a 100644 --- a/test/pytests/test_naive_completion.py +++ b/test/pytests/test_naive_completion.py @@ -96,6 +96,19 @@ def test_column_name_completion(completer, complete_event): assert result == list(map(Completion, completer.all_completions)) +def test_indexed_column_completion_is_not_styled(completer, complete_event): + completer.extend_schemata('test') + completer.set_dbname('test') + completer.extend_relations([('users',)], kind='tables') + completer.extend_columns([('users', 'indexed_id')], kind='tables') + completer.extend_indexed_columns([('users', 'indexed_id')]) + text = 'SELECT indexed_ FROM users' + position = len('SELECT indexed_') + result = list(completer.get_completions(Document(text=text, cursor_position=position), complete_event)) + + assert result == [Completion(text='indexed_id', start_position=-8)] + + def test_special_name_completion(completer, complete_event): text = "\\" position = len("\\") diff --git a/test/pytests/test_schema_prefetcher.py b/test/pytests/test_schema_prefetcher.py index c6ee5c91..13ec25ca 100644 --- a/test/pytests/test_schema_prefetcher.py +++ b/test/pytests/test_schema_prefetcher.py @@ -72,6 +72,7 @@ def make(*_args, **_kwargs): executor = MagicMock() executor.databases.return_value = list(databases) if databases is not None else [] executor.table_columns.side_effect = lambda schema=None: iter(per_schema_tables.get(schema, [])) + executor.indexed_columns.side_effect = lambda schema=None: iter([]) executor.foreign_keys.side_effect = lambda schema=None: iter([]) executor.enum_values.side_effect = lambda schema=None: iter([]) executor.functions.side_effect = lambda schema=None: iter([]) @@ -190,6 +191,7 @@ def _track(schema=None): return iter([]) executor.table_columns.side_effect = _track + executor.indexed_columns.side_effect = lambda schema=None: iter([]) executor.foreign_keys.side_effect = lambda schema=None: iter([]) executor.enum_values.side_effect = lambda schema=None: iter([]) executor.functions.side_effect = lambda schema=None: iter([]) @@ -480,6 +482,11 @@ def test_prefetch_one_loads_foreign_keys_enums_functions_and_procedures(monkeypa executor = MagicMock() executor.table_columns.return_value = iter([('orders', 'id')]) + executor.indexed_columns.return_value = iter([ + ('orders', 'id'), + ('orders', 'id'), + ('order details', 'created at'), + ]) executor.foreign_keys.return_value = iter([('orders', 'user_id', 'users', 'id')]) executor.enum_values.return_value = iter([('orders', 'status', ['pending', 'shipped'])]) executor.functions.return_value = iter([(), ('calc_tax',), (None,)]) @@ -490,6 +497,10 @@ def test_prefetch_one_loads_foreign_keys_enums_functions_and_procedures(monkeypa load_schema_metadata.assert_called_once_with( schema='analytics', table_columns={'orders': ['*', 'id']}, + indexed_columns={ + 'orders': {'id'}, + '`order details`': {'`created at`'}, + }, foreign_keys={ 'tables': {'orders': {'users'}, 'users': {'orders'}}, 'relations': [('orders', 'user_id', 'users', 'id')], diff --git a/test/pytests/test_smart_completion_public_schema_only.py b/test/pytests/test_smart_completion_public_schema_only.py index 7de30fa7..d75197a1 100644 --- a/test/pytests/test_smart_completion_public_schema_only.py +++ b/test/pytests/test_smart_completion_public_schema_only.py @@ -827,6 +827,119 @@ def test_string_no_completion_spaces_inner_2(completer, complete_event): assert result == [] +def test_indexed_column_completion_is_styled(completer, complete_event): + completer.extend_indexed_columns([('users', 'email')]) + text = 'SELECT FROM users' + position = len('SELECT ') + result = { + completion.text: completion + for completion in completer.get_completions( + Document(text=text, cursor_position=position), + complete_event, + ) + } + + assert result['email'] == Completion( + text='email', + start_position=0, + display='email*', + style='class:completion-menu.completion.indexed', + ) + assert result['email'].style == 'class:completion-menu.completion.indexed' + assert result['first_name'] == Completion(text='first_name', start_position=0) + assert result['first_name'].style == '' + + +@pytest.mark.parametrize( + ('marker', 'display'), + [ + (' [indexed]', 'email [indexed]'), + ('', 'email'), + ], +) +def test_indexed_column_completion_uses_configured_marker( + completer, + complete_event, + marker: str, + display: str, +) -> None: + completer.indexed_column_suffix = marker + completer.extend_indexed_columns([('users', 'email')]) + text = 'SELECT FROM users' + position = len('SELECT ') + result = { + completion.text: completion + for completion in completer.get_completions( + Document(text=text, cursor_position=position), + complete_event, + ) + } + + assert result['email'] == Completion( + text='email', + start_position=0, + display=display, + ) + assert result['email'].style == 'class:completion-menu.completion.indexed' + + +def test_indexed_column_completion_uses_all_scoped_tables(completer, complete_event): + completer.extend_indexed_columns([('orders', 'id')]) + text = 'SELECT FROM users, orders' + position = len('SELECT ') + result = { + completion.text: completion + for completion in completer.get_completions( + Document(text=text, cursor_position=position), + complete_event, + ) + } + + assert result['id'] == Completion( + text='id', + start_position=0, + display='id*', + style='class:completion-menu.completion.indexed', + ) + assert result['id'].style == 'class:completion-menu.completion.indexed' + + +def test_indexed_column_completion_ignores_out_of_scope_tables(completer, complete_event): + completer.extend_indexed_columns([('orders', 'id')]) + text = 'SELECT FROM users' + position = len('SELECT ') + result = { + completion.text: completion + for completion in completer.get_completions( + Document(text=text, cursor_position=position), + complete_event, + ) + } + + assert result['id'] == Completion(text='id', start_position=0) + assert result['id'].style == '' + + +def test_backticked_indexed_column_completion_is_styled(completer, complete_event): + completer.extend_indexed_columns([('select', 'insert')]) + text = 'SELECT `ins FROM `select`' + position = len('SELECT `ins') + result = list( + completer.get_completions( + Document(text=text, cursor_position=position), + complete_event, + ) + ) + indexed_completion = next(completion for completion in result if completion.text == '`insert`') + + assert indexed_completion == Completion( + text='`insert`', + start_position=-4, + display='`insert`*', + ) + assert indexed_completion.style == 'class:completion-menu.completion.indexed' + + def test_backticked_column_completion(completer, complete_event): text = 'select `Tim' position = len(text) diff --git a/test/pytests/test_sqlcompleter.py b/test/pytests/test_sqlcompleter.py index 68bf7a97..f3ca1be5 100644 --- a/test/pytests/test_sqlcompleter.py +++ b/test/pytests/test_sqlcompleter.py @@ -374,6 +374,12 @@ def test_init_invalid_keyword_casing_defaults_to_auto() -> None: assert completer.keyword_casing == 'auto' +def test_init_configures_indexed_column_suffix() -> None: + completer = SQLCompleter(indexed_column_suffix=' [indexed]') + + assert completer.indexed_column_suffix == ' [indexed]' + + def test_extend_metadata_helpers_and_logging(caplog) -> None: completer = make_completer() completer.set_dbname('missing') @@ -410,6 +416,12 @@ def test_extend_metadata_helpers_and_logging(caplog) -> None: assert "relname 'missing' was not found in db 'test'" in caplog.text assert completer.dbmetadata['tables']['test']['`select`'] == ['*', '`from`'] + completer.extend_indexed_columns([('select', 'from'), ('select', 'from'), ('orders', 'created at')]) + assert completer.dbmetadata['indexed_columns']['test'] == { + '`select`': {'`from`'}, + 'orders': {'`created at`'}, + } + completer.set_dbname('enumdb') completer.extend_enum_values([('order status', 'select', ['pending'])]) assert completer.dbmetadata['enum_values']['enumdb']['`order status`']['`select`'] == ['pending'] @@ -620,11 +632,33 @@ def test_matches_parent(parent: str, schema: str | None, relname: str, alias: st assert SQLCompleter._matches_parent(parent, schema, relname, alias) is expected +def test_populate_scoped_indexed_columns_uses_current_schema() -> None: + completer = SQLCompleter() + completer.set_dbname('test') + completer.dbmetadata['indexed_columns']['test'] = { + 'users': {'id'}, + 'orders': {'created_at'}, + } + + assert completer.populate_scoped_indexed_columns([]) == {'id', 'created_at'} + + +def test_populate_scoped_indexed_columns_uses_explicit_schema_and_escaped_table() -> None: + completer = SQLCompleter() + completer.set_dbname('test') + completer.dbmetadata['indexed_columns']['analytics'] = { + '`order details`': {'`created at`'}, + } + + assert completer.populate_scoped_indexed_columns([('analytics', 'order details', None)]) == {'`created at`'} + + def test_copy_other_schemas_from_preserves_non_current_metadata() -> None: source = SQLCompleter() source.load_schema_metadata( schema='other', table_columns={'users': ['*', 'id', 'email']}, + indexed_columns={'users': {'id'}}, foreign_keys={'tables': {}, 'relations': []}, enum_values={}, functions={'fn_foo': None}, @@ -634,6 +668,7 @@ def test_copy_other_schemas_from_preserves_non_current_metadata() -> None: source.load_schema_metadata( schema='current', table_columns={'stale_current': ['*']}, + indexed_columns={'stale_current': {'id'}}, foreign_keys={'tables': {}, 'relations': []}, enum_values={}, functions={}, @@ -648,6 +683,7 @@ def test_copy_other_schemas_from_preserves_non_current_metadata() -> None: assert 'other' in dest.dbmetadata['tables'] assert dest.dbmetadata['tables']['other'] == {'users': ['*', 'id', 'email']} + assert dest.dbmetadata['indexed_columns']['other'] == {'users': {'id'}} assert dest.dbmetadata['functions']['other'] == {'fn_foo': None} # The excluded schema is not overwritten with stale source data. assert dest.dbmetadata['tables']['current'] == {} @@ -662,6 +698,7 @@ def test_copy_other_schemas_from_does_not_overwrite_existing_dest() -> None: source.load_schema_metadata( schema='shared', table_columns={'from_source': ['*']}, + indexed_columns={'from_source': {'id'}}, foreign_keys={'tables': {}, 'relations': []}, enum_values={}, functions={}, @@ -684,6 +721,7 @@ def test_load_schema_metadata_ignores_empty_schema() -> None: completer.load_schema_metadata( schema='', table_columns={'users': ['*', 'id']}, + indexed_columns={'users': {'id'}}, foreign_keys={'tables': {'users': []}, 'relations': [('users', 'id')]}, enum_values={'users': {'status': ['pending']}}, functions={'fn_users': None}, @@ -696,5 +734,6 @@ def test_load_schema_metadata_ignores_empty_schema() -> None: assert completer.dbmetadata['procedures'] == {} assert completer.dbmetadata['enum_values'] == {} assert completer.dbmetadata['foreign_keys'] == {} + assert completer.dbmetadata['indexed_columns'] == {} assert 'users' not in completer.all_completions assert 'fn_users' not in completer.all_completions diff --git a/test/pytests/test_sqlexecute.py b/test/pytests/test_sqlexecute.py index 8fd41ceb..294cc926 100644 --- a/test/pytests/test_sqlexecute.py +++ b/test/pytests/test_sqlexecute.py @@ -1161,6 +1161,45 @@ def test_table_columns_returns_empty_generator_when_schema_has_no_tables(monkeyp assert cursor.executed == [(SQLExecute.table_columns_query, ('empty_db',))] +@pytest.mark.parametrize( + ('schema', 'expected_schema'), + [ + (None, 'app_db'), + ('analytics', 'analytics'), + ], +) +def test_indexed_columns_executes_query_and_yields_rows(monkeypatch, schema, expected_schema) -> None: + cursor = FakeMetadataCursor([('users', 'id'), ('orders', 'customer_id')]) + executor = make_executor_for_run_tests(FakeMetadataConnection(cursor)) + executor.dbname = 'app_db' + monkeypatch.setattr(sqlexecute, 'Connection', FakeMetadataConnection) + + result = list(executor.indexed_columns(schema=schema)) + + assert result == [('users', 'id'), ('orders', 'customer_id')] + assert cursor.executed == [(SQLExecute.indexed_columns_query, (expected_schema,))] + assert cursor.entered is True + assert cursor.exited is True + assert 'SEQ_IN_INDEX = 1' in SQLExecute.indexed_columns_query + assert 'COLUMN_NAME IS NOT NULL' in SQLExecute.indexed_columns_query + + +def test_indexed_columns_returns_empty_generator_and_logs_execute_errors(monkeypatch, caplog) -> None: + cursor = FakeMetadataCursor([], execute_error=RuntimeError('boom')) + executor = make_executor_for_run_tests(FakeMetadataConnection(cursor)) + executor.dbname = 'app_db' + monkeypatch.setattr(sqlexecute, 'Connection', FakeMetadataConnection) + + with caplog.at_level('ERROR', logger='mycli.sqlexecute'): + result = list(executor.indexed_columns()) + + assert result == [] + assert cursor.executed == [(SQLExecute.indexed_columns_query, ('app_db',))] + assert cursor.entered is True + assert cursor.exited is True + assert "No indexed-column metadata due to RuntimeError('boom')" in caplog.text + + def test_enum_values_executes_query_and_skips_non_enum_columns(monkeypatch) -> None: cursor = FakeMetadataCursor([ ('orders', 'status', "enum('new','paid')"),