diff --git a/sqlit/domains/connections/providers/adapters/base.py b/sqlit/domains/connections/providers/adapters/base.py index 87770b1f..e8700f9d 100644 --- a/sqlit/domains/connections/providers/adapters/base.py +++ b/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: diff --git a/sqlit/domains/connections/providers/mysql/base.py b/sqlit/domains/connections/providers/mysql/base.py index d65fcb90..2bbeeb53 100644 --- a/sqlit/domains/connections/providers/mysql/base.py +++ b/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 diff --git a/sqlit/domains/query/app/query_service.py b/sqlit/domains/query/app/query_service.py index e22fe345..bce13109 100644 --- a/sqlit/domains/query/app/query_service.py +++ b/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"]) diff --git a/tests/unit/test_mysql_adapter.py b/tests/unit/test_mysql_adapter.py new file mode 100644 index 00000000..39693e7d --- /dev/null +++ b/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 diff --git a/tests/unit/test_sqlite_returning.py b/tests/unit/test_sqlite_returning.py index bb58850d..8d33b715 100644 --- a/tests/unit/test_sqlite_returning.py +++ b/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))