ARRAY of POINT from PostgreSQL in SQLModel about sqlmodel HOT 4 OPEN

from pydantic_geojson import FeatureModel
from geoalchemy2 import WKBElement
from geoalchemy2.shape import to_shape
from shapely import to_geojson
import json

class WKBToGeoJSON(FeatureModel):
    @classmethod
    def validate(cls, value):
        if not isinstance(value, WKBElement):
            raise TypeError("value must be a WKBElement object")
        shape = to_shape(value)
        if shape.type != "Point":
            raise TypeError("value must be a Point")
        return json.loads(to_geojson(shape))

class Asset(BaseModel, table=True):
    name: Optional[str] = Field(default=None)
    description: Optional[str] = Field(default=None)
    location: WKBToGeoJSON = Field(
        default=None, sa_column=Column(Geometry("POINT", srid=4326))
    )

import uuid
from typing import Literal, Optional

from fastapi import FastAPI, HTTPException

import geoalchemy2 as ga
from sqlmodel import Field, Session, SQLModel, create_engine, select
from sqlalchemy.dialects.postgresql import JSONB


engine = create_engine("postgresql://postgres:password@postgres/mydb")
app = FastAPI()


class GeoJSONGeometry(SQLModel):
    """A GeoJSON geometry fragment."""

    type: Literal["Point", "LineString", "Polygon"]
    coordinates: tuple[float, float] | list[tuple[float, float]] | list[list[tuple[float, float]]]


class Feature(SQLModel, table=True):
    id: Optional[uuid.UUID] = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        index=True,
        nullable=False,
    )
    geometry: bytes = Field(
        sa_type=ga.Geometry,
        nullable=False,
    )
    properties: Optional[dict] = Field(
        default=None,
        sa_type=JSONB,
        nullable=True,
    )


class FeatureCreate(SQLModel):
    geometry: GeoJSONGeometry = Field(
        default=None,
        sa_type=JSONB,
        nullable=False,
    )
    properties: Optional[dict]


class FeatureRead(FeatureCreate):
    id: uuid.UUID
    geometry: GeoJSONGeometry
    properties: Optional[dict]


@app.post("/features/", response_model=FeatureRead)
def create_feature(feature: FeatureCreate):
    with Session(engine) as session:
        db_feature = Feature(
            geometry=ga.functions.ST_SetSRID(
                ga.functions.ST_GeomFromGeoJSON(feature.geometry.model_dump_json()),
                0,  # PostGIS >3.0 ST_GeomFromGeoJSON sets the srid to 4326 by default
            ),
            properties=feature.properties,
        )
        session.add(db_feature)
        session.commit()
        session.refresh(db_feature)
        geojson = session.exec(
            select(ga.functions.ST_AsGeoJSON(db_feature.geometry))
        ).one()
    response = Feature(
        id=db_feature.id,
        geometry=GeoJSONGeometry.model_validate_json(geojson),
        properties=feature.properties,
    )
    return response


@app.get("/features/{feature_id}", response_model=FeatureRead)
def read_feature(feature_id: uuid.UUID):
    with Session(engine) as session:
        db_feature = session.get(Feature, feature_id)
        if db_feature is None:
            raise HTTPException(status_code=404, detail="Feature not found")
        geojson = session.exec(
            select(ga.functions.ST_AsGeoJSON(db_feature.geometry))
        ).one()
    response = FeatureRead(
        id=db_feature.id,
        geometry=GeoJSONGeometry.model_validate_json(geojson),
        properties=db_feature.properties,
    )
    return response