diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 8dad4c4..0c0a250 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -28,6 +28,7 @@ from typing import ( cast, Any, Callable, + ContextManager, Dict, Generator, Iterable, @@ -346,6 +347,25 @@ class Database: "Close the SQLite connection, and the underlying database file" self.conn.close() + @contextlib.contextmanager + def ensure_autocommit_off(self): + """ + Ensure autocommit is off for this database connection. + + Example usage:: + + with db.ensure_autocommit_off(): + # do stuff here + + This will reset to the previous autocommit state at the end of the block. + """ + old_isolation_level = self.conn.isolation_level + try: + self.conn.isolation_level = None + yield + finally: + self.conn.isolation_level = old_isolation_level + @contextlib.contextmanager def tracer(self, tracer: Optional[Callable] = None): """ @@ -662,12 +682,14 @@ class Database: Sets ``journal_mode`` to ``'wal'`` to enable Write-Ahead Log mode. """ if self.journal_mode != "wal": - self.execute("PRAGMA journal_mode=wal;") + with self.ensure_autocommit_off(): + self.execute("PRAGMA journal_mode=wal;") def disable_wal(self): "Sets ``journal_mode`` back to ``'delete'`` to disable Write-Ahead Log mode." if self.journal_mode != "delete": - self.execute("PRAGMA journal_mode=delete;") + with self.ensure_autocommit_off(): + self.execute("PRAGMA journal_mode=delete;") def _ensure_counts_table(self): with self.conn: diff --git a/tests/test_cli.py b/tests/test_cli.py index 5360e56..eb569f0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1429,7 +1429,7 @@ def test_enable_wal(): db = Database(dbname) db["t"].create({"pk": int}, pk="pk") assert db.journal_mode == "delete" - result = runner.invoke(cli.cli, ["enable-wal"] + dbs) + result = runner.invoke(cli.cli, ["enable-wal"] + dbs, catch_exceptions=False) assert 0 == result.exit_code for dbname in dbs: db = Database(dbname)