diff --git a/sqlit/domains/connections/providers/sqlite/adapter.py b/sqlit/domains/connections/providers/sqlite/adapter.py index 2b7315f4..f59dbdaf 100644 --- a/sqlit/domains/connections/providers/sqlite/adapter.py +++ b/sqlit/domains/connections/providers/sqlite/adapter.py @@ -220,6 +220,9 @@ def execute_query(self, conn: Any, query: str, max_rows: int | None = None) -> t else: rows = cursor.fetchall() truncated = False + # DML with RETURNING produces a result set but also writes — persist it. + if conn.in_transaction: + conn.commit() return columns, [tuple(row) for row in rows], truncated return [], [], False diff --git a/sqlit/domains/query/app/query_service.py b/sqlit/domains/query/app/query_service.py index b458a5bd..e22fe345 100644 --- a/sqlit/domains/query/app/query_service.py +++ b/sqlit/domains/query/app/query_service.py @@ -18,6 +18,10 @@ # Query types that return result sets (SELECT-like queries) SELECT_KEYWORDS = frozenset(["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"]) +# DML statements that may carry a RETURNING clause — when present they produce a result set. +_DML_KEYWORDS = frozenset(["INSERT", "UPDATE", "DELETE", "MERGE"]) +_RETURNING_RE = re.compile(r"(?is)\bRETURNING\b\s+\S") + # Regex for parsing USE database statements # Matches: USE dbname, USE [dbname], USE `dbname`, USE "dbname" _USE_PATTERN = re.compile( @@ -93,7 +97,12 @@ def classify(self, query: str) -> QueryKind: if non_comment_lines: first_line = non_comment_lines[0].upper() first_word = first_line.split()[0] if first_line else "" - return QueryKind.RETURNS_ROWS if first_word in SELECT_KEYWORDS else QueryKind.NON_QUERY + if first_word in SELECT_KEYWORDS: + return QueryKind.RETURNS_ROWS + # DML with a RETURNING clause produces a result set too. + if first_word in _DML_KEYWORDS and _RETURNING_RE.search(stmt): + return QueryKind.RETURNS_ROWS + return QueryKind.NON_QUERY return QueryKind.NON_QUERY diff --git a/tests/unit/test_sqlite_returning.py b/tests/unit/test_sqlite_returning.py new file mode 100644 index 00000000..bb58850d --- /dev/null +++ b/tests/unit/test_sqlite_returning.py @@ -0,0 +1,71 @@ +"""Regression tests for issue #147: SQLite UPDATE ... RETURNING crashes on commit.""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +import pytest + +from sqlit.domains.connections.providers.sqlite.adapter import SQLiteAdapter +from sqlit.domains.query.app.query_service import KeywordQueryAnalyzer, QueryKind + + +@pytest.fixture +def jobs_db(tmp_path: Path) -> Path: + """A tiny SQLite DB with a `jobs` table for RETURNING tests.""" + db = tmp_path / "jobs.db" + conn = sqlite3.connect(str(db)) + conn.execute("CREATE TABLE jobs (id INTEGER PRIMARY KEY, status TEXT)") + conn.executemany("INSERT INTO jobs (id, status) VALUES (?, ?)", [(1, "new"), (2, "new")]) + conn.commit() + conn.close() + return db + + +def test_classifier_recognizes_update_returning_as_returns_rows(): + """`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(): + 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(): + 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(): + """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): + """UPDATE ... RETURNING via execute_query must return the row AND persist the change.""" + adapter = SQLiteAdapter() + conn = sqlite3.connect(str(jobs_db)) + try: + columns, rows, _ = adapter.execute_query( + conn, + "UPDATE jobs SET status = 'done' WHERE id = 1 RETURNING id, status", + ) + assert columns == ["id", "status"] + assert rows == [(1, "done")] + finally: + conn.close() + + # Verify the write was actually committed by opening a fresh connection. + verify = sqlite3.connect(str(jobs_db)) + try: + result = verify.execute("SELECT status FROM jobs WHERE id = 1").fetchone() + assert result == ("done",), f"UPDATE ... RETURNING did not persist; got {result!r}" + finally: + verify.close()