diff --git a/docs/python-api.rst b/docs/python-api.rst index df9449f..4da350b 100644 --- a/docs/python-api.rst +++ b/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 `__ 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 `__ 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 diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 8b1b5a4..b137d9b 100644 --- a/sqlite_utils/cli.py +++ b/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") diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 06e7557..4ab8862 100644 --- a/sqlite_utils/db.py +++ b/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 `__. 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): diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 0777f30..c2a7c91 100644 --- a/sqlite_utils/utils.py +++ b/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 `__ 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 diff --git a/tests/test_gis.py b/tests/test_gis.py new file mode 100644 index 0000000..3bf0d5f --- /dev/null +++ b/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") diff --git a/tests/test_utils.py b/tests/test_utils.py index 8630e28..66ce413 100644 --- a/tests/test_utils.py +++ b/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", (