From ab8c37473f9b3edc8201bc525b58190aae41a005 Mon Sep 17 00:00:00 2001 From: Peter Adams <18162810+Maxteabag@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:54:40 +0200 Subject: [PATCH] Add SQL Server routine autocomplete --- .../connections/providers/adapters/base.py | 38 ++ .../connections/providers/mssql/adapter.py | 37 ++ sqlit/domains/query/completion/completion.py | 224 ++++++++- sqlit/domains/query/completion/core.py | 1 + .../query/ui/mixins/autocomplete_schema.py | 54 ++- tests/unit/test_mssql_routine_completion.py | 459 ++++++++++++++++++ 6 files changed, 796 insertions(+), 17 deletions(-) create mode 100644 tests/unit/test_mssql_routine_completion.py diff --git a/sqlit/domains/connections/providers/adapters/base.py b/sqlit/domains/connections/providers/adapters/base.py index 573dc475..d7258cdc 100644 --- a/sqlit/domains/connections/providers/adapters/base.py +++ b/sqlit/domains/connections/providers/adapters/base.py @@ -72,6 +72,44 @@ class ColumnInfo: is_primary_key: bool = False +class RoutineInfo(str): + """String-compatible routine metadata used by autocomplete.""" + + schema: str + database: str + name: str + routine_type: str + return_type: str + parameters: tuple[str, ...] + + def __new__( + cls, + name: str, + *, + schema: str = "", + database: str = "", + routine_type: str = "PROCEDURE", + return_type: str = "", + parameters: tuple[str, ...] = (), + ) -> RoutineInfo: + instance = super().__new__(cls, name) + instance.name = name + instance.schema = schema + instance.database = database + instance.routine_type = routine_type.upper() + instance.return_type = return_type.upper() + instance.parameters = parameters + return instance + + @property + def qualified_name(self) -> str: + return f"{self.schema}.{self}" if self.schema else str(self) + + @property + def is_table_valued(self) -> bool: + return self.routine_type == "FUNCTION" and self.return_type == "TABLE" + + @dataclass class IndexInfo: """Information about a database index.""" diff --git a/sqlit/domains/connections/providers/mssql/adapter.py b/sqlit/domains/connections/providers/mssql/adapter.py index 4880a5d3..dc67ba09 100644 --- a/sqlit/domains/connections/providers/mssql/adapter.py +++ b/sqlit/domains/connections/providers/mssql/adapter.py @@ -421,6 +421,43 @@ def get_procedures(self, conn: Any, database: str | None = None) -> list[str]: ) return [row[0] for row in cursor.fetchall()] + def get_completion_routines( + self, conn: Any, database: str | None = None + ) -> list[Any]: + """Get procedures/functions with parameter metadata for autocomplete.""" + from sqlit.domains.connections.providers.adapters.base import RoutineInfo + + cursor = self._get_cursor_for_database(conn, database) + cursor.execute( + "SELECT r.ROUTINE_SCHEMA, r.ROUTINE_NAME, r.ROUTINE_TYPE, r.DATA_TYPE, " + "p.PARAMETER_NAME, p.ORDINAL_POSITION " + "FROM INFORMATION_SCHEMA.ROUTINES r " + "LEFT JOIN INFORMATION_SCHEMA.PARAMETERS p " + "ON p.SPECIFIC_SCHEMA = r.SPECIFIC_SCHEMA " + "AND p.SPECIFIC_NAME = r.SPECIFIC_NAME " + "ORDER BY r.ROUTINE_SCHEMA, r.ROUTINE_NAME, p.ORDINAL_POSITION" + ) + + grouped: dict[tuple[str, str, str, str], list[str]] = {} + for schema, name, routine_type, return_type, parameter, ordinal in cursor.fetchall(): + key = (schema, name, routine_type, return_type or "") + if parameter and (ordinal is None or ordinal > 0): + grouped.setdefault(key, []).append(parameter) + else: + grouped.setdefault(key, []) + + return [ + RoutineInfo( + name, + schema=schema, + database=database or "", + routine_type=routine_type, + return_type=return_type, + parameters=tuple(parameters), + ) + for (schema, name, routine_type, return_type), parameters in grouped.items() + ] + def get_indexes(self, conn: Any, database: str | None = None) -> list[IndexInfo]: """Get indexes from SQL Server.""" cursor = self._get_cursor_for_database(conn, database) diff --git a/sqlit/domains/query/completion/completion.py b/sqlit/domains/query/completion/completion.py index c60cf00b..ee94d23a 100644 --- a/sqlit/domains/query/completion/completion.py +++ b/sqlit/domains/query/completion/completion.py @@ -6,6 +6,9 @@ from __future__ import annotations import re +from collections.abc import Sequence + +from sqlit.domains.query.app.multi_statement import split_statements from .alter_table import get_alter_table_completions from .core import ( @@ -27,6 +30,7 @@ is_inside_string, remove_comments, remove_string_literals, + split_identifier_parts, ) from .create_index import get_create_index_completions from .create_table import get_create_table_completions @@ -63,11 +67,60 @@ def get_context(sql: str, cursor_pos: int) -> list[Suggestion]: if before_cursor.rstrip().endswith(";"): return [] + statements = split_statements(before_cursor) + current_statement = statements[-1] if statements else before_cursor + routine_context = remove_comments(remove_string_literals(current_statement)) + routine_context = re.split( + r"\n(?=\s*(?:SELECT|INSERT|UPDATE|DELETE|MERGE|WITH|EXEC(?:UTE)?|" + r"CREATE|ALTER|DROP)\b)", + routine_context, + flags=re.IGNORECASE, + )[-1] + routine_invocation = re.search( + r"\b(?:EXEC|EXECUTE)\s+" + r"(?:@\w+\s*=\s*)?" + r"([^\s,=;]+)" + r"(?P.*)$", + routine_context, + re.IGNORECASE, + ) + if routine_invocation and routine_invocation.group("suffix").startswith((" ", "\t", "\n")): + return [ + Suggestion( + type=SuggestionType.PARAMETER, + table_scope=routine_invocation.group(1), + ) + ] + if routine_invocation and not routine_invocation.group("suffix"): + target = routine_invocation.group(1) + parts = split_identifier_parts(target) + if target.endswith("."): + qualifiers = parts + else: + qualifiers = parts[:-1] + scope = ".".join(qualifiers) if qualifiers else None + return [Suggestion(type=SuggestionType.PROCEDURE, table_scope=scope)] + + qualified_from = re.search( + r"\b(?:FROM|JOIN|APPLY)\s+" + r"(?:(?:\[[^\]]+\]|\w+)\.){1,2}\[?\w*$", + before_cursor, + re.IGNORECASE, + ) + if qualified_from: + return [Suggestion(type=SuggestionType.TABLE)] + # Check for table.column pattern (alias or table prefix) - dot_match = re.search(r"(\w+)\.\w*$", before_cursor) + dot_match = re.search( + r"((?:\[[^\]]+\]|\w+)(?:\.(?:\[[^\]]+\]|\w+))*)\.\[?\w*$", + before_cursor, + ) if dot_match: prefix = dot_match.group(1) - return [Suggestion(type=SuggestionType.ALIAS_COLUMN, table_scope=prefix)] + return [ + Suggestion(type=SuggestionType.ALIAS_COLUMN, table_scope=prefix), + Suggestion(type=SuggestionType.FUNCTION, table_scope=prefix), + ] # Try statement-specific handlers for handler in [get_insert_context, get_update_context, get_delete_context]: @@ -162,6 +215,77 @@ def get_completions( """ before_cursor = sql[:cursor_pos] current_word = get_current_word(sql, cursor_pos) + routines = procedures or [] + + def routine_display_name(routine: object) -> str: + name = str(routine) + identities = { + ( + str(getattr(candidate, "database", "")).lower(), + str(getattr(candidate, "schema", "")).lower(), + ) + for candidate in routines + if str(candidate).lower() == name.lower() + } + if len(identities) <= 1: + return name + parts = [ + str(getattr(routine, "database", "")), + str(getattr(routine, "schema", "")), + name, + ] + return ".".join(part for part in parts if part) + + def scoped_routine_names( + candidates: Sequence[object], + *, + schema: str, + database: str = "", + ) -> list[str]: + filtered = [ + routine + for routine in candidates + if str(getattr(routine, "schema", "")).lower() == schema.lower() + and ( + not database + or str(getattr(routine, "database", "")).lower() + == database.lower() + ) + ] + counts: dict[str, int] = {} + for routine in filtered: + key = str(routine).lower() + counts[key] = counts.get(key, 0) + 1 + return [ + str(routine) + for routine in filtered + if counts[str(routine).lower()] == 1 + ] + + procedure_names = [ + routine_display_name(routine) + for routine in routines + if getattr(routine, "routine_type", "PROCEDURE") == "PROCEDURE" + ] + function_names = [ + routine_display_name(routine) + for routine in routines + if getattr(routine, "routine_type", "") == "FUNCTION" + ] + legacy_routine_names = [ + str(routine) for routine in routines if not hasattr(routine, "routine_type") + ] + function_routines = [ + routine + for routine in routines + if getattr(routine, "routine_type", "") == "FUNCTION" + ] + table_function_routines = [ + routine + for routine in function_routines + if bool(getattr(routine, "is_table_valued", False)) + ] + table_function_names = [routine_display_name(routine) for routine in table_function_routines] # Don't suggest if inside string literal if is_inside_string(before_cursor): @@ -183,7 +307,12 @@ def get_completions( return fuzzy_match(current_word, result) # Try DROP handler (doesn't need columns) - drop_result = get_drop_completions(before_cursor, tables, procedures) + drop_routines = ( + (function_names if include_functions else []) + legacy_routine_names + if re.search(r"\bDROP\s+FUNCTION\b", before_cursor, re.IGNORECASE) + else procedure_names + ) + drop_result = get_drop_completions(before_cursor, tables, drop_routines) if drop_result is not None: return fuzzy_match(current_word, drop_result) @@ -258,11 +387,27 @@ def get_completions( # Schema.table prefix → suggest tables after schema name # Pattern: FROM/JOIN schema. or schema.partial - namespace_match = re.search(r"\b(?:FROM|JOIN)\s+(\w+)\.\w*$", clean_before, re.IGNORECASE) + namespace_match = re.search( + r"\b(?:FROM|JOIN|APPLY)\s+" + r"(?:(?:\[([^\]]+)\]|(\w+))\.)?" + r"(?:\[([^\]]+)\]|(\w+))\.\[?\w*$", + clean_before, + re.IGNORECASE, + ) if namespace_match: - namespace = namespace_match.group(1) + database = namespace_match.group(1) or namespace_match.group(2) or "" + namespace = namespace_match.group(3) or namespace_match.group(4) names = get_names_for_namespace(namespace, tables) - return fuzzy_match(current_word, names if names else tables) + candidates = list(names if names else tables) + if include_functions: + candidates.extend( + scoped_routine_names( + table_function_routines, + schema=namespace, + database=database, + ) + ) + return fuzzy_match(current_word, candidates) # ANY/ALL/SOME ( → suggest SELECT for subquery if re.search(r"\b(ANY|ALL|SOME)\s*\(\s*\w*$", clean_before, re.IGNORECASE): @@ -333,6 +478,8 @@ def get_completions( if suggestion.type == SuggestionType.TABLE: results.extend(get_identifier_namespaces(tables)) results.extend(tables) + if include_functions: + results.extend(table_function_names) results.extend(cte_names) elif suggestion.type == SuggestionType.COLUMN: @@ -358,6 +505,7 @@ def get_completions( if include_functions: results.extend(get_all_functions()) + results.extend(function_names) elif suggestion.type == SuggestionType.ALIAS_COLUMN: scope = suggestion.table_scope @@ -372,14 +520,74 @@ def get_completions( results.extend(columns[scope_lower]) elif suggestion.type == SuggestionType.PROCEDURE: - if procedures: - results.extend(procedures) + scope_parts = split_identifier_parts(suggestion.table_scope or "") + schema_scope = scope_parts[-1].lower() if scope_parts else "" + database_scope = scope_parts[-2].lower() if len(scope_parts) > 1 else "" + if not schema_scope: + results.extend(procedure_names) + else: + results.extend( + scoped_routine_names( + [ + routine + for routine in routines + if getattr(routine, "routine_type", "PROCEDURE") + == "PROCEDURE" + ], + schema=schema_scope, + database=database_scope, + ) + ) + + elif suggestion.type == SuggestionType.FUNCTION: + scope_parts = split_identifier_parts(suggestion.table_scope or "") + schema_scope = scope_parts[-1].lower() if scope_parts else "" + database_scope = scope_parts[-2].lower() if len(scope_parts) > 1 else "" + is_alias = len(scope_parts) == 1 and schema_scope in alias_map + if include_functions and not is_alias: + if schema_scope: + results.extend( + scoped_routine_names( + function_routines, + schema=schema_scope, + database=database_scope, + ) + ) + else: + results.extend( + routine_display_name(routine) + for routine in function_routines + ) + + elif suggestion.type == SuggestionType.PARAMETER: + parts = split_identifier_parts(suggestion.table_scope or "") + name_part = parts[-1].lower() if parts else "" + schema_part = parts[-2].lower() if len(parts) > 1 else "" + database_part = parts[-3].lower() if len(parts) > 2 else "" + matches = [] + for routine in routines: + name = str(routine).lower() + schema = str(getattr(routine, "schema", "")).lower() + database = str(getattr(routine, "database", "")).lower() + if ( + name == name_part + and (not schema_part or schema == schema_part) + and (not database_part or database == database_part) + ): + matches.append(routine) + if not database_part and len(matches) > 1: + # The user's default schema is not available here, so choosing + # any same-named routine would risk suggesting invalid params. + matches = [] + for routine in matches: + results.extend(getattr(routine, "parameters", ())) elif suggestion.type == SuggestionType.KEYWORD: if include_keywords: results.extend(get_all_keywords()) if include_functions: results.extend(get_all_functions()) + results.extend(function_names) elif suggestion.type == SuggestionType.OPERATOR: results.extend(SQL_OPERATORS) diff --git a/sqlit/domains/query/completion/core.py b/sqlit/domains/query/completion/core.py index 8b18a563..16dadc21 100644 --- a/sqlit/domains/query/completion/core.py +++ b/sqlit/domains/query/completion/core.py @@ -23,6 +23,7 @@ class SuggestionType(Enum): SCHEMA = auto() DATABASE = auto() PROCEDURE = auto() + PARAMETER = auto() ALIAS_COLUMN = auto() # Column for a specific table/alias OPERATOR = auto() # Comparison operators (=, <, >, etc.) diff --git a/sqlit/domains/query/ui/mixins/autocomplete_schema.py b/sqlit/domains/query/ui/mixins/autocomplete_schema.py index 1da9f2b1..3c481769 100644 --- a/sqlit/domains/query/ui/mixins/autocomplete_schema.py +++ b/sqlit/domains/query/ui/mixins/autocomplete_schema.py @@ -12,6 +12,28 @@ SCHEMA_PROCESS_BATCH_SIZE = 200 +def _dedupe_routines(routines: list[Any]) -> list[Any]: + """Deduplicate routines without collapsing equal names in other schemas.""" + unique: dict[tuple[str, str, str, str], Any] = {} + for routine in routines: + key = ( + str(getattr(routine, "database", "")).lower(), + str(getattr(routine, "schema", "")).lower(), + str(routine).lower(), + str(getattr(routine, "routine_type", "PROCEDURE")).upper(), + ) + unique.setdefault(key, routine) + return list(unique.values()) + + +def _completion_routine_loader(inspector: Any) -> tuple[str, Any]: + """Return a cache key and loader that cannot collide with explorer data.""" + rich_loader = getattr(inspector, "get_completion_routines", None) + if callable(rich_loader): + return "completion_routines", rich_loader + return "procedures", inspector.get_procedures + + class AutocompleteSchemaMixin: """Mixin providing schema loading and caching for autocomplete.""" @@ -522,12 +544,13 @@ def _load_procedures_job(self: AutocompleteMixinHost, database: str | None) -> N self._schema_job_complete() return procedure_inspector = cast(ProcedureInspector, inspector) + cache_field, loader = _completion_routine_loader(procedure_inspector) cache_key = database or "__default__" # Check shared cache first (may have been populated by tree expansion) - if cache_key in self._db_object_cache and "procedures" in self._db_object_cache[cache_key]: - self._process_procedures_result(self._db_object_cache[cache_key]["procedures"], cache_key) + if cache_key in self._db_object_cache and cache_field in self._db_object_cache[cache_key]: + self._process_procedures_result(self._db_object_cache[cache_key][cache_field], cache_key) return # Offload DB call to thread @@ -536,18 +559,30 @@ def work() -> None: db_arg = database if hasattr(self, "_get_metadata_db_arg"): db_arg = self._get_metadata_db_arg(database) - procedures = self._run_db_call(procedure_inspector.get_procedures, connection, db_arg) - self.call_from_thread(self._on_procedures_loaded, procedures, database, cache_key) + procedures = self._run_db_call(loader, connection, db_arg) + self.call_from_thread( + self._on_procedures_loaded, + procedures, + database, + cache_key, + cache_field, + ) except Exception as e: self.call_from_thread(self._on_procedures_error, e, database) self.run_worker(work, thread=True, name=f"load-procedures-{cache_key}") - def _on_procedures_loaded(self: AutocompleteMixinHost, procedures: list, database: str | None, cache_key: str) -> None: + def _on_procedures_loaded( + self: AutocompleteMixinHost, + procedures: list, + database: str | None, + cache_key: str, + cache_field: str = "procedures", + ) -> None: """Handle procedures loaded from thread.""" if cache_key not in self._db_object_cache: self._db_object_cache[cache_key] = {} - self._db_object_cache[cache_key]["procedures"] = procedures + self._db_object_cache[cache_key][cache_field] = procedures self._process_procedures_result(procedures, cache_key) def _on_procedures_error(self: AutocompleteMixinHost, error: Exception, database: str | None) -> None: @@ -681,7 +716,7 @@ def work() -> None: try: tables_dedup = list(dict.fromkeys(tables)) views_dedup = list(dict.fromkeys(views)) - procedures_dedup = list(dict.fromkeys(procedures)) + procedures_dedup = _dedupe_routines(procedures) except Exception as e: self.call_from_thread(self._on_schema_dedup_error, e) return @@ -840,7 +875,8 @@ async def run_db_call(fn: Any, *args: Any, **kwargs: Any) -> Any: # Get procedures if caps.supports_stored_procedures and isinstance(inspector, ProcedureInspector): - procedures = await run_db_call(inspector.get_procedures, connection, db_arg) + _cache_field, loader = _completion_routine_loader(inspector) + procedures = await run_db_call(loader, connection, db_arg) schema_cache["procedures"].extend(procedures) except Exception: @@ -849,7 +885,7 @@ async def run_db_call(fn: Any, *args: Any, **kwargs: Any) -> Any: # Deduplicate schema_cache["tables"] = list(dict.fromkeys(schema_cache["tables"])) schema_cache["views"] = list(dict.fromkeys(schema_cache["views"])) - schema_cache["procedures"] = list(dict.fromkeys(schema_cache["procedures"])) + schema_cache["procedures"] = _dedupe_routines(schema_cache["procedures"]) # Update cache - columns will be lazy-loaded when needed self._update_schema_cache(schema_cache, table_metadata) diff --git a/tests/unit/test_mssql_routine_completion.py b/tests/unit/test_mssql_routine_completion.py new file mode 100644 index 00000000..3e002f94 --- /dev/null +++ b/tests/unit/test_mssql_routine_completion.py @@ -0,0 +1,459 @@ +"""SQL Server procedure/function autocomplete regressions.""" + +from unittest.mock import MagicMock + +from sqlit.domains.connections.providers.adapters.base import RoutineInfo +from sqlit.domains.connections.providers.mssql.adapter import SQLServerAdapter +from sqlit.domains.query.completion import get_completions +from sqlit.domains.query.ui.mixins.autocomplete_schema import ( + _completion_routine_loader, + _dedupe_routines, +) + + +def _routines() -> list[RoutineInfo]: + return [ + RoutineInfo( + "dt_setpropertybyid", + schema="dbo", + routine_type="PROCEDURE", + parameters=("@direction", "@value"), + ), + RoutineInfo( + "fn_score", + schema="dbo", + routine_type="FUNCTION", + return_type="int", + parameters=("@user_id",), + ), + RoutineInfo( + "fn_orders", + schema="dbo", + routine_type="FUNCTION", + return_type="TABLE", + parameters=("@user_id",), + ), + ] + + +def test_mssql_loads_routine_types_and_parameters() -> None: + connection = MagicMock() + cursor = connection.cursor.return_value + cursor.fetchall.return_value = [ + ("dbo", "dt_setpropertybyid", "PROCEDURE", None, "@direction", 1), + ("dbo", "dt_setpropertybyid", "PROCEDURE", None, "@value", 2), + ("dbo", "fn_orders", "FUNCTION", "TABLE", None, 0), + ("dbo", "fn_orders", "FUNCTION", "TABLE", "@user_id", 1), + ("dbo", "fn_score", "FUNCTION", "int", None, 0), + ("dbo", "fn_score", "FUNCTION", "int", "@user_id", 1), + ] + + routines = SQLServerAdapter().get_completion_routines(connection, database="AppDB") + + by_name = {str(routine): routine for routine in routines} + assert by_name["dt_setpropertybyid"].parameters == ("@direction", "@value") + assert by_name["fn_orders"].is_table_valued is True + assert by_name["fn_score"].return_type == "INT" + assert cursor.execute.call_args_list[0].args[0] == "USE [AppDB]" + assert "INFORMATION_SCHEMA.PARAMETERS" in cursor.execute.call_args_list[1].args[0] + + +def test_exec_parameter_completion() -> None: + sql = "EXEC dt_setpropertybyid dir" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["@direction"] + + +def test_schema_qualified_exec_parameter_completion() -> None: + sql = "EXEC dbo.dt_setpropertybyid val" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["@value"] + + +def test_schema_qualified_procedure_name_completion() -> None: + sql = "EXEC dbo.dt_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["dt_setpropertybyid"] + + +def test_schema_qualified_procedure_completion_after_dot() -> None: + sql = "EXEC dbo." + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["dt_setpropertybyid"] + + +def test_partially_bracketed_procedure_name_completion() -> None: + sql = "EXEC [dbo].[dt_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["dt_setpropertybyid"] + + +def test_bracket_qualified_exec_parameter_completion() -> None: + sql = "EXEC [dbo].[dt_setpropertybyid] dir" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["@direction"] + + +def test_later_exec_parameter_completion() -> None: + sql = "EXEC dt_setpropertybyid @direction = 'up', val" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["@value"] + + +def test_parameter_completion_uses_current_exec_statement() -> None: + routines = _routines() + [ + RoutineInfo( + "first_proc", + schema="dbo", + routine_type="PROCEDURE", + parameters=("@first_only",), + ) + ] + sql = "EXEC first_proc @first_only = 1; EXEC dt_setpropertybyid val" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == ["@value"] + + +def test_parameter_completion_after_return_status_assignment() -> None: + sql = "EXEC @status = dbo.dt_setpropertybyid dir" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["@direction"] + + +def test_scalar_function_completion_in_select_expression() -> None: + sql = "SELECT fn_sc" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_score" in completions + assert "dt_setpropertybyid" not in completions + + +def test_table_valued_function_completion_after_from() -> None: + sql = "SELECT * FROM fn_ord" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_orders" in completions + assert "fn_score" not in completions + + +def test_schema_qualified_function_completion() -> None: + sql = "SELECT dbo.fn_sc" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_score" in completions + + +def test_schema_qualified_table_function_excludes_scalar_functions() -> None: + sql = "SELECT * FROM dbo.fn_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_orders" in completions + assert "fn_score" not in completions + + +def test_schema_qualified_table_function_excludes_other_schemas() -> None: + routines = _routines() + [ + RoutineInfo( + "fn_sales", + schema="sales", + routine_type="FUNCTION", + return_type="TABLE", + ) + ] + sql = "SELECT * FROM dbo.fn_" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert "fn_orders" in completions + assert "fn_sales" not in completions + + +def test_table_alias_dot_does_not_offer_scalar_functions() -> None: + sql = "SELECT * FROM orders o WHERE o." + + completions = get_completions( + sql, + len(sql), + ["orders"], + {"orders": ["id", "total"]}, + _routines(), + ) + + assert "id" in completions + assert "fn_score" not in completions + + +def test_schema_qualified_apply_excludes_other_schemas() -> None: + routines = _routines() + [ + RoutineInfo( + "fn_sales", + schema="sales", + routine_type="FUNCTION", + return_type="TABLE", + ) + ] + sql = "SELECT * FROM users CROSS APPLY dbo.fn_" + + completions = get_completions(sql, len(sql), ["users"], {}, routines) + + assert "fn_orders" in completions + assert "fn_sales" not in completions + + +def test_bracket_qualified_table_function_excludes_other_schemas() -> None: + routines = _routines() + [ + RoutineInfo( + "fn_sales", + schema="sales", + routine_type="FUNCTION", + return_type="TABLE", + ) + ] + sql = "SELECT * FROM [dbo].fn_" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert "fn_orders" in completions + assert "fn_sales" not in completions + + +def test_partially_bracketed_table_function_completion() -> None: + sql = "SELECT * FROM [dbo].[fn_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_orders" in completions + assert "fn_score" not in completions + + +def test_database_qualified_table_function_completion() -> None: + routines = [ + RoutineInfo( + "fn_orders", + database="AppDB", + schema="dbo", + routine_type="FUNCTION", + return_type="TABLE", + ), + RoutineInfo( + "fn_orders", + database="AuditDB", + schema="dbo", + routine_type="FUNCTION", + return_type="TABLE", + ), + ] + sql = "SELECT * FROM AppDB.dbo.fn_" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == ["fn_orders"] + + +def test_database_qualified_scalar_function_completion() -> None: + routines = [ + RoutineInfo( + "fn_score", + database="AppDB", + schema="dbo", + routine_type="FUNCTION", + return_type="INT", + ), + RoutineInfo( + "fn_score", + database="AuditDB", + schema="dbo", + routine_type="FUNCTION", + return_type="INT", + ), + ] + sql = "SELECT AppDB.dbo.fn_" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == ["fn_score"] + + +def test_exec_inside_string_does_not_trigger_parameter_completion() -> None: + sql = "SELECT 'EXEC dbo.dt_setpropertybyid dir', fn_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_score" in completions + assert "@direction" not in completions + + +def test_exec_inside_comment_does_not_trigger_parameter_completion() -> None: + sql = "-- EXEC dbo.dt_setpropertybyid dir\nSELECT fn_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_score" in completions + assert "@direction" not in completions + + +def test_semicolon_free_statement_after_exec_uses_new_statement_context() -> None: + sql = "EXEC dbo.dt_setpropertybyid @direction = 'up'\nSELECT fn_sc" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert "fn_score" in completions + assert "@direction" not in completions + + +def test_routine_deduplication_preserves_same_name_in_different_schemas() -> None: + dbo = RoutineInfo( + "refresh", + schema="dbo", + parameters=("@dbo_value",), + ) + sales = RoutineInfo( + "refresh", + schema="sales", + parameters=("@sales_value",), + ) + + assert _dedupe_routines([dbo, sales, dbo]) == [dbo, sales] + + +def test_routine_deduplication_preserves_same_name_in_different_databases() -> None: + app_db = RoutineInfo("refresh", database="AppDB", schema="dbo") + audit_db = RoutineInfo("refresh", database="AuditDB", schema="dbo") + + assert _dedupe_routines([app_db, audit_db, app_db]) == [app_db, audit_db] + + +def test_duplicate_cross_database_routines_are_qualified() -> None: + routines = [ + RoutineInfo("refresh", database="AppDB", schema="dbo"), + RoutineInfo("refresh", database="AuditDB", schema="dbo"), + ] + sql = "EXEC ref" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == ["AppDB.dbo.refresh", "AuditDB.dbo.refresh"] + + +def test_database_qualified_parameter_completion() -> None: + routines = [ + RoutineInfo( + "refresh", + database="AppDB", + schema="dbo", + parameters=("@app_value",), + ), + RoutineInfo( + "refresh", + database="AuditDB", + schema="dbo", + parameters=("@audit_value",), + ), + ] + sql = "EXEC AuditDB.dbo.refresh audit" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == ["@audit_value"] + + +def test_unqualified_cross_schema_routine_requires_qualification() -> None: + routines = [ + RoutineInfo("refresh", schema="sales", parameters=("@sales_value",)), + RoutineInfo("refresh", schema="dbo", parameters=("@dbo_value",)), + ] + sql = "EXEC refresh dbo" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == [] + + +def test_unqualified_ambiguous_non_dbo_routine_has_no_parameters() -> None: + routines = [ + RoutineInfo("refresh", schema="sales", parameters=("@sales_value",)), + RoutineInfo("refresh", schema="reporting", parameters=("@report_value",)), + ] + sql = "EXEC refresh value" + + completions = get_completions(sql, len(sql), [], {}, routines) + + assert completions == [] + + +def test_rich_completion_metadata_uses_separate_cache_key() -> None: + inspector = MagicMock() + inspector.get_completion_routines = MagicMock() + + cache_field, loader = _completion_routine_loader(inspector) + + assert cache_field == "completion_routines" + assert loader is inspector.get_completion_routines + + +def test_drop_procedure_excludes_functions() -> None: + sql = "DROP PROCEDURE dt_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["dt_setpropertybyid"] + + +def test_drop_function_excludes_procedures() -> None: + sql = "DROP FUNCTION fn_" + + completions = get_completions(sql, len(sql), [], {}, _routines()) + + assert completions == ["fn_score", "fn_orders"] + + +def test_plain_legacy_procedures_are_not_scalar_function_suggestions() -> None: + sql = "SELECT get_" + + completions = get_completions(sql, len(sql), [], {}, ["get_user"]) + + assert "get_user" not in completions + + +def test_include_functions_false_hides_scalar_and_table_functions() -> None: + select_sql = "SELECT fn_" + from_sql = "SELECT * FROM fn_" + + assert "fn_score" not in get_completions( + select_sql, + len(select_sql), + [], + {}, + _routines(), + include_functions=False, + ) + assert "fn_orders" not in get_completions( + from_sql, + len(from_sql), + [], + {}, + _routines(), + include_functions=False, + )