Replace SQL prefix matching with a first-token scanner

query() and execute() classified statements using
sql.lstrip().upper().startswith(...), which a leading SQL comment
defeated: db.query('/* c */ COMMIT') inside db.atomic() committed
the caller's transaction and masked the ValueError with a 'no such
savepoint' error, comment-prefixed PRAGMAs ran inside the savepoint
guard where journal mode changes are refused, and a comment-prefixed
BEGIN passed to db.execute() was instantly auto-committed.

The new _first_keyword() helper skips leading whitespace and -- or
/* */ comments - the only things SQLite's tokenizer allows before
the first token - and returns that token for exact comparison, so
keyword detection now matches what SQLite itself will see.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

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

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

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",