diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 8e6a194..f796a63 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -324,14 +324,51 @@ CREATE TABLE IF NOT EXISTS "{}"( """.strip() -_TRANSACTION_CONTROL_PREFIXES = ( +_TRANSACTION_CONTROL_KEYWORDS = { "BEGIN", "COMMIT", "END", "ROLLBACK", "SAVEPOINT", "RELEASE", -) +} + +# Statements that never return rows and cannot run inside (or would break +# out of) the savepoint guard used by query() +_QUERY_REJECTED_KEYWORDS = _TRANSACTION_CONTROL_KEYWORDS | { + "VACUUM", + "ATTACH", + "DETACH", +} + + +def _first_keyword(sql: str) -> str: + """ + Return the first keyword of a SQL statement, uppercased, skipping any + leading whitespace and ``--`` or ``/* ... */`` comments - the only + things SQLite's tokenizer allows before the first token. Returns an + empty string if there is no leading keyword. + """ + i, n = 0, len(sql) + while i < n: + if sql[i].isspace(): + i += 1 + elif sql.startswith("--", i): + newline = sql.find("\n", i) + if newline == -1: + return "" + i = newline + 1 + elif sql.startswith("/*", i): + end = sql.find("*/", i + 2) + if end == -1: + return "" + i = end + 2 + else: + break + j = i + while j < n and (sql[j].isalpha() or sql[j] == "_"): + j += 1 + return sql[i:j].upper() class Database: @@ -667,16 +704,14 @@ class Database: "query() can only be used with SQL that returns rows - " "use execute() for other statements" ) - prefix = sql.lstrip().upper() - if prefix.startswith( - _TRANSACTION_CONTROL_PREFIXES + ("VACUUM", "ATTACH", "DETACH") - ): + keyword = _first_keyword(sql) + if keyword in _QUERY_REJECTED_KEYWORDS: # None of these return rows - reject them without executing anything raise ValueError(message) if self._tracer: self._tracer(sql, params) args: tuple = (params,) if params is not None else () - if prefix.startswith("PRAGMA"): + if keyword == "PRAGMA": # Some PRAGMA statements refuse to run inside a transaction, so # execute these without the savepoint guard used below. PRAGMAs # never open an implicit transaction, so there is nothing to @@ -741,7 +776,7 @@ class Database: not was_in_transaction and self.conn.in_transaction and cursor.description is None - and not sql.lstrip().upper().startswith(_TRANSACTION_CONTROL_PREFIXES) + and _first_keyword(sql) not in _TRANSACTION_CONTROL_KEYWORDS ): # The statement opened an implicit transaction - commit it, so # that execute() behaves consistently with the rest of the diff --git a/tests/test_atomic.py b/tests/test_atomic.py index a8e3563..f75c4ee 100644 --- a/tests/test_atomic.py +++ b/tests/test_atomic.py @@ -247,6 +247,17 @@ def test_execute_write_respects_explicit_transaction(fresh_db): assert [r["id"] for r in fresh_db["t"].rows] == [1] +def test_execute_comment_prefixed_begin_leaves_transaction_open(fresh_db): + # A BEGIN hidden behind a leading comment must not be auto-committed + # out from under the caller + fresh_db["t"].insert({"id": 1}, pk="id") + fresh_db.execute("-- start a transaction\nbegin") + assert fresh_db.conn.in_transaction + fresh_db.execute("insert into t (id) values (2)") + fresh_db.rollback() + assert [r["id"] for r in fresh_db["t"].rows] == [1] + + def test_query_returning_commits_after_iteration(tmpdir): if sqlite3.sqlite_version_info < (3, 35, 0): import pytest as _pytest diff --git a/tests/test_query.py b/tests/test_query.py index cf8d9f5..1d5c1f6 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -48,7 +48,18 @@ def test_query_rejected_write_inside_transaction_is_rolled_back(fresh_db): @pytest.mark.parametrize( - "sql", ["begin", "commit", "rollback", "vacuum", "detach database foo"] + "sql", + [ + "begin", + "commit", + "rollback", + "vacuum", + "detach database foo", + "/* comment */ commit", + "-- comment\nbegin", + "/* multi\nline */ -- and another\n vacuum", + "\t /* a */ /* b */ savepoint s1", + ], ) def test_query_rejects_transaction_control_and_vacuum(fresh_db, sql): with pytest.raises(ValueError) as ex: @@ -57,6 +68,21 @@ def test_query_rejects_transaction_control_and_vacuum(fresh_db, sql): assert not fresh_db.conn.in_transaction +def test_query_comment_prefixed_commit_does_not_commit_transaction(fresh_db): + # A COMMIT hidden behind a leading comment must not slip past the + # keyword check - previously it committed the caller's open + # transaction before the ValueError was raised + fresh_db["dogs"].insert({"name": "Cleo"}) + fresh_db.begin() + fresh_db.execute("insert into dogs (name) values ('Pancakes')") + with pytest.raises(ValueError): + fresh_db.query("/* comment */ COMMIT") + # The explicit transaction is still open and can still be rolled back + assert fresh_db.conn.in_transaction + fresh_db.rollback() + assert [row["name"] for row in fresh_db["dogs"].rows] == ["Cleo"] + + def test_query_error_leaves_no_transaction_open(fresh_db): with pytest.raises(sqlite3.OperationalError): fresh_db.query("select * from missing_table") @@ -75,6 +101,41 @@ def test_query_pragma(tmpdir): db.close() +def test_query_comment_prefixed_pragma(tmpdir): + from sqlite_utils import Database + + db = Database(str(tmpdir / "test.db")) + # A leading comment must not stop a PRAGMA being recognized as one - + # previously it was executed inside the savepoint guard, where + # journal mode changes are refused + assert list(db.query("-- set WAL mode\npragma journal_mode = wal")) == [ + {"journal_mode": "wal"} + ] + db.close() + + +@pytest.mark.parametrize( + "sql,expected", + [ + ("select 1", "SELECT"), + (" \t\n select 1", "SELECT"), + ("-- comment\nbegin", "BEGIN"), + ("/* one */ /* two */ pragma user_version", "PRAGMA"), + ("/* multi\nline */vacuum", "VACUUM"), + ("insert into t values (1)", "INSERT"), + ("-- only a comment", ""), + ("/* unterminated", ""), + ("", ""), + (" ", ""), + ("123", ""), + ], +) +def test_first_keyword(sql, expected): + from sqlite_utils.db import _first_keyword + + assert _first_keyword(sql) == expected + + @pytest.mark.skipif( sqlite3.sqlite_version_info < (3, 35, 0), reason="RETURNING requires SQLite 3.35.0 or higher",