Split out _iter_complete_sql_statements

To make it clearer why we are avoiding executescript()

docs/changelog.rst

@@ -9,6 +9,7 @@
 Unreleased
 ----------
 
+- New ``db.atomic()`` context manager for transactions, with nested transaction support using SQLite savepoints. Internal multi-step operations such as ``table.transform()`` now use this mechanism to avoid unexpectedly committing an existing transaction.
 - New :ref:`database migrations system <migrations>`, incorporating functionality that was previously provided by the separate `sqlite-migrate <https://github.com/simonw/sqlite-migrate>`__ plugin. Define migration sets using the new :class:`sqlite_utils.Migrations` class and apply them using the ``sqlite-utils migrate`` command or the :ref:`migrations Python API <migrations_python>`.  (:issue:`752`)
 
 .. _v3_39:

docs/python-api.rst

@@ -239,6 +239,35 @@ The ``db.execute()`` and ``db.executescript()`` methods provide wrappers around
 
 Other cursor methods such as ``.fetchone()`` and ``.fetchall()`` are also available, see the `standard library documentation <https://docs.python.org/3/library/sqlite3.html#sqlite3.Cursor>`__.
 
+.. _python_api_atomic:
+
+Transactions with db.atomic()
+-----------------------------
+
+Use ``db.atomic()`` to group multiple operations in a transaction:
+
+.. code-block:: python
+
+    with db.atomic():
+        db.table("dogs").insert({"id": 1, "name": "Cleo"}, pk="id")
+        db.table("dogs").insert({"id": 2, "name": "Pancakes"})
+
+If an exception is raised, changes made inside the block will be rolled back.
+
+``db.atomic()`` can be nested. Nested blocks use SQLite savepoints, so an exception in an inner block can roll back to that savepoint without rolling back the entire outer transaction:
+
+.. code-block:: python
+
+    with db.atomic():
+        db.table("dogs").insert({"id": 1, "name": "Cleo"}, pk="id")
+        try:
+            with db.atomic():
+                db.table("dogs").insert({"id": 2, "name": "Pancakes"})
+                raise ValueError("skip this one")
+        except ValueError:
+            pass
+        db.table("dogs").insert({"id": 3, "name": "Marnie"})
+
 .. _python_api_parameters:
 
 Passing parameters
@@ -955,7 +984,7 @@ You can delete all records in a table that match a specific WHERE statement usin
 
     >>> db = sqlite_utils.Database("dogs.db")
     >>> # Delete every dog with age less than 3
-    >>> with db.conn:
+    >>> with db.atomic():
     >>>     db.table("dogs").delete_where("age < ?", [3])
 
 Calling ``table.delete_where()`` with no other arguments will delete every row in the table.

sqlite_utils/db.py

@@ -196,6 +196,19 @@ DEFAULT = Default()
 Tracer = Callable[[str, Optional[Union[Sequence[Any], Dict[str, Any]]]], None]
 
 
+def _iter_complete_sql_statements(sql: str) -> Generator[str, None, None]:
+    statement = []
+    for char in sql:
+        statement.append(char)
+        statement_sql = "".join(statement).strip()
+        if statement_sql and sqlite3.complete_statement(statement_sql):
+            yield statement_sql
+            statement = []
+    statement_sql = "".join(statement).strip()
+    if statement_sql:
+        yield statement_sql
+
+
 COLUMN_TYPE_MAPPING: Dict[Any, str] = {
     float: "REAL",
     int: "INTEGER",
@@ -406,6 +419,38 @@ class Database:
         "Close the SQLite connection, and the underlying database file"
         self.conn.close()
 
+    @contextlib.contextmanager
+    def atomic(self) -> Generator["Database", None, None]:
+        """
+        Context manager for wrapping multiple database operations in a transaction.
+
+        Nested blocks use SQLite savepoints.
+        """
+        if self.conn.in_transaction:
+            savepoint = "sqlite_utils_{}".format(secrets.token_hex(16))
+            self.conn.execute("SAVEPOINT {};".format(savepoint))
+            try:
+                yield self
+            except BaseException:
+                self.conn.execute("ROLLBACK TO SAVEPOINT {};".format(savepoint))
+                self.conn.execute("RELEASE SAVEPOINT {};".format(savepoint))
+                raise
+            else:
+                self.conn.execute("RELEASE SAVEPOINT {};".format(savepoint))
+        else:
+            self.conn.execute("BEGIN")
+            try:
+                yield self
+            except BaseException:
+                self.conn.rollback()
+                raise
+            else:
+                try:
+                    self.conn.commit()
+                except BaseException:
+                    self.conn.rollback()
+                    raise
+
     @contextlib.contextmanager
     def ensure_autocommit_off(self) -> Generator[None, None, None]:
         """
@@ -581,6 +626,15 @@ class Database:
         """
         if self._tracer:
             self._tracer(sql, None)
+        return self._executescript(sql)
+
+    def _executescript(self, sql: str) -> sqlite3.Cursor:
+        if self.conn.in_transaction:
+            cursor = self.conn.cursor()
+            # avoid sqlite3.executescript()'s implicit commit:
+            for statement in _iter_complete_sql_statements(sql):
+                cursor.execute(statement)
+            return cursor
         return self.conn.executescript(sql)
 
     def table(self, table_name: str, **kwargs: Any) -> "Table":
@@ -729,7 +783,7 @@ class Database:
         if not hasattr(self, "_supports_strict"):
             try:
                 table_name = "t{}".format(secrets.token_hex(16))
-                with self.conn:
+                with self.atomic():
                     self.conn.execute(
                         "create table {} (name text) strict".format(table_name)
                     )
@@ -745,7 +799,7 @@ class Database:
         if not hasattr(self, "_supports_on_conflict"):
             table_name = "t{}".format(secrets.token_hex(16))
             try:
-                with self.conn:
+                with self.atomic():
                     self.conn.execute(
                         "create table {} (id integer primary key, name text)".format(
                             table_name
@@ -797,7 +851,7 @@ class Database:
                 self.execute("PRAGMA journal_mode=delete;")
 
     def _ensure_counts_table(self) -> None:
-        with self.conn:
+        with self.atomic():
             self.execute(_COUNTS_TABLE_CREATE_SQL.format(self._counts_table_name))
 
     def enable_counts(self) -> None:
@@ -833,7 +887,7 @@ class Database:
     def reset_counts(self) -> None:
         "Re-calculate cached counts for tables."
         tables = [table for table in self.tables if table.has_counts_triggers]
-        with self.conn:
+        with self.atomic():
             self._ensure_counts_table()
             counts_table = self.table(self._counts_table_name)
             counts_table.delete_where()
@@ -1288,6 +1342,7 @@ class Database:
         for table, fks in by_table.items():
             self.table(table).transform(add_foreign_keys=fks)
 
+        if not self.conn.in_transaction:
             self.vacuum()
 
     def index_foreign_keys(self) -> None:
@@ -1820,7 +1875,7 @@ class Table(Queryable):
             self._defaults["strict"] = strict
 
         columns = {name: value for (name, value) in columns.items()}
-        with self.db.conn:
+        with self.db.atomic():
             self.db.create_table(
                 self.name,
                 columns,
@@ -1848,7 +1903,7 @@ class Table(Queryable):
         """
         if not self.exists():
             raise NoTable(f"Table {self.name} does not exist")
-        with self.db.conn:
+        with self.db.atomic():
             sql = "CREATE TABLE {} AS SELECT * FROM {};".format(
                 quote_identifier(new_name),
                 quote_identifier(self.name),
@@ -1905,20 +1960,40 @@ class Table(Queryable):
             column_order=column_order,
             keep_table=keep_table,
         )
-        pragma_foreign_keys_was_on = self.db.execute("PRAGMA foreign_keys").fetchone()[
+        pragma_foreign_keys_was_on = bool(
-            0
+            self.db.execute("PRAGMA foreign_keys").fetchone()[0]
-        ]
+        )
+        already_in_transaction = self.db.conn.in_transaction
+        should_disable_foreign_keys = (
+            pragma_foreign_keys_was_on and not already_in_transaction
+        )
+        should_defer_foreign_keys = (
+            pragma_foreign_keys_was_on and already_in_transaction
+        )
+        defer_foreign_keys_was_on = False
         try:
-            if pragma_foreign_keys_was_on:
+            if should_disable_foreign_keys:
                 self.db.execute("PRAGMA foreign_keys=0;")
-            with self.db.conn:
+            elif should_defer_foreign_keys:
+                defer_foreign_keys_was_on = bool(
+                    self.db.execute("PRAGMA defer_foreign_keys").fetchone()[0]
+                )
+                if not defer_foreign_keys_was_on:
+                    self.db.execute("PRAGMA defer_foreign_keys=ON;")
+            with self.db.atomic():
                 for sql in sqls:
                     self.db.execute(sql)
                 # Run the foreign_key_check before we commit
                 if pragma_foreign_keys_was_on:
-                    self.db.execute("PRAGMA foreign_key_check;")
+                    foreign_key_violations = self.db.execute(
+                        "PRAGMA foreign_key_check;"
+                    ).fetchall()
+                    if foreign_key_violations:
+                        raise sqlite3.IntegrityError("FOREIGN KEY constraint failed")
         finally:
-            if pragma_foreign_keys_was_on:
+            if should_defer_foreign_keys and not defer_foreign_keys_was_on:
+                self.db.execute("PRAGMA defer_foreign_keys=OFF;")
+            if should_disable_foreign_keys:
                 self.db.execute("PRAGMA foreign_keys=1;")
         return self
 
@@ -2158,6 +2233,7 @@ class Table(Queryable):
                     columns, list(self.columns_dict.keys())
                 )
             )
+        with self.db.atomic():
             table = table or "_".join(columns)
             lookup_table = self.db.table(table)
             fk_column = fk_column or "{}_id".format(table)
@@ -2194,7 +2270,9 @@ class Table(Queryable):
                 "INSERT OR IGNORE INTO {} ({lookup_columns}) SELECT DISTINCT {table_cols} FROM {}".format(
                     quote_identifier(table),
                     quote_identifier(self.name),
-                lookup_columns=", ".join(quote_identifier(c) for c in lookup_columns),
+                    lookup_columns=", ".join(
+                        quote_identifier(c) for c in lookup_columns
+                    ),
                     table_cols=", ".join(quote_identifier(c) for c in columns),
                 )
             )
@@ -2521,8 +2599,8 @@ class Table(Queryable):
                 ),
             )
         )
-        with self.db.conn:
+        with self.db.atomic():
-            self.db.conn.executescript(sql)
+            self.db._executescript(sql)
         self.db.use_counts_table = True
 
     @property
@@ -2663,7 +2741,7 @@ class Table(Queryable):
         trigger_names = []
         for row in self.db.execute(sql).fetchall():
             trigger_names.append(row[0])
-        with self.db.conn:
+        with self.db.atomic():
             for trigger_name in trigger_names:
                 self.db.execute(
                     "DROP TRIGGER IF EXISTS {}".format(quote_identifier(trigger_name))
@@ -2862,7 +2940,7 @@ class Table(Queryable):
         sql = "delete from {} where {wheres}".format(
             quote_identifier(self.name), wheres=" and ".join(wheres)
         )
-        with self.db.conn:
+        with self.db.atomic():
             self.db.execute(sql, pk_values)
         return self
 
@@ -2935,7 +3013,7 @@ class Table(Queryable):
             sets=", ".join(sets),
             wheres=" and ".join(wheres),
         )
-        with self.db.conn:
+        with self.db.atomic():
             try:
                 rowcount = self.db.execute(sql, args).rowcount
             except OperationalError as e:
@@ -3024,7 +3102,7 @@ class Table(Queryable):
                 ),
                 where=" where {}".format(where) if where is not None else "",
             )
-            with self.db.conn:
+            with self.db.atomic():
                 self.db.execute(sql, where_args or [])
                 if drop:
                     self.transform(drop=columns)
@@ -3072,7 +3150,7 @@ class Table(Queryable):
         with progressbar(
             length=self.count, silent=not show_progress, label="2: Updating"
         ) as bar:
-            with self.db.conn:
+            with self.db.atomic():
                 for pk, updates in pk_to_values.items():
                     self.update(pk, updates)
                     bar.update(1)
@@ -3291,7 +3369,7 @@ class Table(Queryable):
             list_mode,
         )
         result = None
-        with self.db.conn:
+        with self.db.atomic():
             for query, params in queries_and_params:
                 try:
                     result = self.db.execute(query, params)
@@ -3528,6 +3606,7 @@ class Table(Queryable):
         self.last_rowid = None
         self.last_pk = None
         if truncate and self.exists():
+            with self.db.atomic():
                 self.db.execute("DELETE FROM {};".format(quote_identifier(self.name)))
         result = None
         for chunk in chunks(itertools.chain([first_record], records_iter), batch_size):

tests/test_atomic.py

@@ -0,0 +1,174 @@
+import pytest
+
+from sqlite_utils.db import _iter_complete_sql_statements
+from sqlite_utils.utils import sqlite3
+
+
+@pytest.mark.parametrize(
+    "sql,expected",
+    (
+        (
+            "CREATE TABLE t(id); INSERT INTO t VALUES (1)",
+            ["CREATE TABLE t(id);", "INSERT INTO t VALUES (1)"],
+        ),
+        (
+            "INSERT INTO t VALUES ('a;b');",
+            ["INSERT INTO t VALUES ('a;b');"],
+        ),
+        (
+            "-- comment;\nCREATE TABLE t(id);",
+            ["-- comment;\nCREATE TABLE t(id);"],
+        ),
+        (
+            """
+            CREATE TRIGGER t_ai AFTER INSERT ON t
+            BEGIN
+                UPDATE t SET value = 'a;b' WHERE id = new.id;
+                INSERT INTO log VALUES ('x;y');
+            END;
+            """,
+            [
+                "CREATE TRIGGER t_ai AFTER INSERT ON t\n"
+                "            BEGIN\n"
+                "                UPDATE t SET value = 'a;b' WHERE id = new.id;\n"
+                "                INSERT INTO log VALUES ('x;y');\n"
+                "            END;"
+            ],
+        ),
+    ),
+)
+def test_iter_complete_sql_statements(sql, expected):
+    assert list(_iter_complete_sql_statements(sql)) == expected
+
+
+def test_atomic_commits(fresh_db):
+    with fresh_db.atomic():
+        fresh_db["dogs"].insert({"id": 1, "name": "Cleo"}, pk="id")
+
+    assert list(fresh_db["dogs"].rows) == [{"id": 1, "name": "Cleo"}]
+
+
+def test_atomic_rolls_back(fresh_db):
+    with pytest.raises(RuntimeError):
+        with fresh_db.atomic():
+            fresh_db["dogs"].insert({"id": 1, "name": "Cleo"}, pk="id")
+            raise RuntimeError("boom")
+
+    assert not fresh_db["dogs"].exists()
+
+
+def test_nested_atomic_rolls_back_to_savepoint(fresh_db):
+    fresh_db["dogs"].create({"id": int, "name": str}, pk="id")
+
+    with fresh_db.atomic():
+        fresh_db["dogs"].insert({"id": 1, "name": "Cleo"})
+        with pytest.raises(RuntimeError):
+            with fresh_db.atomic():
+                fresh_db["dogs"].insert({"id": 2, "name": "Pancakes"})
+                raise RuntimeError("boom")
+        fresh_db["dogs"].insert({"id": 3, "name": "Marnie"})
+
+    assert list(fresh_db["dogs"].rows) == [
+        {"id": 1, "name": "Cleo"},
+        {"id": 3, "name": "Marnie"},
+    ]
+
+
+def test_outer_atomic_rolls_back_released_savepoint(fresh_db):
+    with pytest.raises(RuntimeError):
+        with fresh_db.atomic():
+            fresh_db["dogs"].insert({"id": 1, "name": "Cleo"}, pk="id")
+            with fresh_db.atomic():
+                fresh_db["dogs"].insert({"id": 2, "name": "Pancakes"})
+            raise RuntimeError("boom")
+
+    assert not fresh_db["dogs"].exists()
+
+
+def test_executescript_does_not_commit_open_atomic_block(fresh_db):
+    with pytest.raises(RuntimeError):
+        with fresh_db.atomic():
+            fresh_db.executescript("""
+                CREATE TABLE dogs(id INTEGER PRIMARY KEY, name TEXT);
+                CREATE TRIGGER dogs_ai AFTER INSERT ON dogs
+                BEGIN
+                    UPDATE dogs SET name = upper(new.name) || '; updated' WHERE id = new.id;
+                END;
+                -- This comment has a semicolon;
+                INSERT INTO dogs VALUES (1, 'Cleo; the first');
+            """)
+            raise RuntimeError("boom")
+
+    assert not fresh_db["dogs"].exists()
+
+
+def test_transform_does_not_commit_open_atomic_block(fresh_db):
+    fresh_db["dogs"].insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id")
+
+    with pytest.raises(RuntimeError):
+        with fresh_db.atomic():
+            fresh_db["dogs"].insert({"id": 2, "name": "Pancakes", "age": "6"})
+            fresh_db["dogs"].transform(rename={"age": "dog_age"})
+            raise RuntimeError("boom")
+
+    assert (
+        fresh_db["dogs"].schema
+        == 'CREATE TABLE "dogs" (\n   "id" INTEGER PRIMARY KEY,\n   "name" TEXT,\n   "age" TEXT\n)'
+    )
+    assert list(fresh_db["dogs"].rows) == [
+        {"id": 1, "name": "Cleo", "age": "5"},
+    ]
+
+
+def test_transform_parent_table_with_foreign_keys_in_atomic(fresh_db):
+    fresh_db.conn.execute("PRAGMA foreign_keys=ON")
+    fresh_db["authors"].insert({"id": 1, "name": "Tina"}, pk="id")
+    fresh_db["books"].insert(
+        {"id": 1, "title": "Book", "author_id": 1},
+        pk="id",
+        foreign_keys={"author_id"},
+    )
+
+    with fresh_db.atomic():
+        fresh_db["authors"].transform(rename={"name": "full_name"})
+        assert fresh_db.conn.execute("PRAGMA foreign_keys").fetchone()[0]
+
+    assert (
+        fresh_db["authors"].schema
+        == 'CREATE TABLE "authors" (\n   "id" INTEGER PRIMARY KEY,\n   "full_name" TEXT\n)'
+    )
+    assert fresh_db.execute("PRAGMA foreign_key_check").fetchall() == []
+
+
+def test_transform_parent_table_with_foreign_keys_rolls_back(fresh_db):
+    fresh_db.conn.execute("PRAGMA foreign_keys=ON")
+    fresh_db["authors"].insert({"id": 1, "name": "Tina"}, pk="id")
+    fresh_db["books"].insert(
+        {"id": 1, "title": "Book", "author_id": 1},
+        pk="id",
+        foreign_keys={"author_id"},
+    )
+
+    with pytest.raises(RuntimeError):
+        with fresh_db.atomic():
+            fresh_db["authors"].transform(rename={"name": "full_name"})
+            raise RuntimeError("boom")
+
+    assert (
+        fresh_db["authors"].schema
+        == 'CREATE TABLE "authors" (\n   "id" INTEGER PRIMARY KEY,\n   "name" TEXT\n)'
+    )
+    assert fresh_db.conn.execute("PRAGMA foreign_keys").fetchone()[0]
+    assert fresh_db.execute("PRAGMA foreign_key_check").fetchall() == []
+
+
+def test_transform_detects_foreign_key_check_violations(fresh_db):
+    fresh_db.conn.execute("PRAGMA foreign_keys=ON")
+    fresh_db["authors"].insert({"id": 1, "name": "Tina"}, pk="id")
+    fresh_db["books"].insert({"id": 1, "author_id": 2}, pk="id")
+
+    with pytest.raises(sqlite3.IntegrityError):
+        fresh_db["books"].transform(add_foreign_keys=(("author_id", "authors", "id"),))
+
+    assert fresh_db["books"].foreign_keys == []
+    assert fresh_db.conn.execute("PRAGMA foreign_keys").fetchone()[0]