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()`.

.github/workflows/test.yml

@@ -35,6 +35,9 @@ jobs:
     - name: Install SpatiaLite
       if: matrix.os == 'ubuntu-latest'
       run: sudo apt-get install libsqlite3-mod-spatialite
+    - name: On macOS with Python 3.10 test with sqlean.py
+      if: matrix.os == 'macos-latest' && matrix.python-version == '3.10'
+      run: pip install sqlean.py sqlite-dump
     - name: Build extension for --load-extension test
       if: matrix.os == 'ubuntu-latest'
       run: |-

docs/installation.rst

@@ -39,6 +39,32 @@ Using pipx
 
     pipx install sqlite-utils
 
+.. _installation_sqlite3_alternatives:
+
+Alternatives to sqlite3
+=======================
+
+By default, ``sqlite-utils`` uses the ``sqlite3`` package bundled with the Python standard library.
+
+Depending on your operating system, this may come with some limitations.
+
+On some platforms the ability to load additional extensions (via ``conn.load_extension(...)`` or ``--load-extension=/path/to/extension``) may be disabled.
+
+You may also see the error ``sqlite3.OperationalError: table sqlite_master may not be modified`` when trying to alter an existing table.
+
+You can work around these limitations by installing either the `pysqlite3 <https://pypi.org/project/pysqlite3/>`__ package or the `sqlean.py <https://pypi.org/project/sqlean.py/>`__ package, both of which provide drop-in replacements for the standard library ``sqlite3`` module but with a recent version of SQLite and full support for loading extensions.
+
+To install ``sqlean.py`` (which has compiled binary wheels available for all major platforms) run the following:
+
+.. code-block:: bash
+
+    sqlite-utils install sqlean.py
+
+``pysqlite3`` and ``sqlean.py`` do not provide implementations of the ``.iterdump()`` method. To use that method (see :ref:`python_api_itedump`) or the ``sqlite-utils dump`` command you should also install the ``sqlite-dump`` package:
+
+.. code-block:: bash
+
+    sqlite-utils install sqlite-dump
 
 .. _installation_completion:
 

docs/python-api.rst

@@ -1779,6 +1779,25 @@ The ``db.sqlite_version`` property returns a tuple of integers representing the
     >>> db.sqlite_version
     (3, 36, 0)
 
+.. _python_api_itedump:
+
+Dumping the database to SQL
+===========================
+
+The ``db.iterdump()`` method returns a sequence of SQL strings representing a complete dump of the database. Use it like this:
+
+.. code-block:: python
+
+    full_sql = "".join(db.iterdump())
+
+This uses the `sqlite3.Connection.iterdump() <https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.iterdump>`__ method.
+
+If you are using ``pysqlite3`` or ``sqlean.py`` the underlying method may be missing. If you install the `sqlite-dump <https://pypi.org/project/sqlite-dump/>`__ package then the ``db.iterdump()`` method will use that implementation instead:
+
+.. code-block:: bash
+
+    pip install sqlite-dump
+
 .. _python_api_introspection:
 
 Introspecting tables and views

mypy.ini

@@ -0,0 +1,4 @@
+[mypy]
+
+[mypy-pysqlite3,sqlean,sqlite_dump]
+ignore_missing_imports = True

sqlite_utils/cli.py

@@ -395,7 +395,7 @@ def dump(path, load_extension):
     """
     db = sqlite_utils.Database(path)
     _load_extensions(db, load_extension)
-    for line in db.conn.iterdump():
+    for line in db.iterdump():
         click.echo(line)
 
 
@@ -1893,7 +1893,7 @@ def memory(
         return
 
     if dump:
-        for line in db.conn.iterdump():
+        for line in db.iterdump():
             click.echo(line)
         return
 
@@ -1903,7 +1903,7 @@ def memory(
 
     if save:
         db2 = sqlite_utils.Database(save)
-        for line in db.conn.iterdump():
+        for line in db.iterdump():
             db2.execute(line)
         return
 

sqlite_utils/db.py

@@ -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.

sqlite_utils/utils.py

@@ -14,15 +14,22 @@ from typing import Dict, cast, BinaryIO, Iterable, Optional, Tuple, Type
 import click
 
 try:
-    import pysqlite3 as sqlite3  # type: ignore
-    import pysqlite3.dbapi2  # type: ignore
+    import pysqlite3 as sqlite3  # noqa: F401
+    from pysqlite3 import dbapi2  # noqa: F401
 
-    OperationalError = pysqlite3.dbapi2.OperationalError
+    OperationalError = dbapi2.OperationalError
 except ImportError:
-    # https://github.com/python/mypy/issues/1153#issuecomment-253842414
-    import sqlite3  # type: ignore
+    try:
+        import sqlean as sqlite3  # noqa: F401
+        from sqlean import dbapi2  # noqa: F401
+
+        OperationalError = dbapi2.OperationalError
+    except ImportError:
+        import sqlite3  # noqa: F401
+        from sqlite3 import dbapi2  # noqa: F401
+
+        OperationalError = dbapi2.OperationalError
 
-    OperationalError = sqlite3.OperationalError
 
 SPATIALITE_PATHS = (
     "/usr/lib/x86_64-linux-gnu/mod_spatialite.so",

tests/test_analyze.py

@@ -14,7 +14,9 @@ def db(fresh_db):
 def test_analyze_whole_database(db):
     assert set(db.table_names()) == {"one_index", "two_indexes"}
     db.analyze()
-    assert set(db.table_names()) == {"one_index", "two_indexes", "sqlite_stat1"}
+    assert set(db.table_names()).issuperset(
+        {"one_index", "two_indexes", "sqlite_stat1"}
+    )
     assert list(db["sqlite_stat1"].rows) == [
         {"tbl": "two_indexes", "idx": "idx_two_indexes_species", "stat": "1 1"},
         {"tbl": "two_indexes", "idx": "idx_two_indexes_name", "stat": "1 1"},
@@ -24,13 +26,15 @@ def test_analyze_whole_database(db):
 
 @pytest.mark.parametrize("method", ("db_method_with_name", "table_method"))
 def test_analyze_one_table(db, method):
-    assert set(db.table_names()) == {"one_index", "two_indexes"}
+    assert set(db.table_names()).issuperset({"one_index", "two_indexes"})
     if method == "db_method_with_name":
         db.analyze("one_index")
     elif method == "table_method":
         db["one_index"].analyze()
 
-    assert set(db.table_names()) == {"one_index", "two_indexes", "sqlite_stat1"}
+    assert set(db.table_names()).issuperset(
+        {"one_index", "two_indexes", "sqlite_stat1"}
+    )
     assert list(db["sqlite_stat1"].rows) == [
         {"tbl": "one_index", "idx": "idx_one_index_name", "stat": "1 1"}
     ]
@@ -39,7 +43,9 @@ def test_analyze_one_table(db, method):
 def test_analyze_index_by_name(db):
     assert set(db.table_names()) == {"one_index", "two_indexes"}
     db.analyze("idx_two_indexes_species")
-    assert set(db.table_names()) == {"one_index", "two_indexes", "sqlite_stat1"}
+    assert set(db.table_names()).issuperset(
+        {"one_index", "two_indexes", "sqlite_stat1"}
+    )
     assert list(db["sqlite_stat1"].rows) == [
         {"tbl": "two_indexes", "idx": "idx_two_indexes_species", "stat": "1 1"},
     ]

tests/test_analyze_tables.py

@@ -135,7 +135,8 @@ def test_analyze_column(db_to_analyze, column, extra_kwargs, expected):
 def db_to_analyze_path(db_to_analyze, tmpdir):
     path = str(tmpdir / "test.db")
     db = sqlite3.connect(path)
-    db.executescript("\n".join(db_to_analyze.conn.iterdump()))
+    sql = "\n".join(db_to_analyze.iterdump())
+    db.executescript(sql)
     return path
 
 

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)

tests/test_cli_memory.py

@@ -164,9 +164,9 @@ def test_memory_dump(extra_args):
         input="id,name\n1,Cleo\n2,Bants",
     )
     assert result.exit_code == 0
-    assert result.output.strip() == (
+    expected = (
         "BEGIN TRANSACTION;\n"
-        'CREATE TABLE "stdin" (\n'
+        'CREATE TABLE IF NOT EXISTS "stdin" (\n'
         "   [id] INTEGER,\n"
         "   [name] TEXT\n"
         ");\n"
@@ -176,6 +176,9 @@ def test_memory_dump(extra_args):
         "CREATE VIEW t AS select * from [stdin];\n"
         "COMMIT;"
     )
+    # Using sqlite-dump it won't have IF NOT EXISTS
+    expected_alternative = expected.replace("IF NOT EXISTS ", "")
+    assert result.output.strip() in (expected, expected_alternative)
 
 
 @pytest.mark.parametrize("extra_args", ([], ["select 1"]))

tests/test_gis.py

@@ -6,6 +6,12 @@ from sqlite_utils.cli import cli
 from sqlite_utils.db import Database
 from sqlite_utils.utils import find_spatialite, sqlite3
 
+try:
+    import sqlean
+except ImportError:
+    sqlean = None
+
+
 pytestmark = [
     pytest.mark.skipif(
         not find_spatialite(), reason="Could not find SpatiaLite extension"
@@ -14,6 +20,9 @@ pytestmark = [
         not hasattr(sqlite3.Connection, "enable_load_extension"),
         reason="sqlite3.Connection missing enable_load_extension",
     ),
+    pytest.mark.skipif(
+        sqlean is not None, reason="sqlean.py is not compatible with SpatiaLite"
+    ),
 ]