Return status. Handle double calls.

This commit is contained in:
Chris Amico 2022-01-18 09:59:03 -05:00
commit 0df57367e9
3 changed files with 55 additions and 11 deletions

View file

@ -2431,7 +2431,7 @@ Spatialite helpers
.. _init_spatialite:
Initialize Spatialite
-----------------------
----------------------
.. autofunction:: sqlite_utils.gis.init_spatialite
@ -2439,7 +2439,7 @@ Initialize Spatialite
.. _find_spatialite:
Finding Spatialite
-----------------
------------------
.. autofunction:: sqlite_utils.gis.find_spatialite

View file

@ -30,12 +30,14 @@ def find_spatialite() -> str:
return None
def init_spatialite(db: Database, path: str) -> None:
def init_spatialite(db: Database, path: str) -> bool:
"""
The ``init_spatialite`` function will load and initalize 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 initalized.
.. code-block:: python
from sqlite_utils.gis import find_spatialite, init_spatialite
@ -60,8 +62,10 @@ def init_spatialite(db: Database, path: str) -> None:
db.conn.load_extension(path)
# Initialize SpatiaLite if not yet initialized
if "spatial_ref_sys" in db.table_names():
return
db.execute("select InitSpatialMetadata(1)")
return False
cursor = db.execute("select InitSpatialMetadata(1)")
result = cursor.fetchone()
return result and bool(result[0])
def add_geometry_column(
@ -71,7 +75,7 @@ def add_geometry_column(
srid: int = 4326,
coord_dimension: str = "XY",
not_null: bool = False,
) -> None:
) -> bool:
"""
In Spatialite, a geometry column can only be added to an existing table.
To do so, use ``add_geometry_column``, passing in a :ref:`table <reference_db_table>`
@ -81,6 +85,8 @@ def add_geometry_column(
`SRID 4326 <https://spatialreference.org/ref/epsg/wgs-84/>`__. These can be customized using
the ``column_name`` and ``srid`` arguments.
Returns True if the column was successfully added, False if not.
.. code-block:: python
from sqlite_utils.gis import find_spatialite, init_spatialite, add_geometry_column
@ -93,18 +99,24 @@ def add_geometry_column(
add_geometry_column(db["locations"], "POINT")
"""
table.db.execute(
cursor = table.db.execute(
"SELECT AddGeometryColumn(?, ?, ?, ?, ?, ?);",
[table.name, column_name, srid, geometry_type, coord_dimension, int(not_null)],
)
result = cursor.fetchone()
return result and bool(result[0])
def create_spatial_index(table: Table, column_name: str = "geometry") -> None:
def create_spatial_index(table: Table, column_name: str = "geometry") -> bool:
"""
A spatial index allows for significantly faster bounding box queries.
To create on, use ``create_spatial_index`` with a :ref:`table <reference_db_table>`
and 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
from sqlite_utils.gis import add_geometry_column, create_spatial_index
@ -122,4 +134,11 @@ def create_spatial_index(table: Table, column_name: str = "geometry") -> None:
# CREATE VIRTUAL TABLE "idx_locations_geometry" USING rtree(pkid, xmin, xmax, ymin, ymax)
"""
table.db.execute("select CreateSpatialIndex(?, ?)", [table.name, column_name])
if f"idx_{table.name}_{column_name}" in table.db.table_names():
return False
cursor = table.db.execute(
"select CreateSpatialIndex(?, ?)", [table.name, column_name]
)
result = cursor.fetchone()
return result and bool(result[0])

View file

@ -70,6 +70,28 @@ def test_add_geometry_column():
reason="sqlite3.Connection missing enable_load_extension",
)
def test_create_spatial_index():
db = Database(memory=True)
spatialite = gis.find_spatialite()
assert gis.init_spatialite(db, spatialite)
# create a table, add a geometry column with default values
db.create_table("locations", {"id": str, "properties": str})
assert gis.add_geometry_column(db["locations"], "Point", "geometry")
# index it
assert gis.create_spatial_index(db["locations"], "geometry")
assert "idx_locations_geometry" in db.table_names()
@pytest.mark.skipif(
not gis.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_double_create_spatial_index():
db = Database(memory=True)
spatialite = gis.find_spatialite()
gis.init_spatialite(db, spatialite)
@ -78,7 +100,10 @@ def test_create_spatial_index():
db.create_table("locations", {"id": str, "properties": str})
gis.add_geometry_column(db["locations"], "Point", "geometry")
# index it
gis.create_spatial_index(db["locations"], "geometry")
# index it, return True
assert gis.create_spatial_index(db["locations"], "geometry")
assert "idx_locations_geometry" in db.table_names()
# call it again, return False
assert not gis.create_spatial_index(db["locations"], "geometry")