fix: display stored procedure result sets

sqlit/domains/connections/providers/adapters/base.py

@@ -13,6 +13,33 @@
 SELECT_KEYWORDS = frozenset(["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"])
 
 
+def _strip_leading_sql_comments(query: str) -> str:
+    """Strip whitespace and consecutive SQL comments from a statement."""
+    remaining = query.lstrip()
+    while remaining:
+        if remaining.startswith("--") or remaining.startswith("#"):
+            newline = remaining.find("\n")
+            if newline < 0:
+                return ""
+            remaining = remaining[newline + 1 :].lstrip()
+            continue
+        if remaining.startswith("/*"):
+            end = remaining.find("*/", 2)
+            if end < 0:
+                return ""
+            remaining = remaining[end + 2 :].lstrip()
+            continue
+        break
+    return remaining
+
+
+def _first_keyword(query: str) -> str:
+    """Return the first SQL token, ignoring whitespace and leading comments."""
+    remaining = _strip_leading_sql_comments(query)
+    parts = remaining.split(maxsplit=1)
+    return parts[0].rstrip(";").upper() if parts else ""
+
+
 def resolve_file_path(path_str: str) -> Path:
     """Resolve a file path for file-based databases (SQLite, DuckDB).
 
@@ -254,8 +281,7 @@ def test_query(self) -> str:
 
     def classify_query(self, query: str) -> bool:
         """Return True if the query is expected to return rows."""
-        query_type = query.strip().upper().split()[0] if query.strip() else ""
-        return query_type in SELECT_KEYWORDS
+        return _first_keyword(query) in SELECT_KEYWORDS
 
     def execute_test_query(self, conn: Any) -> None:
         """Execute a simple query to verify the connection works.

sqlit/domains/connections/providers/mssql/adapter.py

@@ -13,6 +13,7 @@
     SequenceInfo,
     TableInfo,
     TriggerInfo,
+    _strip_leading_sql_comments,
 )
 from sqlit.domains.connections.providers.tls import (
     TLS_MODE_DEFAULT,
@@ -124,6 +125,30 @@ def supports_cross_database_queries(self) -> bool:
     def supports_stored_procedures(self) -> bool:
         return True
 
+    def classify_query(self, query: str) -> bool:
+        """Treat T-SQL procedure execution as potentially row-returning."""
+        return self._is_procedure_call(query) or super().classify_query(query)
+
+    @staticmethod
+    def _is_procedure_call(query: str) -> bool:
+        """Distinguish procedure calls from EXECUTE AS and dynamic SQL."""
+        statement = _strip_leading_sql_comments(query)
+        parts = statement.split(maxsplit=1)
+        if not parts or parts[0].upper() not in {"EXEC", "EXECUTE"} or len(parts) == 1:
+            return False
+        target = _strip_leading_sql_comments(parts[1])
+        if not target:
+            return False
+        upper_target = target.upper()
+        if upper_target.split(maxsplit=1)[0] == "AS":
+            return False
+        if target.startswith(("(", "'", '"')) or upper_target.startswith("N'"):
+            return False
+        if target.startswith("@"):
+            _return_variable, separator, procedure = target.partition("=")
+            return bool(separator and procedure.strip())
+        return True
+
     @property
     def system_databases(self) -> frozenset[str]:
         return frozenset({"master", "tempdb", "model", "msdb"})
@@ -645,6 +670,25 @@ def execute_query(self, conn: Any, query: str, max_rows: int | None = None) -> t
         """Execute a query on SQL Server with optional row limit."""
         cursor = conn.cursor()
         cursor.execute(query)
+        if self._is_procedure_call(query):
+            result: tuple[list[str], list[tuple], bool] | None = None
+            while True:
+                if cursor.description and result is None:
+                    columns = [col[0] for col in cursor.description]
+                    if max_rows is not None:
+                        rows = cursor.fetchmany(max_rows + 1)
+                        truncated = len(rows) > max_rows
+                        rows = rows[:max_rows]
+                    else:
+                        rows = cursor.fetchall()
+                        truncated = False
+                    result = (columns, [tuple(row) for row in rows], truncated)
+
+                nextset = getattr(cursor, "nextset", None)
+                if not callable(nextset) or not nextset():
+                    break
+            return result or ([], [], False)
+
         if cursor.description:
             columns = [col[0] for col in cursor.description]
             if max_rows is not None:

sqlit/domains/connections/providers/mysql/base.py

@@ -15,6 +15,8 @@
     SequenceInfo,
     TableInfo,
     TriggerInfo,
+    _first_keyword,
+    _sanitize_row,
 )
 
 
@@ -42,6 +44,44 @@ def supports_stored_procedures(self) -> bool:
     def supports_foreign_keys(self) -> bool:
         return True
 
+    def classify_query(self, query: str) -> bool:
+        """Treat MySQL/MariaDB CALL statements as potentially row-returning."""
+        return _first_keyword(query) == "CALL" or super().classify_query(query)
+
+    def execute_query(
+        self, conn: Any, query: str, max_rows: int | None = None
+    ) -> tuple[list[str], list[tuple], bool]:
+        """Return the first row-bearing result from a stored procedure call.
+
+        MySQL-compatible drivers expose procedure output as multiple DB-API
+        result sets, which can include leading and trailing status-only sets.
+        Every set must be consumed before the connection can be reused.
+        """
+        if _first_keyword(query) != "CALL":
+            return super().execute_query(conn, query, max_rows)
+
+        cursor = conn.cursor()
+        cursor.execute(query)
+        result: tuple[list[str], list[tuple], bool] | None = None
+
+        while True:
+            if cursor.description and result is None:
+                columns = [col[0] for col in cursor.description]
+                if max_rows is None:
+                    rows = cursor.fetchall()
+                    truncated = False
+                else:
+                    rows = cursor.fetchmany(max_rows + 1)
+                    truncated = len(rows) > max_rows
+                    rows = rows[:max_rows]
+                result = (columns, [_sanitize_row(row) for row in rows], truncated)
+
+            nextset = getattr(cursor, "nextset", None)
+            if not callable(nextset) or not nextset():
+                break
+
+        return result or ([], [], False)
+
     def apply_database_override(self, config: ConnectionConfig, database: str) -> ConnectionConfig:
         """Apply a default database for unqualified queries."""
         if not database:

tests/unit/test_stored_procedure_results.py

@@ -0,0 +1,123 @@
+"""Regression coverage for stored procedures that return DB-API result sets."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from sqlit.domains.connections.providers.mssql.adapter import SQLServerAdapter
+from sqlit.domains.connections.providers.mysql.adapter import MySQLAdapter
+from sqlit.domains.query.app.query_service import DialectQueryAnalyzer, QueryResult, QueryService
+
+
+class ResultSetCursor:
+    def __init__(self, result_sets: list[tuple[Any, list[tuple[Any, ...]]]]) -> None:
+        self._sets = result_sets
+        self._index = 0
+        self.description = result_sets[0][0]
+        self.executed: str | None = None
+        self.nextset_calls = 0
+
+    def execute(self, query: str) -> None:
+        self.executed = query
+
+    def fetchall(self) -> list[tuple[Any, ...]]:
+        return self._sets[self._index][1]
+
+    def fetchmany(self, size: int) -> list[tuple[Any, ...]]:
+        return self._sets[self._index][1][:size]
+
+    def nextset(self) -> bool:
+        self.nextset_calls += 1
+        self._index += 1
+        if self._index >= len(self._sets):
+            return False
+        self.description = self._sets[self._index][0]
+        return True
+
+
+def _description(*names: str) -> tuple[tuple[str], ...]:
+    return tuple((name,) for name in names)
+
+
+@pytest.mark.parametrize(
+    ("adapter", "statement"),
+    [
+        (MySQLAdapter(), "CALL get_users()"),
+        (MySQLAdapter(), "CALL\nget_users()"),
+        (MySQLAdapter(), "-- report\nCALL get_users()"),
+        (SQLServerAdapter(), "EXEC get_users"),
+        (SQLServerAdapter(), "EXEC\tget_users"),
+        (SQLServerAdapter(), "/* report */ EXEC get_users"),
+        (SQLServerAdapter(), "EXECUTE get_users"),
+    ],
+)
+def test_procedure_call_is_routed_to_row_execution(adapter: Any, statement: str) -> None:
+    cursor = ResultSetCursor([(None, []), (_description("id", "name"), [(1, "Alice")]), (None, [])])
+    connection = MagicMock()
+    connection.cursor.return_value = cursor
+    executor = MagicMock(wraps=adapter)
+    service = QueryService(analyzer=DialectQueryAnalyzer(adapter))
+
+    result = service.execute(connection, executor, statement, save_to_history=False)
+
+    assert result == QueryResult(columns=["id", "name"], rows=[(1, "Alice")], row_count=1, truncated=False)
+    executor.execute_query.assert_called_once_with(connection, statement, None)
+    executor.execute_non_query.assert_not_called()
+    assert cursor.nextset_calls == 3
+
+
+@pytest.mark.parametrize(
+    ("adapter", "statement"),
+    [(MySQLAdapter(), "CALL update_users()"), (SQLServerAdapter(), "EXEC update_users")],
+)
+def test_procedure_without_rows_consumes_all_status_sets(adapter: Any, statement: str) -> None:
+    cursor = ResultSetCursor([(None, []), (None, [])])
+    connection = MagicMock()
+    connection.cursor.return_value = cursor
+
+    columns, rows, truncated = adapter.execute_query(connection, statement)
+
+    assert (columns, rows, truncated) == ([], [], False)
+    assert cursor.nextset_calls == 2
+
+
+def test_procedure_results_respect_row_limit_and_drain_connection() -> None:
+    adapter = MySQLAdapter()
+    cursor = ResultSetCursor([(_description("id"), [(1,), (2,), (3,)]), (None, [])])
+    connection = MagicMock()
+    connection.cursor.return_value = cursor
+
+    columns, rows, truncated = adapter.execute_query(connection, "CALL get_users()", max_rows=2)
+
+    assert columns == ["id"]
+    assert rows == [(1,), (2,)]
+    assert truncated is True
+    assert cursor.nextset_calls == 2
+
+
+def test_stored_procedure_keywords_are_provider_specific() -> None:
+    """CALL must not globally bypass non-query execution for other providers."""
+    from sqlit.domains.connections.providers.firebird.adapter import FirebirdAdapter
+    from sqlit.domains.connections.providers.oracle.adapter import OracleAdapter
+
+    assert MySQLAdapter().classify_query("CALL update_users()") is True
+    assert SQLServerAdapter().classify_query("EXEC update_users") is True
+    assert OracleAdapter().classify_query("CALL update_users()") is False
+    assert FirebirdAdapter().classify_query("EXECUTE PROCEDURE update_users") is False
+
+
+@pytest.mark.parametrize(
+    "statement",
+    [
+        "EXECUTE AS USER = 'reporter'",
+        "EXECUTE AS\nUSER = 'reporter'",
+        "EXECUTE /* context */ AS\tLOGIN = 'reporter'",
+        "EXECUTE ('UPDATE users SET active = 1')",
+        "EXEC @sql",
+    ],
+)
+def test_sql_server_non_procedure_execute_forms_remain_non_queries(statement: str) -> None:
+    assert SQLServerAdapter().classify_query(statement) is False