From 138914a7d8458db1c20a925f49f996314c787975 Mon Sep 17 00:00:00 2001 From: James Bannister Date: Tue, 7 May 2024 10:15:38 +0100 Subject: [PATCH] SQLAlchemy current iteration --- application/routers/tiles.py | 185 ++++++++++++----------------------- 1 file changed, 61 insertions(+), 124 deletions(-) diff --git a/application/routers/tiles.py b/application/routers/tiles.py index 1ef1bc2e..50325a02 100644 --- a/application/routers/tiles.py +++ b/application/routers/tiles.py @@ -1,134 +1,68 @@ -import logging - -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, Depends from fastapi.responses import StreamingResponse - -import psycopg2 +from sqlalchemy.orm import Session +from sqlalchemy import func from io import BytesIO -from application.settings import get_settings +from application.db.models import EntityOrm +from application.db.session import get_session router = APIRouter() -logger = logging.getLogger(__name__) - -DATABASE = {"user": "", "password": "", "host": "", "port": "5432", "database": ""} - -DATABASE_CONNECTION = None - -QUERY_PARAMS = { - "table1": "entity t1", - "srid": "4326", - "geomColumn": "t1.geometry", - "attrColumns": "t1.entity, t1.name, t1.reference", -} - # ============================================================ # Helper Funcs # ============================================================ -def get_db_connection(): - conn_str = get_settings() - - DATABASE["user"] = conn_str.READ_DATABASE_URL.user - DATABASE["password"] = conn_str.READ_DATABASE_URL.password - DATABASE["host"] = conn_str.READ_DATABASE_URL.host - DATABASE["database"] = conn_str.READ_DATABASE_URL.path.split("/")[1] - -get_db_connection() - -# Do the tile x/y coordinates make sense at this zoom level? +# Validate tile x/y coordinates at the given zoom level def tile_is_valid(tile): - if not ("x" in tile and "y" in tile and "zoom" in tile): - return False - - if "format" not in tile or tile["format"] not in ["pbf", "mvt"]: - return False - size = 2 ** tile["zoom"] - - if tile["x"] >= size or tile["y"] >= size: - return False - - if tile["x"] < 0 or tile["y"] < 0: - return False - - return True - - -def build_db_query(tile): - qry_params = QUERY_PARAMS.copy() - qry_params["dataset"] = tile["dataset"] - qry_params["x"] = tile["x"] - qry_params["y"] = tile["y"] - qry_params["z"] = tile["zoom"] - - query = """ - WITH - webmercator(envelope) AS ( - SELECT ST_TileEnvelope({z}, {x}, {y}) - ), - wgs84(envelope) AS ( - SELECT ST_Transform((SELECT envelope FROM webmercator), {srid}) - ), - b(bounds) AS ( - SELECT ST_MakeEnvelope(-180, -85.0511287798066, 180, 85.0511287798066, {srid}) - ), - geometries(entity, name, reference, wkb_geometry) AS ( - SELECT - {attrColumns}, - CASE WHEN ST_Covers(b.bounds, {geomColumn}) - THEN ST_Transform({geomColumn},{srid}) - ELSE ST_Transform(ST_Intersection(b.bounds, {geomColumn}),{srid}) - END - FROM - {table1} - CROSS JOIN - b - WHERE - {geomColumn} && (SELECT envelope FROM wgs84) - AND - t1.dataset = '{dataset}' - ) - SELECT - ST_AsMVT(tile, '{dataset}') as mvt - FROM ( - SELECT - entity, - name, - reference, - ST_AsMVTGeom(wkb_geometry, (SELECT envelope FROM wgs84)) - FROM geometries - ) AS tile - """.format( - **qry_params + return ( + 0 <= tile["x"] < size + and 0 <= tile["y"] < size + and tile["format"] in ["pbf", "mvt"] ) - return query - - -def sql_to_pbf(sql): - global DATABASE_CONNECTION - # Make and hold connection to database - if not DATABASE_CONNECTION: - try: - DATABASE_CONNECTION = psycopg2.connect(**DATABASE) - except (Exception, psycopg2.Error) as error: - logger.warning(error) - return None - - # Query for MVT - with DATABASE_CONNECTION.cursor() as cur: - cur.execute(sql) - if not cur: - logger.warning(f"sql query failed: {sql}") - return None +# Build the database query using SQLAlchemy ORM to match the direct SQL logic +def build_db_query(tile, session: Session): + envelope = func.ST_TileEnvelope(tile["zoom"], tile["x"], tile["y"]) + webmercator = envelope + srid = 4326 # WGS 84 + wgs84 = func.ST_Transform(webmercator, srid) + bounds = func.ST_MakeEnvelope(-180, -85.0511287798066, 180, 85.0511287798066, srid) + + geometries = ( + session.query( + EntityOrm.entity, + EntityOrm.name, + EntityOrm.reference, + func.ST_AsMVTGeom( + func.CASE( + [ + ( + func.ST_Covers(bounds, EntityOrm.geometry), + func.ST_Transform(EntityOrm.geometry, srid), + ) + ], + else_=func.ST_Transform( + func.ST_Intersection(bounds, EntityOrm.geometry), srid + ), + ), + wgs84, + ), + ) + .filter( + EntityOrm.geometry.ST_Intersects(wgs84), + EntityOrm.dataset == tile["dataset"], + ) + .subquery() + ) - return cur.fetchone()[0] + # Build vector tile + tile_data = session.query(func.ST_AsMVT(geometries, tile["dataset"])).scalar() - return None + return tile_data # ============================================================ @@ -136,21 +70,24 @@ def sql_to_pbf(sql): # ============================================================ -@router.get("/-/tiles/{dataset}/{z}/{x}/{y}.vector.{fmt}") -async def read_tiles_from_postgres(dataset: str, z: int, x: int, y: int, fmt: str): +@router.get("/tiles/{dataset}/{z}/{x}/{y}.{fmt}") +async def read_tiles_from_postgres( + dataset: str, + z: int, + x: int, + y: int, + fmt: str, + session: Session = Depends(get_session), +): tile = {"dataset": dataset, "zoom": z, "x": x, "y": y, "format": fmt} - if not tile_is_valid(tile): - raise HTTPException(status_code=400, detail=f"invalid tile path: {tile}") - - sql = build_db_query(tile) - - pbf = sql_to_pbf(sql) + raise HTTPException(status_code=400, detail=f"Invalid tile path: {tile}") - pbf_buffer = BytesIO() - pbf_buffer.write(pbf) - pbf_buffer.seek(0) + tile_data = build_db_query(tile, session) + if not tile_data: + raise HTTPException(status_code=404, detail="Tile data not found") + pbf_buffer = BytesIO(tile_data) resp_headers = { "Access-Control-Allow-Origin": "*", "Content-Type": "application/vnd.mapbox-vector-tile",