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>

docs/python-api.rst

@@ -352,7 +352,7 @@ Counting rows
 
 To count the number of rows that would be returned by a where filter, use ``.count_where(where, where_args)``:
 
-    >>> db["dogs"].count_where("age > ?", [1]):
+    >>> db["dogs"].count_where("age > ?", [1])
     2
 
 .. _python_api_pks_and_rows_where:
@@ -2422,26 +2422,6 @@ For example:
     #    [thumbnail] BLOB
     # )
 
-.. _find_spatialite:
-
-Finding SpatiaLite
-==================
-
-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)
-
 .. _python_api_register_function:
 
 Registering custom SQL functions
@@ -2522,3 +2502,38 @@ If that option isn't relevant to your use-case you can to quote a string for use
     "'hello'"
     >>> db.quote("hello'this'has'quotes")
     "'hello''this''has''quotes'"
+
+.. _python_api_gis:
+
+Spatialite helpers
+==================
+
+`SpatiaLite <https://www.gaia-gis.it/fossil/libspatialite/index>`__ is a geographic extension to SQLite (similar to PostgreSQL + PostGIS). Using requires finding, loading and initializing the extension, adding geometry columns to existing tables and optionally creating spatial indexes. The utilities here help streamline that setup.
+
+.. _python_api_gis_init_spatialite:
+
+Initialize Spatialite
+---------------------
+
+.. automethod:: sqlite_utils.db.Database.init_spatialite
+
+.. _python_api_gis_find_spatialite:
+
+Finding Spatialite
+------------------
+
+.. autofunction:: sqlite_utils.utils.find_spatialite
+
+.. _python_api_gis_add_geometry_column:
+
+Adding geometry columns
+-----------------------
+
+.. automethod:: sqlite_utils.db.Table.add_geometry_column
+
+.. _python_api_gis_create_spatial_index:
+
+Creating a spatial index
+------------------------
+
+.. automethod:: sqlite_utils.db.Table.create_spatial_index

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

tests/test_gis.py

@@ -0,0 +1,83 @@
+import pytest
+from sqlite_utils.utils import find_spatialite
+from sqlite_utils.db import Database
+from sqlite_utils.utils import sqlite3
+
+pytestmark = [
+    pytest.mark.skipif(
+        not find_spatialite(), reason="Could not find SpatiaLite extension"
+    ),
+    pytest.mark.skipif(
+        not hasattr(sqlite3.Connection, "enable_load_extension"),
+        reason="sqlite3.Connection missing enable_load_extension",
+    ),
+]
+
+
+def test_find_spatialite():
+    spatialite = find_spatialite()
+    assert spatialite is None or isinstance(spatialite, str)
+
+
+def test_init_spatialite():
+    db = Database(memory=True)
+    spatialite = find_spatialite()
+    db.init_spatialite(spatialite)
+    assert "spatial_ref_sys" in db.table_names()
+
+
+def test_add_geometry_column():
+    db = Database(memory=True)
+    spatialite = find_spatialite()
+    db.init_spatialite(spatialite)
+
+    # create a table first
+    table = db.create_table("locations", {"id": str, "properties": str})
+    table.add_geometry_column(
+        column_name="geometry",
+        geometry_type="Point",
+        srid=4326,
+        coord_dimension=2,
+    )
+
+    assert db["geometry_columns"].get(["locations", "geometry"]) == {
+        "f_table_name": "locations",
+        "f_geometry_column": "geometry",
+        "geometry_type": 1,  # point
+        "coord_dimension": 2,
+        "srid": 4326,
+        "spatial_index_enabled": 0,
+    }
+
+
+def test_create_spatial_index():
+    db = Database(memory=True)
+    spatialite = find_spatialite()
+    assert db.init_spatialite(spatialite)
+
+    # create a table, add a geometry column with default values
+    table = db.create_table("locations", {"id": str, "properties": str})
+    assert table.add_geometry_column("geometry", "Point")
+
+    # index it
+    assert table.create_spatial_index("geometry")
+
+    assert "idx_locations_geometry" in db.table_names()
+
+
+def test_double_create_spatial_index():
+    db = Database(memory=True)
+    spatialite = find_spatialite()
+    db.init_spatialite(spatialite)
+
+    # create a table, add a geometry column with default values
+    table = db.create_table("locations", {"id": str, "properties": str})
+    table.add_geometry_column("geometry", "Point")
+
+    # index it, return True
+    assert table.create_spatial_index("geometry")
+
+    assert "idx_locations_geometry" in db.table_names()
+
+    # call it again, return False
+    assert not table.create_spatial_index("geometry")

tests/test_utils.py

@@ -22,11 +22,6 @@ def test_decode_base64_values(input, expected, should_be_is):
         assert actual == expected
 
 
-def test_find_spatialite():
-    spatialite = utils.find_spatialite()
-    assert spatialite is None or isinstance(spatialite, str)
-
-
 @pytest.mark.parametrize(
     "size,expected",
     (