Add SpatiaLite helpers to CLI (#407)

* Add SpatiaLite CLI helpers
* Add docs for spaitalite helpers
* Fix flake8 issues and add more detail on spatial types
* Run cog and add some help text.
* Use SpatiaLite when calculating coverage, refs #407

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

.github/workflows/test-coverage.yml

@@ -24,6 +24,8 @@ jobs:
         key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}
         restore-keys: |
           ${{ runner.os }}-pip-
+    - name: Install SpatiaLite
+      run: sudo apt-get install libsqlite3-mod-spatialite
     - name: Install Python dependencies
       run: |
         python -m pip install --upgrade pip

docs/cli-reference.rst

@@ -772,8 +772,10 @@ See :ref:`cli_create_database`.
           sqlite-utils create-database trees.db
 
     Options:
-      --enable-wal  Enable WAL mode on the created database
-      -h, --help    Show this message and exit.
+      --enable-wal           Enable WAL mode on the created database
+      --init-spatialite      Enable SpatiaLite on the created database
+      --load-extension TEXT  SQLite extensions to load
+      -h, --help             Show this message and exit.
 
 
 create-table
@@ -1226,4 +1228,49 @@ See :ref:`cli_drop_view`.
       -h, --help             Show this message and exit.
 
 
+add-geometry-column
+===================
+
+::
+
+    Usage: sqlite-utils add-geometry-column [OPTIONS] DB_PATH TABLE COLUMN_NAME
+
+      Add a SpatiaLite geometry column to an existing table. Requires SpatiaLite
+      extension.
+
+      By default, this command will try to load the SpatiaLite extension from usual
+      paths. To load it from a specific path, use --load-extension.
+
+    Options:
+      -t, --type [POINT|LINESTRING|POLYGON|MULTIPOINT|MULTILINESTRING|MULTIPOLYGON|GEOMETRYCOLLECTION|GEOMETRY]
+                                      Specify a geometry type for this column.
+                                      [default: GEOMETRY]
+      --srid INTEGER                  Spatial Reference ID. See
+                                      https://spatialreference.org for details on
+                                      specific projections.  [default: 4326]
+      --dimensions TEXT               Coordinate dimensions. Use XYZ for three-
+                                      dimensional geometries.
+      --not-null                      Add a NOT NULL constraint.
+      --load-extension TEXT           SQLite extensions to load
+      -h, --help                      Show this message and exit.
+
+
+create-spatial-index
+====================
+
+::
+
+    Usage: sqlite-utils create-spatial-index [OPTIONS] DB_PATH TABLE COLUMN_NAME
+
+      Create a spatial index on a SpatiaLite geometry column. The table and geometry
+      column must already exist before trying to add a spatial index.
+
+      By default, this command will try to load the SpatiaLite extension from usual
+      paths. To load it from a specific path, use --load-extension.
+
+    Options:
+      --load-extension TEXT  SQLite extensions to load
+      -h, --help             Show this message and exit.
+
+
 .. [[[end]]]

docs/cli.rst

@@ -724,6 +724,14 @@ To enable :ref:`cli_wal` on the newly created database add the ``--enable-wal``
 
     $ sqlite-utils create-database empty.db --enable-wal
 
+To enable SpatiaLite metadata on a newly created database, add the ``--init-spatialite`` flag::
+
+    $ sqlite-utils create-database empty.db --init-spatialite
+
+That will look for SpatiaLite in a set of predictable locations. To load it from somewhere else, use the ``--load-extension`` option::
+
+    $ sqlite-utils create-database empty.db --init-spatialite --load-extension /path/to/spatialite.so
+
 .. _cli_inserting_data:
 
 Inserting JSON data
@@ -1975,3 +1983,34 @@ Since `SpatiaLite <https://www.gaia-gis.it/fossil/libspatialite/index>`__ is com
 
     $ sqlite-utils memory "select spatialite_version()" --load-extension=spatialite
     [{"spatialite_version()": "4.3.0a"}]
+
+
+SpatiaLite helpers
+==================
+
+`SpatiaLite <https://www.gaia-gis.it/fossil/libspatialite/home>`_ adds geographic capability to SQLite (similar to how PostGIS builds on PostgreSQL). The `SpatiaLite cookbook <http://www.gaia-gis.it/gaia-sins/spatialite-cookbook-5/index.html>`_ is a good resource for learning what's possible with it.
+
+You can convert an existing table to a geographic table by adding a geometry column, use the `sqlite-utils add-geometry-column` command::
+
+    $ sqlite-utils add-geometry-column spatial.db locations geometry --type POLYGON --srid 4326
+
+The table (``locations`` in the example above) must already exist before adding a geometry column. Use ``sqlite-utils create-table`` first, then ``add-geometry-column``.
+
+Use the ``--type`` option to specify a geometry type. By default, ``add-geometry-column`` uses a generic ``GEOMETRY``, which will work with any type, though it may not be supported by some desktop GIS applications. 
+
+Eight (case-insensitive) types are allowed:
+
+ * POINT
+ * LINESTRING
+ * POLYGON
+ * MULTIPOINT
+ * MULTILINESTRING
+ * MULTIPOLYGON
+ * GEOMETRYCOLLECTION
+ * GEOMETRY
+
+Once you have a geometry column, you can speed up bounding box queries by adding a spatial index::
+
+    $ sqlite-utils create-spatial-index spatial.db locations geometry
+
+See the `SpatiaLite Cookbook <http://www.gaia-gis.it/gaia-sins/spatialite-cookbook-5/cookbook_topics.03.html#topic_Wonderful_RTree_Spatial_Index>`_ for examples of how to use a spatial index.

sqlite_utils/cli.py

@@ -1363,7 +1363,11 @@ def bulk(
 @click.option(
     "--enable-wal", is_flag=True, help="Enable WAL mode on the created database"
 )
-def create_database(path, enable_wal):
+@click.option(
+    "--init-spatialite", is_flag=True, help="Enable SpatiaLite on the created database"
+)
+@load_extension_option
+def create_database(path, enable_wal, init_spatialite, load_extension):
     """Create a new empty database file
 
     Example:
@@ -1374,6 +1378,15 @@ def create_database(path, enable_wal):
     db = sqlite_utils.Database(path)
     if enable_wal:
         db.enable_wal()
+
+    # load spatialite or another extension from a custom location
+    if load_extension:
+        _load_extensions(db, load_extension)
+
+    # load spatialite from expected locations and initialize metadata
+    if init_spatialite:
+        db.init_spatialite()
+
     db.vacuum()
 
 
@@ -2544,7 +2557,7 @@ def _analyze(db, tables, columns, save):
                     total=len(todo),
                     most_common_rendered=most_common_rendered,
                     least_common_rendered=least_common_rendered,
-                    **column_details._asdict()
+                    **column_details._asdict(),
                 )
             )
             + "\n"
@@ -2701,6 +2714,116 @@ def convert(
             )
 
 
+@cli.command("add-geometry-column")
+@click.argument(
+    "db_path",
+    type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
+    required=True,
+)
+@click.argument("table", type=str)
+@click.argument("column_name", type=str)
+@click.option(
+    "-t",
+    "--type",
+    "geometry_type",
+    type=click.Choice(
+        [
+            "POINT",
+            "LINESTRING",
+            "POLYGON",
+            "MULTIPOINT",
+            "MULTILINESTRING",
+            "MULTIPOLYGON",
+            "GEOMETRYCOLLECTION",
+            "GEOMETRY",
+        ],
+        case_sensitive=False,
+    ),
+    default="GEOMETRY",
+    help="Specify a geometry type for this column.",
+    show_default=True,
+)
+@click.option(
+    "--srid",
+    type=int,
+    default=4326,
+    show_default=True,
+    help="Spatial Reference ID. See https://spatialreference.org for details on specific projections.",
+)
+@click.option(
+    "--dimensions",
+    "coord_dimension",
+    type=str,
+    default="XY",
+    help="Coordinate dimensions. Use XYZ for three-dimensional geometries.",
+)
+@click.option("--not-null", "not_null", is_flag=True, help="Add a NOT NULL constraint.")
+@load_extension_option
+def add_geometry_column(
+    db_path,
+    table,
+    column_name,
+    geometry_type,
+    srid,
+    coord_dimension,
+    not_null,
+    load_extension,
+):
+    """Add a SpatiaLite geometry column to an existing table. Requires SpatiaLite extension.
+    \n\n
+    By default, this command will try to load the SpatiaLite extension from usual paths.
+    To load it from a specific path, use --load-extension."""
+    db = sqlite_utils.Database(db_path)
+    if not db[table].exists():
+        raise click.ClickException(
+            "You must create a table before adding a geometry column"
+        )
+
+    # load spatialite, one way or another
+    if load_extension:
+        _load_extensions(db, load_extension)
+    db.init_spatialite()
+
+    if db[table].add_geometry_column(
+        column_name, geometry_type, srid, coord_dimension, not_null
+    ):
+        click.echo(f"Added {geometry_type} column {column_name} to {table}")
+
+
+@cli.command("create-spatial-index")
+@click.argument(
+    "db_path",
+    type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
+    required=True,
+)
+@click.argument("table", type=str)
+@click.argument("column_name", type=str)
+@load_extension_option
+def create_spatial_index(db_path, table, column_name, load_extension):
+    """Create a spatial index on a SpatiaLite geometry column.
+    The table and geometry column must already exist before trying to add a spatial index.
+    \n\n
+    By default, this command will try to load the SpatiaLite extension from usual paths.
+    To load it from a specific path, use --load-extension."""
+    db = sqlite_utils.Database(db_path)
+    if not db[table].exists():
+        raise click.ClickException(
+            "You must create a table and add a geometry column before creating a spatial index"
+        )
+
+    # load spatialite
+    if load_extension:
+        _load_extensions(db, load_extension)
+    db.init_spatialite()
+
+    if column_name not in db[table].columns_dict:
+        raise click.ClickException(
+            "You must add a geometry column before creating a spatial index"
+        )
+
+    db[table].create_spatial_index(column_name)
+
+
 def _render_common(title, values):
     if values is None:
         return ""

tests/test_cli.py

@@ -7,7 +7,6 @@ from unittest import mock
 import json
 import os
 import pytest
-from sqlite_utils.utils import sqlite3, find_spatialite
 import textwrap
 
 from .utils import collapse_whitespace
@@ -792,34 +791,6 @@ def test_query_raw(db_path, content, is_binary):
         assert result.output == str(content)
 
 
-@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",
-)
-@pytest.mark.parametrize("use_spatialite_shortcut", [True, False])
-def test_query_load_extension(use_spatialite_shortcut):
-    # Without --load-extension:
-    result = CliRunner().invoke(cli.cli, [":memory:", "select spatialite_version()"])
-    assert result.exit_code == 1
-    assert "no such function: spatialite_version" in result.output
-    # With --load-extension:
-    if use_spatialite_shortcut:
-        load_extension = "spatialite"
-    else:
-        load_extension = find_spatialite()
-    result = CliRunner().invoke(
-        cli.cli,
-        [
-            ":memory:",
-            "select spatialite_version()",
-            "--load-extension={}".format(load_extension),
-        ],
-    )
-    assert result.exit_code == 0, result.stdout
-    assert ["spatialite_version()"] == list(json.loads(result.output)[0].keys())
-
-
 def test_query_memory_does_not_create_file(tmpdir):
     owd = os.getcwd()
     try:

tests/test_gis.py

@@ -1,7 +1,10 @@
+import json
 import pytest
-from sqlite_utils.utils import find_spatialite
+
+from click.testing import CliRunner
+from sqlite_utils.cli import cli
 from sqlite_utils.db import Database
-from sqlite_utils.utils import sqlite3
+from sqlite_utils.utils import find_spatialite, sqlite3
 
 pytestmark = [
     pytest.mark.skipif(
@@ -14,6 +17,7 @@ pytestmark = [
 ]
 
 
+# python API tests
 def test_find_spatialite():
     spatialite = find_spatialite()
     assert spatialite is None or isinstance(spatialite, str)
@@ -81,3 +85,152 @@ def test_double_create_spatial_index():
 
     # call it again, return False
     assert not table.create_spatial_index("geometry")
+
+
+# cli tests
+@pytest.mark.parametrize("use_spatialite_shortcut", [True, False])
+def test_query_load_extension(use_spatialite_shortcut):
+    # Without --load-extension:
+    result = CliRunner().invoke(cli, [":memory:", "select spatialite_version()"])
+    assert result.exit_code == 1
+    assert "no such function: spatialite_version" in result.output
+    # With --load-extension:
+    if use_spatialite_shortcut:
+        load_extension = "spatialite"
+    else:
+        load_extension = find_spatialite()
+    result = CliRunner().invoke(
+        cli,
+        [
+            ":memory:",
+            "select spatialite_version()",
+            "--load-extension={}".format(load_extension),
+        ],
+    )
+    assert result.exit_code == 0, result.stdout
+    assert ["spatialite_version()"] == list(json.loads(result.output)[0].keys())
+
+
+def test_cli_create_spatialite(tmpdir):
+    # sqlite-utils create test.db --init-spatialite
+    db_path = tmpdir / "created.db"
+    result = CliRunner().invoke(
+        cli, ["create-database", str(db_path), "--init-spatialite"]
+    )
+
+    assert 0 == result.exit_code
+    assert db_path.exists()
+    assert db_path.read_binary()[:16] == b"SQLite format 3\x00"
+
+    db = Database(str(db_path))
+    assert "spatial_ref_sys" in db.table_names()
+
+
+def test_cli_add_geometry_column(tmpdir):
+    # create a rowid table with one column
+    db_path = tmpdir / "spatial.db"
+    db = Database(str(db_path))
+    db.init_spatialite()
+
+    table = db["locations"].create({"name": str})
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "add-geometry-column",
+            str(db_path),
+            table.name,
+            "geometry",
+            "--type",
+            "POINT",
+        ],
+    )
+
+    assert 0 == result.exit_code
+
+    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_cli_add_geometry_column_options(tmpdir):
+    # create a rowid table with one column
+    db_path = tmpdir / "spatial.db"
+    db = Database(str(db_path))
+    db.init_spatialite()
+    table = db["locations"].create({"name": str})
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "add-geometry-column",
+            str(db_path),
+            table.name,
+            "geometry",
+            "-t",
+            "POLYGON",
+            "--srid",
+            "3857",  # https://epsg.io/3857
+            "--not-null",
+        ],
+    )
+
+    assert 0 == result.exit_code
+
+    assert db["geometry_columns"].get(["locations", "geometry"]) == {
+        "f_table_name": "locations",
+        "f_geometry_column": "geometry",
+        "geometry_type": 3,  # polygon
+        "coord_dimension": 2,
+        "srid": 3857,
+        "spatial_index_enabled": 0,
+    }
+
+    column = table.columns[1]
+    assert column.notnull
+
+
+def test_cli_add_geometry_column_invalid_type(tmpdir):
+    # create a rowid table with one column
+    db_path = tmpdir / "spatial.db"
+    db = Database(str(db_path))
+    db.init_spatialite()
+
+    table = db["locations"].create({"name": str})
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "add-geometry-column",
+            str(db_path),
+            table.name,
+            "geometry",
+            "--type",
+            "NOT-A-TYPE",
+        ],
+    )
+
+    assert 2 == result.exit_code
+
+
+def test_cli_create_spatial_index(tmpdir):
+    # create a rowid table with one column
+    db_path = tmpdir / "spatial.db"
+    db = Database(str(db_path))
+    db.init_spatialite()
+
+    table = db["locations"].create({"name": str})
+    table.add_geometry_column("geometry", "POINT")
+
+    result = CliRunner().invoke(
+        cli, ["create-spatial-index", str(db_path), table.name, "geometry"]
+    )
+
+    assert 0 == result.exit_code
+
+    assert "idx_locations_geometry" in db.table_names()