Add SpatiaLite CLI helpers

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 from expected locations
+    if init_spatialite:
+        db.init_spatialite()
+
+    # load spatialite or another extension from a custom location
+    if load_extension:
+        _load_extensions(db, load_extension)
+
     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,106 @@ def convert(
             )
 
 
+@cli.command(
+    "add-geometry-column",
+    help="""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.""",
+)
+@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",
+)
+@click.option("--srid", type=int, default=4326)
+@click.option("--dimensions", "coord_dimension", type=str, default="XY")
+@click.option("--not-null", "not_null", is_flag=True)
+@load_extension_option
+def add_geometry_column(
+    db_path,
+    table,
+    column_name,
+    geometry_type,
+    srid,
+    coord_dimension,
+    not_null,
+    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",
+    help="""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.""",
+)
+@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):
+    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

@@ -792,34 +792,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,5 +1,9 @@
+import json
 import pytest
+
+from click.testing import CliRunner
 from sqlite_utils.utils import find_spatialite
+from sqlite_utils.cli import cli
 from sqlite_utils.db import Database
 from sqlite_utils.utils import sqlite3
 
@@ -14,6 +18,7 @@ pytestmark = [
 ]
 
 
+# python API tests
 def test_find_spatialite():
     spatialite = find_spatialite()
     assert spatialite is None or isinstance(spatialite, str)
@@ -81,3 +86,150 @@ 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 "idx_locations_geometry" in db.table_names()