Use sqlean if available in environment (#560)

Closes #559
Closes #235

Refs https://github.com/simonw/llm/issues/60

- Uses `sqlean` in place of `sqlite3` if `sqlean.py` is installed
- Uses `sqlite-dump` if available and `conn.iterdump()` does not exist
- New `with db.ensure_autocommit_off()` method for ensuring autocommit is off, used by `enable_wal()` and `disable_wal()`.
This commit is contained in:
Simon Willison 2023-06-25 16:25:51 -07:00 committed by GitHub
commit f5c63088e1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 136 additions and 19 deletions

View file

@ -38,6 +38,12 @@ from typing import (
)
import uuid
try:
from sqlite_dump import iterdump
except ImportError:
iterdump = None
SQLITE_MAX_VARS = 999
_quote_fts_re = re.compile(r'\s+|(".*?")')
@ -340,6 +346,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):
"""
@ -656,12 +681,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:
@ -1149,6 +1176,18 @@ class Database:
sql += " [{}]".format(name)
self.execute(sql)
def iterdump(self) -> Generator[str, None, None]:
"A sequence of strings representing a SQL dump of the database"
if iterdump:
yield from iterdump(self.conn)
else:
try:
yield from self.conn.iterdump()
except AttributeError:
raise AttributeError(
"conn.iterdump() not found - try pip install sqlite-dump"
)
def init_spatialite(self, path: Optional[str] = None) -> bool:
"""
The ``init_spatialite`` method will load and initialize the SpatiaLite extension.