fix: display result sets from stored procedure calls (CALL/EXEC/EXECUTE)

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

@@ -10,7 +10,7 @@
 if TYPE_CHECKING:
     from sqlit.domains.connections.domain.config import ConnectionConfig
 
-SELECT_KEYWORDS = frozenset(["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"])
+SELECT_KEYWORDS = frozenset(["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA", "CALL", "EXEC", "EXECUTE"])
 
 
 def resolve_file_path(path_str: str) -> Path:

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

@@ -14,6 +14,7 @@
     SequenceInfo,
     TableInfo,
     TriggerInfo,
+    _sanitize_row,
 )
 
 
@@ -37,7 +38,40 @@ def system_databases(self) -> frozenset[str]:
     def supports_stored_procedures(self) -> bool:
         return True
 
-    def apply_database_override(self, config: "ConnectionConfig", database: str) -> "ConnectionConfig":
+    def execute_query(self, conn: Any, query: str, max_rows: int | None = None) -> tuple[list[str], list[tuple], bool]:
+        """Execute a row-returning query, handling MySQL stored procedure result sets.
+
+        MySQL/MariaDB stored procedures can produce multiple result sets. The base
+        implementation only reads the first one and may leave the connection in a
+        "commands out of sync" state. This override skips status-only result sets,
+        returns the first result set that has columns, and consumes the rest.
+        """
+        cursor = conn.cursor()
+        cursor.execute(query)
+
+        # Stored procedures may produce a status-only result set before the data.
+        # Advance to the first result set that actually has columns.
+        while cursor.description is None and getattr(cursor, "nextset", lambda: False)():
+            pass
+
+        if cursor.description:
+            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
+                if truncated:
+                    rows = rows[:max_rows]
+            else:
+                rows = cursor.fetchall()
+                truncated = False
+            # Consume any remaining result sets to keep the connection clean.
+            while getattr(cursor, "nextset", lambda: False)():
+                pass
+            return columns, [_sanitize_row(row) for row in rows], truncated
+
+        return [], [], False
+
+    def apply_database_override(self, config: ConnectionConfig, database: str) -> ConnectionConfig:
         """Apply a default database for unqualified queries."""
         if not database:
             return config

sqlit/domains/query/app/query_service.py

@@ -16,7 +16,8 @@
     from sqlit.shared.core.protocols import HistoryStoreProtocol, QueryExecutorProtocol
 
 # Query types that return result sets (SELECT-like queries)
-SELECT_KEYWORDS = frozenset(["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"])
+# CALL/EXEC/EXECUTE are included so stored procedure calls that return result sets are displayed.
+SELECT_KEYWORDS = frozenset(["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA", "CALL", "EXEC", "EXECUTE"])
 
 # DML statements that may carry a RETURNING clause — when present they produce a result set.
 _DML_KEYWORDS = frozenset(["INSERT", "UPDATE", "DELETE", "MERGE"])

tests/unit/test_mysql_adapter.py

@@ -0,0 +1,124 @@
+"""Unit tests for MySQL/MariaDB adapter behavior."""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock
+
+from sqlit.domains.connections.providers.mysql.adapter import MySQLAdapter
+
+
+# DB-API column descriptions are tuples of 7-tuples; we only care about the name
+# in the first position, so a single-element tuple is sufficient for tests.
+def _make_cursor(result_sets: list[Any]) -> MagicMock:
+    """Build a mock cursor that yields the given result sets.
+
+    Each entry is (description, rows). description is a tuple of column
+    description tuples (DB-API style) or None to indicate a status-only
+    result set. nextset() advances to the next entry and returns True until
+    exhausted.
+    """
+    cursor = MagicMock()
+    cursor.description = result_sets[0][0]
+    cursor.fetchall.return_value = result_sets[0][1]
+    cursor.fetchmany.return_value = result_sets[0][1]
+
+    state: dict[str, Any] = {"index": 0, "sets": result_sets}
+
+    def _advance() -> bool:
+        state["index"] += 1
+        if state["index"] >= len(state["sets"]):
+            return False
+        desc, rows = state["sets"][state["index"]]
+        cursor.description = desc
+        cursor.fetchall.return_value = rows
+        cursor.fetchmany.return_value = rows
+        return True
+
+    cursor.nextset = MagicMock(side_effect=_advance)
+    return cursor
+
+
+def _desc(*names: str) -> tuple[tuple[str, ...], ...]:
+    """Build a DB-API style description tuple from column names."""
+    return tuple((name,) for name in names)
+
+
+def test_execute_query_returns_first_result_set() -> None:
+    adapter = MySQLAdapter()
+    cursor = _make_cursor([
+        (_desc("id", "name"), [(1, "Alice"), (2, "Bob")]),
+    ])
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+
+    columns, rows, truncated = adapter.execute_query(conn, "SELECT * FROM users")
+
+    assert columns == ["id", "name"]
+    assert rows == [(1, "Alice"), (2, "Bob")]
+    assert truncated is False
+
+
+def test_execute_query_skips_status_only_result_sets() -> None:
+    """Stored procedures may emit a status-only result set before the data."""
+    adapter = MySQLAdapter()
+    cursor = _make_cursor([
+        (None, []),
+        (_desc("id", "name"), [(1, "Alice")]),
+    ])
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+
+    columns, rows, _ = adapter.execute_query(conn, "CALL getalltheme()")
+
+    assert columns == ["id", "name"]
+    assert rows == [(1, "Alice")]
+
+
+def test_execute_query_consumes_remaining_result_sets() -> None:
+    """Remaining result sets should be consumed to keep the connection clean."""
+    adapter = MySQLAdapter()
+    cursor = _make_cursor([
+        (_desc("id", "name"), [(1, "Alice")]),
+        (_desc("code"), [("a",), ("b",)]),
+        (None, []),
+    ])
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+
+    adapter.execute_query(conn, "CALL getalltheme()")
+
+    # nextset should have been called until all result sets are consumed
+    assert cursor.nextset.call_count == 3
+
+
+def test_execute_query_with_max_rows_truncates() -> None:
+    adapter = MySQLAdapter()
+    cursor = _make_cursor([
+        (_desc("id"), [(1,), (2,), (3,)]),
+    ])
+    cursor.fetchmany.side_effect = lambda size: [(1,), (2,), (3,)][:size]
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+
+    columns, rows, truncated = adapter.execute_query(conn, "SELECT * FROM users", max_rows=2)
+
+    assert columns == ["id"]
+    assert rows == [(1,), (2,)]
+    assert truncated is True
+    cursor.fetchmany.assert_called_once_with(3)
+
+
+def test_execute_query_no_result_set_returns_empty() -> None:
+    adapter = MySQLAdapter()
+    cursor = _make_cursor([
+        (None, []),
+    ])
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+
+    columns, rows, truncated = adapter.execute_query(conn, "CALL update_only()")
+
+    assert columns == []
+    assert rows == []
+    assert truncated is False

tests/unit/test_sqlite_returning.py

@@ -23,32 +23,47 @@ def jobs_db(tmp_path: Path) -> Path:
     return db
 
 
-def test_classifier_recognizes_update_returning_as_returns_rows():
+def test_classifier_recognizes_update_returning_as_returns_rows() -> None:
     """`UPDATE ... RETURNING` produces a result set, so the analyzer must classify it as RETURNS_ROWS."""
     analyzer = KeywordQueryAnalyzer()
     sql = "UPDATE jobs SET status = status WHERE id = 1 RETURNING id"
     assert analyzer.classify(sql) == QueryKind.RETURNS_ROWS
 
 
-def test_classifier_recognizes_insert_returning_as_returns_rows():
+def test_classifier_recognizes_insert_returning_as_returns_rows() -> None:
     analyzer = KeywordQueryAnalyzer()
     sql = "INSERT INTO jobs (id, status) VALUES (3, 'new') RETURNING id"
     assert analyzer.classify(sql) == QueryKind.RETURNS_ROWS
 
 
-def test_classifier_recognizes_delete_returning_as_returns_rows():
+def test_classifier_recognizes_delete_returning_as_returns_rows() -> None:
     analyzer = KeywordQueryAnalyzer()
     sql = "DELETE FROM jobs WHERE id = 1 RETURNING id"
     assert analyzer.classify(sql) == QueryKind.RETURNS_ROWS
 
 
-def test_classifier_plain_update_is_non_query():
+def test_classifier_plain_update_is_non_query() -> None:
     """Plain DML without RETURNING must still be NON_QUERY (sanity check we don't over-correct)."""
     analyzer = KeywordQueryAnalyzer()
     assert analyzer.classify("UPDATE jobs SET status = 'done'") == QueryKind.NON_QUERY
 
 
-def test_sqlite_execute_query_runs_update_returning_and_persists(jobs_db: Path):
+def test_classifier_recognizes_call_as_returns_rows() -> None:
+    analyzer = KeywordQueryAnalyzer()
+    assert analyzer.classify("CALL getalltheme()") == QueryKind.RETURNS_ROWS
+
+
+def test_classifier_recognizes_exec_as_returns_rows() -> None:
+    analyzer = KeywordQueryAnalyzer()
+    assert analyzer.classify("EXEC getalltheme") == QueryKind.RETURNS_ROWS
+
+
+def test_classifier_recognizes_execute_as_returns_rows() -> None:
+    analyzer = KeywordQueryAnalyzer()
+    assert analyzer.classify("EXECUTE getalltheme") == QueryKind.RETURNS_ROWS
+
+
+def test_sqlite_execute_query_runs_update_returning_and_persists(jobs_db: Path) -> None:
     """UPDATE ... RETURNING via execute_query must return the row AND persist the change."""
     adapter = SQLiteAdapter()
     conn = sqlite3.connect(str(jobs_db))