WIP Testing geopackage support

.gitignore

@@ -18,6 +18,7 @@ ansible/*.retry
 .venv
 bin/
 lib/
+var/
 pyvenv.cfg
 
 # Unit test / coverage reports

setup.py

@@ -17,7 +17,7 @@ def get_version():
 
 
 install_requires = [
-    "GeoAlchemy2>=0.9,!=0.11.*",
+    "GeoAlchemy2>=0.14",
     "SQLAlchemy>=1.4",
     "alembic>=1.8,<2",
 ]

threedi_schema/application/schema.py

@@ -1,5 +1,6 @@
 import warnings
 
+import geoalchemy2.alembic_helpers  # NOQA
 from alembic import command as alembic_command
 from alembic.config import Config
 from alembic.environment import EnvironmentContext

threedi_schema/application/threedi_database.py

@@ -1,11 +1,13 @@
+import os
 import shutil
 import tempfile
 import uuid
 from contextlib import contextmanager
 from pathlib import Path
 
-from sqlalchemy import create_engine, event, inspect, text
-from sqlalchemy.engine import Engine
+from geoalchemy2.admin.dialects.geopackage import load_geopackage_driver
+from geoalchemy2.admin.dialects.sqlite import load_spatialite_driver
+from sqlalchemy import create_engine, inspect, text
 from sqlalchemy.event import listen
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.pool import NullPool
@@ -15,7 +17,6 @@
 __all__ = ["ThreediDatabase"]
 
 
-@event.listens_for(Engine, "connect")
 def set_sqlite_pragma(dbapi_connection, connection_record):
     """Switch on legacy_alter_table setting to fix our migrations.
 
@@ -35,40 +36,6 @@ def set_sqlite_pragma(dbapi_connection, connection_record):
     cursor.close()
 
 
-def load_spatialite(con, connection_record):
-    """Load spatialite extension as described in
-    https://geoalchemy-2.readthedocs.io/en/latest/spatialite_tutorial.html"""
-    import sqlite3
-
-    con.enable_load_extension(True)
-    cur = con.cursor()
-    libs = [
-        # SpatiaLite >= 4.2 and Sqlite >= 3.7.17, should work on all platforms
-        ("mod_spatialite", "sqlite3_modspatialite_init"),
-        # SpatiaLite >= 4.2 and Sqlite < 3.7.17 (Travis)
-        ("mod_spatialite.so", "sqlite3_modspatialite_init"),
-        # SpatiaLite < 4.2 (linux)
-        ("libspatialite.so", "sqlite3_extension_init"),
-    ]
-    found = False
-    for lib, entry_point in libs:
-        try:
-            cur.execute("select load_extension('{}', '{}')".format(lib, entry_point))
-        except sqlite3.OperationalError:
-            continue
-        else:
-            found = True
-            break
-    try:
-        cur.execute("select EnableGpkgAmphibiousMode()")
-    except sqlite3.OperationalError:
-        pass
-    if not found:
-        raise RuntimeError("Cannot find any suitable spatialite module")
-    cur.close()
-    con.enable_load_extension(False)
-
-
 class ThreediDatabase:
     def __init__(self, path, echo=False):
         self.path = path
@@ -89,6 +56,10 @@ def engine(self):
     def base_path(self):
         return Path(self.path).absolute().parent
 
+    @property
+    def _driver(self) -> str:
+        return "gpkg" if Path(self.path).suffix.lower() == ".gpkg" else "sqlite"
+
     def get_engine(self, get_seperate_engine=False):
         if self._engine is None or get_seperate_engine:
             if self.path == "":
@@ -98,9 +69,17 @@ def get_engine(self, get_seperate_engine=False):
             else:
                 poolclass = NullPool
             engine = create_engine(
-                "sqlite:///{0}".format(self.path), echo=self.echo, poolclass=poolclass
+                f"{self._driver}:///{self.path}", echo=self.echo, poolclass=poolclass
+            )
+            os.environ.setdefault("SPATIALITE_LIBRARY_PATH", "mod_spatialite")
+            listen(
+                engine,
+                "connect",
+                load_spatialite_driver
+                if self._driver == "sqlite"
+                else load_geopackage_driver,
             )
-            listen(engine, "connect", load_spatialite)
+            listen(engine, "connect", set_sqlite_pragma)
             if get_seperate_engine:
                 return engine
             else:
@@ -139,14 +118,16 @@ def file_transaction(self, start_empty=False):
         On contextmanager exit, the database is copied back and the real
         database is overwritten. On error, nothing happens.
         """
+        if self._engine is not None:
+            self._engine.dispose()
         with tempfile.TemporaryDirectory() as tempdir:
-            work_file = Path(tempdir) / f"work-{uuid.uuid4()}.sqlite"
+            work_file = Path(tempdir) / f"work-{uuid.uuid4()}{Path(self.path).suffix}"
             # copy the database to the temporary directory
             if not start_empty:
                 shutil.copy(self.path, str(work_file))
             # yield a new ThreediDatabase refering to the backup
             try:
-                yield self.__class__(str(work_file))
+                yield self.__class__(str(work_file), echo=self.echo)
             except Exception as e:
                 raise e
             else:

threedi_schema/domain/custom_types.py

@@ -1,6 +1,5 @@
 import geoalchemy2
-from packaging import version
-from sqlalchemy.types import Integer, TypeDecorator, VARCHAR
+from sqlalchemy.types import Integer, TEXT, TypeDecorator
 
 
 class Geometry(geoalchemy2.types.Geometry):
@@ -13,8 +12,6 @@ def __init__(self, geometry_type, from_text="ST_GeomFromEWKT"):
             "spatial_index": True,
             "from_text": from_text,
         }
-        if version.parse(geoalchemy2.__version__) < version.parse("0.13.0"):
-            kwargs["management"] = True
         super().__init__(**kwargs)
 
 
@@ -68,4 +65,4 @@ class IntegerEnum(CustomEnum):
 
 class VarcharEnum(CustomEnum):
     cache_ok = True
-    impl = VARCHAR
+    impl = TEXT