geolocation_mixin.py

geolocation_mixin.md

Geolocation with SQLModel

Need to use postgis with SQLModel? This example will show you how.

This gist demonstrates how to integrate a geolocation field into your SQLModel projects by directly adding a geometry column representing a point on Earth. The GeolocationMixin class provides an easy way to store and serialize geospatial data, specifically using PostGIS to handle POINT geometry types.

To run this example, you'll need a PostGIS-enabled database. You can set up PostGIS using Docker with the following command:

docker run -d --rm \
    -e POSTGRES_DB=dbuser \
    -e POSTGRES_USER=dbpassword \
    -e POSTGRES_PASSWORD=yourdb \
    -e PGDATA=/var/lib/postgresql/data/pgdata \
    -v pgvolume:/var/lib/postgresql/data \
    -p 5532:5432 \
    --name pgvector \
    --platform linux/amd64 \
    ivanlonel/postgis-with-extensions

After launching the container, create the PostGIS extension in your PostgreSQL database:

CREATE EXTENSION postgis;

Since you're mapping the container to your local filesystem, this should persist across runs. With this setup, you're ready to add geospatial capabilities to your SQLModel models.

Dependencies

To work with geospatial data in this example, you'll need to install the following Python packages. I use poetry:

poetry add geojson[shapely] shapely geoalchemy2

GeolocationMixin Example

The following code defines a GeolocationMixin class that introduces a geolocation field geom, which stores a geometric point using the POINT type in PostGIS, represented by latitude and longitude. The serialize_geom method serializes this geometric data into GeoJSON format, handling different input types such as WKT strings and WKB elements.

from typing import Any, Optional

import geojson
from geoalchemy2 import Geometry, WKBElement
from pydantic import field_serializer
from shapely import wkt
from shapely.geometry import mapping
from shapely.wkb import loads as wkb_loads
from sqlalchemy import Column
from sqlmodel import Field

class GeolocationMixin:
    """Mixin class for geolocation fields."""

    geom: Optional[Any] = Field(
        default=None,
        sa_column=Column(Geometry(geometry_type="POINT", srid=4326)),
        description="The geometry point representing the geolocation, created with a latitude and longitude pair stored as a PostGIS POINT.",
    )

    @field_serializer("geom")
    def serialize_geom(self, geom: Any, _info) -> Optional[str]:
        """Serialize the geometry point to a GeoJSON object."""
        if geom is None:
            return None
        if isinstance(geom, str):
            return geojson.dumps(mapping(wkt.loads(geom)))
        elif isinstance(geom, WKBElement):
            return geojson.dumps(wkb_loads(bytes(geom.data)))

This class can be easily integrated into your SQLModel-based projects to handle geospatial data effectively. Let's take a look how.

First, add the mixin. The following will provie AddressBase with postgis powers.

class AddressBase(SQLModel, GeolocationMixin):
    """Represents a mailing address, physical address, or location."""
    ...

Next, given an example address, validate the model:

# Example dictionary for Address
example_address = {
    "location_name": "Ramsey Park",
    "street_number": "4301",
    "street_name": "Rosedale Ave.",
    "unit": "N/A",
    "city": "Austin",
    "state": "TX",
    "postal_code": "78756",
    "country": "USA",
    "geom": "POINT (-97.7435376 30.313271999999998)",
}

# validate and instantiate the model from the example address
address_instance = AddressBase.model_validate(example_address)

# model dump
print(address_instance.model_dump_json(indent=4))

# schema dump
print(address_instance.model_json_schema())