Split load_spatialite into load_spatialite_driver and init_spatialite (#459)

adrien-berchet · web-flow · commit 3dd0582ba577 · 2023-07-05T11:11:06.000+02:00
* Split load_spatialite into load_spatialite_driver and init_spatialite

* Add transaction and journal_mode arguments and improve docs

* Skip PyPy when journal_mode='OFF'
diff --git a/geoalchemy2/admin/dialects/geopackage.py b/geoalchemy2/admin/dialects/geopackage.py
@@ -2,7 +2,6 @@
 
 See GeoPackage specifications here: http://www.geopackage.org/spec/
 """
-import os
 import re
 
 from sqlalchemy import text
@@ -19,6 +18,7 @@
 from geoalchemy2.admin.dialects.common import setup_create_drop
 from geoalchemy2.admin.dialects.sqlite import _SQLITE_FUNCTIONS
 from geoalchemy2.admin.dialects.sqlite import get_col_dim
+from geoalchemy2.admin.dialects.sqlite import load_spatialite_driver
 from geoalchemy2.types import Geography
 from geoalchemy2.types import Geometry
 from geoalchemy2.types import _DummyGeometry
@@ -37,31 +37,48 @@ class GeoPackageDialect(SQLiteDialect_pysqlite):
 registry.register("gpkg", "geoalchemy2.admin.dialects.geopackage", "GeoPackageDialect")
 
 
-def load_spatialite_gpkg(dbapi_conn, connection_record):
-    """Load SpatiaLite extension in GeoPackage.
+def load_geopackage_driver(dbapi_conn, *args):
+    """Load SpatiaLite extension in GeoPackage connection and set VirtualGpkg and Amphibious modes.
 
-    The path to the SpatiaLite module should be set in the `SPATIALITE_LIBRARY_PATH` environment
-    variable.
+    .. Warning::
+        The path to the SpatiaLite module should be set in the `SPATIALITE_LIBRARY_PATH`
+        environment variable.
 
-    .. Note::
+    Args:
+        dbapi_conn: The DBAPI connection.
+    """
+    load_spatialite_driver(dbapi_conn, *args)
+
+    dbapi_conn.execute("SELECT AutoGpkgStart();")
+    dbapi_conn.execute("SELECT EnableGpkgAmphibiousMode();")
+
+
+def init_geopackage(dbapi_conn, *args):
+    """Initialize GeoPackage tables.
 
+    Args:
+        dbapi_conn: The DBAPI connection.
+
+    .. Warning::
         No EPSG SRID is loaded in the `gpkg_spatial_ref_sys` table after initialization but
         it is possible to load other EPSG SRIDs afterwards using the
         `gpkgInsertEpsgSRID(srid)`.
         Nevertheless, SRIDs of newly created tables are automatically added.
     """
-    if "SPATIALITE_LIBRARY_PATH" not in os.environ:
-        raise RuntimeError("The SPATIALITE_LIBRARY_PATH environment variable is not set.")
-    dbapi_conn.enable_load_extension(True)
-    dbapi_conn.load_extension(os.environ["SPATIALITE_LIBRARY_PATH"])
-    dbapi_conn.enable_load_extension(False)
-
     if not dbapi_conn.execute("SELECT CheckGeoPackageMetaData();").fetchone()[0]:
         # This only works on the main database
         dbapi_conn.execute("SELECT gpkgCreateBaseTables();")
 
-    dbapi_conn.execute("SELECT AutoGpkgStart();")
-    dbapi_conn.execute("SELECT EnableGpkgAmphibiousMode();")
+
+def load_spatialite_gpkg(*args, **kwargs):
+    """Load SpatiaLite extension in GeoPackage and initialize internal tables.
+
+    See :func:`geoalchemy2.admin.dialects.geopackage.load_geopackage_driver` and
+    :func:`geoalchemy2.admin.dialects.geopackage.init_geopackage` functions for details about
+    arguments.
+    """
+    load_geopackage_driver(*args)
+    init_geopackage(*args, **kwargs)
 
 
 def _get_spatialite_attrs(bind, table_name, col_name):
diff --git a/geoalchemy2/admin/dialects/sqlite.py b/geoalchemy2/admin/dialects/sqlite.py
@@ -1,5 +1,6 @@
 """This module defines specific functions for SQLite dialect."""
 import os
+from typing import Optional
 
 from sqlalchemy import text
 from sqlalchemy.ext.compiler import compiles
@@ -16,36 +17,119 @@
 from geoalchemy2.types import _DummyGeometry
 
 
-def load_spatialite(dbapi_conn, connection_record, init_mode=None):
-    """Load SpatiaLite extension in SQLite DB.
+def load_spatialite_driver(dbapi_conn, *args):
+    """Load SpatiaLite extension in SQLite connection.
 
-    The path to the SpatiaLite module should be set in the `SPATIALITE_LIBRARY_PATH` environment
-    variable.
+    .. Warning::
+        The path to the SpatiaLite module should be set in the `SPATIALITE_LIBRARY_PATH`
+        environment variable.
 
-    The init_mode argument can be `'NONE'` to load all EPSG SRIDs, `'WGS84'` to load only the ones
-    related to WGS84 or `'EMPTY'` to not load any EPSG SRID.
-
-    .. Note::
-
-        It is possible to load other EPSG SRIDs afterwards using the `InsertEpsgSrid(srid)`.
+    Args:
+        dbapi_conn: The DBAPI connection.
     """
     if "SPATIALITE_LIBRARY_PATH" not in os.environ:
         raise RuntimeError("The SPATIALITE_LIBRARY_PATH environment variable is not set.")
     dbapi_conn.enable_load_extension(True)
     dbapi_conn.load_extension(os.environ["SPATIALITE_LIBRARY_PATH"])
     dbapi_conn.enable_load_extension(False)
 
-    init_mode_values = [None, "WGS84", "EMPTY"]
+
+def init_spatialite(
+    dbapi_conn,
+    *args,
+    transaction: bool = False,
+    init_mode: Optional[str] = None,
+    journal_mode: Optional[str] = None,
+):
+    """Initialize internal SpatiaLite tables.
+
+    Args:
+        dbapi_conn: The DBAPI connection.
+        init_mode: Can be `'NONE'` to load all EPSG SRIDs, `'WGS84'` to load only the ones related
+            to WGS84 or `'EMPTY'` to not load any EPSG SRID.
+
+            .. Note::
+
+                It is possible to load other EPSG SRIDs afterwards using `InsertEpsgSrid(srid)`.
+
+        transaction: If set to `True` the whole operation will be handled as a single Transaction
+            (faster). The default value is `False` (slower, but safer).
+        journal_mode: Change the journal mode to the given value. This can make the table creation
+            much faster. The possible values are the following: 'DELETE', 'TRUNCATE', 'PERSIST',
+            'MEMORY', 'WAL' and 'OFF'. See https://www.sqlite.org/pragma.html#pragma_journal_mode
+            for more details.
+
+            .. Warning::
+                Some values, like 'MEMORY' or 'OFF', can lead to corrupted databases if the process
+                is interrupted during initialization.
+
+            .. Note::
+                The original value is restored after the initialization.
+
+    .. Note::
+        When using this function as a listener it is not possible to pass the `transaction`,
+        `init_mode` or `journal_mode` arguments directly. To do this you can either create another
+        function that calls `load_spatialite` with an hard-coded `init_mode` or just use a lambda::
+
+            >>> sqlalchemy.event.listen(
+            ...     engine,
+            ...     "connect",
+            ...     lambda x, y: load_spatialite(
+            ...         x,
+            ...         y,
+            ...         transaction=True,
+            ...         init_mode="EMPTY",
+            ...         journal_mode="OFF",
+            ...     )
+            ... )
+    """
+    func_args = []
+
+    # Check the value of the 'transaction' parameter
+    if not isinstance(transaction, (bool, int)):
+        raise ValueError("The 'transaction' argument must be True or False.")
+    else:
+        func_args.append(str(transaction))
+
+    # Check the value of the 'init_mode' parameter
+    init_mode_values = ["WGS84", "EMPTY"]
     if isinstance(init_mode, str):
         init_mode = init_mode.upper()
-    if init_mode not in init_mode_values:
-        raise ValueError("The 'init_mode' must be in {}".format(init_mode_values))
+    if init_mode is not None:
+        if init_mode not in init_mode_values:
+            raise ValueError("The 'init_mode' argument must be one of {}.".format(init_mode_values))
+        func_args.append(f"'{init_mode}'")
+
+    # Check the value of the 'journal_mode' parameter
+    journal_mode_values = ["DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"]
+    if isinstance(journal_mode, str):
+        journal_mode = journal_mode.upper()
+    if journal_mode is not None:
+        if journal_mode not in journal_mode_values:
+            raise ValueError(
+                "The 'journal_mode' argument must be one of {}.".format(journal_mode_values)
+            )
 
     if dbapi_conn.execute("SELECT CheckSpatialMetaData();").fetchone()[0] < 1:
-        if init_mode is not None:
-            dbapi_conn.execute("SELECT InitSpatialMetaData('{}');".format(init_mode))
-        else:
-            dbapi_conn.execute("SELECT InitSpatialMetaData();")
+        if journal_mode is not None:
+            current_journal_mode = dbapi_conn.execute("PRAGMA journal_mode").fetchone()[0]
+            dbapi_conn.execute("PRAGMA journal_mode = {}".format(journal_mode))
+
+        dbapi_conn.execute("SELECT InitSpatialMetaData({});".format(", ".join(func_args)))
+
+        if journal_mode is not None:
+            dbapi_conn.execute("PRAGMA journal_mode = {}".format(current_journal_mode))
+
+
+def load_spatialite(*args, **kwargs):
+    """Load SpatiaLite extension in SQLite DB and initialize internal tables.
+
+    See :func:`geoalchemy2.admin.dialects.sqlite.load_spatialite_driver` and
+    :func:`geoalchemy2.admin.dialects.sqlite.init_spatialite` functions for details about
+    arguments.
+    """
+    load_spatialite_driver(*args)
+    init_spatialite(*args, **kwargs)
 
 
 def _get_spatialite_attrs(bind, table_name, col_name):
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -1,4 +1,5 @@
 import os
+import platform
 import re
 import shutil
 import sys
@@ -17,8 +18,6 @@
 from geoalchemy2 import load_spatialite
 from geoalchemy2 import load_spatialite_gpkg
 
-# from geoalchemy2 import load_spatialite_gpkg
-
 
 class test_only_with_dialects:
     def __init__(self, *dialects):
@@ -76,6 +75,11 @@ def skip_pg12_sa1217(bind):
         pytest.skip("Reflection for PostgreSQL-12 is only supported by sqlalchemy>=1.2.17")
 
 
+def skip_pypy(msg=None):
+    if platform.python_implementation() == "PyPy":
+        pytest.skip(msg if msg is not None else "Incompatible with PyPy")
+
+
 def select(args):
     if version.parse(SA_VERSION) < version.parse("1.4"):
         return raw_select(args)
@@ -101,15 +105,6 @@ def copy_and_connect_sqlite_db(input_db, tmp_db, engine_echo, dialect):
         db_url, echo=engine_echo, execution_options={"schema_translate_map": {"gis": None}}
     )
 
-    if input_db.lower().endswith("spatialite_lt_4.sqlite"):
-        engine._spatialite_version = 3
-    elif input_db.lower().endswith("spatialite_ge_4.sqlite"):
-        engine._spatialite_version = 4
-    elif input_db.lower().endswith(".gpkg"):
-        engine._spatialite_version = -1
-    else:
-        engine._spatialite_version = None
-
     if dialect == "gpkg":
         listen(engine, "connect", load_spatialite_gpkg)
     else:
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -172,7 +172,6 @@ def engine(tmpdir, db_url, _engine_echo):
     # For other dialects the engine is directly returned
     engine = create_engine(db_url, echo=_engine_echo)
     engine.update_execution_options(search_path=["gis", "public"])
-    engine._spatialite_version = None
     return engine
 
 
diff --git a/tests/test_functional_sqlite.py b/tests/test_functional_sqlite.py
@@ -24,6 +24,7 @@
 
 from . import select
 from . import skip_case_insensitivity
+from . import skip_pypy
 from . import test_only_with_dialects
 from .schema_fixtures import TransformedGeometry
 
@@ -207,8 +208,31 @@ def test_all_indexes(self, conn, TableWithIndexes, setup_tables):
 
 class TestMiscellaneous:
     @test_only_with_dialects("sqlite-spatialite3", "sqlite-spatialite4")
-    @pytest.mark.parametrize("init_mode", [None, "WGS84", "EMPTY"])
-    def test_load_spatialite(self, tmpdir, _engine_echo, init_mode, check_spatialite):
+    @pytest.mark.parametrize(
+        [
+            "transaction",
+            "init_mode",
+            "journal_mode",
+        ],
+        [
+            pytest.param(False, "WGS84", None),
+            pytest.param(False, "WGS84", "OFF"),
+            pytest.param(False, "EMPTY", None),
+            pytest.param(False, "EMPTY", "OFF"),
+            pytest.param(True, None, None),
+            pytest.param(True, None, "OFF"),
+            pytest.param(True, "WGS84", None),
+            pytest.param(True, "WGS84", "OFF"),
+            pytest.param(True, "EMPTY", None),
+            pytest.param(True, "EMPTY", "OFF"),
+        ],
+    )
+    def test_load_spatialite(
+        self, tmpdir, _engine_echo, check_spatialite, transaction, init_mode, journal_mode
+    ):
+        if journal_mode == "OFF":
+            skip_pypy("The journal mode can not be OFF with PyPy.")
+
         # Create empty DB
         tmp_db = tmpdir / "test_spatial_db.sqlite"
         db_url = f"sqlite:///{tmp_db}"
@@ -219,13 +243,21 @@ def test_load_spatialite(self, tmpdir, _engine_echo, init_mode, check_spatialite
 
         assert not conn.execute(text("PRAGMA main.table_info('geometry_columns')")).fetchall()
         assert not conn.execute(text("PRAGMA main.table_info('spatial_ref_sys')")).fetchall()
+        assert conn.execute(text("PRAGMA journal_mode")).fetchone()[0].upper() == "DELETE"
 
-        load_spatialite(conn.connection.dbapi_connection, None, init_mode)
+        load_spatialite(
+            conn.connection.dbapi_connection,
+            transaction=transaction,
+            init_mode=init_mode,
+            journal_mode=journal_mode,
+        )
 
         assert conn.execute(text("SELECT CheckSpatialMetaData();")).scalar() == 3
         assert conn.execute(text("PRAGMA main.table_info('geometry_columns')")).fetchall()
         assert conn.execute(text("PRAGMA main.table_info('spatial_ref_sys')")).fetchall()
 
+        assert conn.execute(text("PRAGMA journal_mode")).fetchone()[0].upper() == "DELETE"
+
         # Check that spatial_ref_sys table was properly populated
         nb_srid = conn.execute(text("""SELECT COUNT(*) FROM spatial_ref_sys;""")).scalar()
         if init_mode is None:
@@ -235,16 +267,47 @@ def test_load_spatialite(self, tmpdir, _engine_echo, init_mode, check_spatialite
         elif init_mode == "EMPTY":
             assert nb_srid == 0
 
+        # Check that the journal mode is properly reset even when an error is returned by the
+        # InitSpatialMetaData() function
+        assert conn.execute(text("PRAGMA journal_mode")).fetchone()[0].upper() == "DELETE"
+
+        load_spatialite(
+            conn.connection.dbapi_connection,
+            transaction=transaction,
+            init_mode=init_mode,
+            journal_mode=journal_mode,
+        )
+
+        assert conn.execute(text("PRAGMA journal_mode")).fetchone()[0].upper() == "DELETE"
+
+    @test_only_with_dialects("sqlite-spatialite3", "sqlite-spatialite4")
+    def test_load_spatialite_unknown_transaction(self, conn):
+        with pytest.raises(ValueError, match=r"The 'transaction' argument must be True or False\."):
+            load_spatialite(conn.connection.dbapi_connection, transaction="UNKNOWN MODE")
+
     @test_only_with_dialects("sqlite-spatialite3", "sqlite-spatialite4")
-    def test_load_spatialite_unknown_init_type(self, monkeypatch, conn):
-        with pytest.raises(ValueError):
-            load_spatialite(conn.connection.dbapi_connection, None, "UNKNOWN TYPE")
+    def test_load_spatialite_unknown_init_type(self, conn):
+        with pytest.raises(
+            ValueError, match=r"The 'init_mode' argument must be one of \['WGS84', 'EMPTY'\]\."
+        ):
+            load_spatialite(conn.connection.dbapi_connection, init_mode="UNKNOWN TYPE")
+
+    @test_only_with_dialects("sqlite-spatialite3", "sqlite-spatialite4")
+    def test_load_spatialite_unknown_journal_mode(self, conn):
+        with pytest.raises(
+            ValueError,
+            match=(
+                r"The 'journal_mode' argument must be one of "
+                r"\['DELETE', 'TRUNCATE', 'PERSIST', 'MEMORY', 'WAL', 'OFF'\]\."
+            ),
+        ):
+            load_spatialite(conn.connection.dbapi_connection, journal_mode="UNKNOWN MODE")
 
     @test_only_with_dialects("sqlite-spatialite3", "sqlite-spatialite4")
     def test_load_spatialite_no_env_variable(self, monkeypatch, conn):
         monkeypatch.delenv("SPATIALITE_LIBRARY_PATH")
         with pytest.raises(RuntimeError):
-            load_spatialite(conn.connection.dbapi_connection, None)
+            load_spatialite(conn.connection.dbapi_connection)
 
 
 class TestInsertionORM:
