From 9d8ca533a89468dfd4b844151ea6ff478cf03b1f Mon Sep 17 00:00:00 2001 From: Peter Adams <18162810+Maxteabag@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:21:22 +0200 Subject: [PATCH] fix: display stored procedure result sets --- .../connections/providers/adapters/base.py | 30 ++++- .../connections/providers/mssql/adapter.py | 44 +++++++ .../connections/providers/mysql/base.py | 40 ++++++ tests/unit/test_stored_procedure_results.py | 123 ++++++++++++++++++ 4 files changed, 235 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_stored_procedure_results.py diff --git a/sqlit/domains/connections/providers/adapters/base.py b/sqlit/domains/connections/providers/adapters/base.py index 078ec0f0..573dc475 100644 --- a/sqlit/domains/connections/providers/adapters/base.py +++ b/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. diff --git a/sqlit/domains/connections/providers/mssql/adapter.py b/sqlit/domains/connections/providers/mssql/adapter.py index 5a3fd10b..4880a5d3 100644 --- a/sqlit/domains/connections/providers/mssql/adapter.py +++ b/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: diff --git a/sqlit/domains/connections/providers/mysql/base.py b/sqlit/domains/connections/providers/mysql/base.py index 867dd790..d5343722 100644 --- a/sqlit/domains/connections/providers/mysql/base.py +++ b/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: diff --git a/tests/unit/test_stored_procedure_results.py b/tests/unit/test_stored_procedure_results.py new file mode 100644 index 00000000..0f7b79f1 --- /dev/null +++ b/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