New spatialite helper methods, closes #79

- db.init_spatialite()
- table.add_geometry_column()
- table.create_spatial_index()

Co-authored-by: Simon Willison <swillison@gmail.com>

sqlite_utils/cli.py

@@ -29,6 +29,7 @@ from .utils import (
     TypeTracker,
 )
 
+
 CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
 
 VALID_COLUMN_TYPES = ("INTEGER", "TEXT", "FLOAT", "BLOB")

sqlite_utils/db.py

@@ -6,6 +6,7 @@ from .utils import (
     types_for_column_types,
     column_affinity,
     progressbar,
+    find_spatialite,
 )
 from collections import namedtuple
 from collections.abc import Mapping
@@ -931,6 +932,46 @@ class Database:
             sql += " [{}]".format(name)
         self.execute(sql)
 
+    def init_spatialite(self, path: str = None) -> bool:
+        """
+        The ``init_spatialite`` method will load and initialize the SpatiaLite extension.
+        The ``path`` argument should be an absolute path to the compiled extension, which
+        can be found using ``find_spatialite``.
+
+        Returns true if SpatiaLite was successfully initialized.
+
+        .. code-block:: python
+
+            from sqlite_utils.db import Database
+            from sqlite_utils.utils import find_spatialite
+
+            db = Database("mydb.db")
+            db.init_spatialite(find_spatialite())
+
+        If you've installed SpatiaLite somewhere unexpected (for testing an alternate version, for example)
+        you can pass in an absolute path:
+
+        .. code-block:: python
+
+            from sqlite_utils.db import Database
+            from sqlite_utils.utils import find_spatialite
+
+            db = Database("mydb.db")
+            db.init_spatialite("./local/mod_spatialite.dylib")
+
+        """
+        if path is None:
+            path = find_spatialite()
+
+        self.conn.enable_load_extension(True)
+        self.conn.load_extension(path)
+        # Initialize SpatiaLite if not yet initialized
+        if "spatial_ref_sys" in self.table_names():
+            return False
+        cursor = self.execute("select InitSpatialMetadata(1)")
+        result = cursor.fetchone()
+        return result and bool(result[0])
+
 
 class Queryable:
     def exists(self) -> bool:
@@ -3012,6 +3053,84 @@ class Table(Queryable):
             least_common,
         )
 
+    def add_geometry_column(
+        self,
+        column_name: str,
+        geometry_type: str,
+        srid: int = 4326,
+        coord_dimension: str = "XY",
+        not_null: bool = False,
+    ) -> bool:
+        """
+        In SpatiaLite, a geometry column can only be added to an existing table.
+        To do so, use ``table.add_geometry_column``, passing in a geometry type.
+
+        By default, this will add a nullable column using
+        `SRID 4326 <https://spatialreference.org/ref/epsg/wgs-84/>`__. This can
+        be customized using the ``column_name``, ``srid`` and ``not_null`` arguments.
+
+        Returns True if the column was successfully added, False if not.
+
+        .. code-block:: python
+
+            from sqlite_utils.db import Database
+            from sqlite_utils.utils import find_spatialite
+
+            db = Database("mydb.db")
+            db.init_spatialite(find_spatialite())
+
+            # the table must exist before adding a geometry column
+            table = db["locations"].create({"name": str})
+            table.add_geometry_column("geometry", "POINT")
+
+        """
+        cursor = self.db.execute(
+            "SELECT AddGeometryColumn(?, ?, ?, ?, ?, ?);",
+            [
+                self.name,
+                column_name,
+                srid,
+                geometry_type,
+                coord_dimension,
+                int(not_null),
+            ],
+        )
+
+        result = cursor.fetchone()
+        return result and bool(result[0])
+
+    def create_spatial_index(self, column_name) -> bool:
+        """
+        A spatial index allows for significantly faster bounding box queries.
+        To create one, use ``create_spatial_index`` with the name of an existing geometry column.
+
+        Returns ``True`` if the index was successfully created, ``False`` if not. Calling this
+        function if an index already exists is a no-op.
+
+        .. code-block:: python
+
+            # assuming SpatiaLite is loaded, create the table, add the column
+            table = db["locations"].create({"name": str})
+            table.add_geometry_column("geometry", "POINT")
+
+            # now we can index it
+            table.create_spatial_index("geometry")
+
+            # the spatial index is a virtual table, which we can inspect
+            print(db["idx_locations_geometry"].schema)
+            # outputs:
+            # CREATE VIRTUAL TABLE "idx_locations_geometry" USING rtree(pkid, xmin, xmax, ymin, ymax)
+
+        """
+        if f"idx_{self.name}_{column_name}" in self.db.table_names():
+            return False
+
+        cursor = self.db.execute(
+            "select CreateSpatialIndex(?, ?)", [self.name, column_name]
+        )
+        result = cursor.fetchone()
+        return result and bool(result[0])
+
 
 class View(Queryable):
     def exists(self):

sqlite_utils/utils.py

@@ -22,12 +22,40 @@ except ImportError:
 
     OperationalError = sqlite3.OperationalError
 
+
 SPATIALITE_PATHS = (
     "/usr/lib/x86_64-linux-gnu/mod_spatialite.so",
     "/usr/local/lib/mod_spatialite.dylib",
 )
 
 
+def find_spatialite() -> str:
+    """
+    The ``find_spatialite()`` function searches for the `SpatiaLite <https://www.gaia-gis.it/fossil/libspatialite/index>`__ SQLite extension in some common places. It returns a string path to the location, or ``None`` if SpatiaLite was not found.
+
+    You can use it in code like this:
+
+    .. code-block:: python
+
+        from sqlite_utils import Database
+        from sqlite_utils.utils import find_spatialite
+
+        db = Database("mydb.db")
+        spatialite = find_spatialite()
+        if spatialite:
+            db.conn.enable_load_extension(True)
+            db.conn.load_extension(spatialite)
+
+        # or use with db.init_spatialite like this
+        db.init_spatialite(find_spatialite())
+
+    """
+    for path in SPATIALITE_PATHS:
+        if os.path.exists(path):
+            return path
+    return None
+
+
 def suggest_column_types(records):
     all_column_types = {}
     for record in records:
@@ -96,13 +124,6 @@ def decode_base64_values(doc):
     return dict(doc, **{k: base64.b64decode(doc[k]["encoded"]) for k in to_fix})
 
 
-def find_spatialite():
-    for path in SPATIALITE_PATHS:
-        if os.path.exists(path):
-            return path
-    return None
-
-
 class UpdateWrapper:
     def __init__(self, wrapped, update):
         self._wrapped = wrapped