Merge pull request #312 from Maxteabag/fix/issue-255-mssql-routine-completion

Add SQL Server routine autocomplete

Author: Maxteabag

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."""

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)

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<suffix>.*)$",
+        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)

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