From 414d0cf7327008d2d56d9ef5b9391cc04b6fa5ca Mon Sep 17 00:00:00 2001 From: Chris Amico Date: Sun, 13 Feb 2022 19:36:52 -0500 Subject: [PATCH] Add SpatiaLite CLI helpers --- sqlite_utils/cli.py | 117 +++++++++++++++++++++++++++++++++- tests/test_cli.py | 28 -------- tests/test_gis.py | 152 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 267 insertions(+), 30 deletions(-) diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 9e0289f..0809177 100644 --- a/sqlite_utils/cli.py +++ b/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 "" diff --git a/tests/test_cli.py b/tests/test_cli.py index 244029d..a97c047 100644 --- a/tests/test_cli.py +++ b/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: diff --git a/tests/test_gis.py b/tests/test_gis.py index 3bf0d5f..4eccbf2 100644 --- a/tests/test_gis.py +++ b/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()