from __future__ import annotations from dataclasses import dataclass import math import time from xml.etree import ElementTree import httpx from app.core.config import settings PROJECT_LAYER_NAMES = ( "junctions", "pipes", "pumps", "reservoirs", "scada_devices", "tanks", "valves", ) class GeoServerProvisioningError(RuntimeError): pass @dataclass(frozen=True) class GeoServerDatabaseConfig: host: str port: str user: str password: str def _longitude(web_mercator_x: float) -> float: return web_mercator_x * 180.0 / 20037508.342789244 def _latitude(web_mercator_y: float) -> float: degrees = web_mercator_y * 180.0 / 20037508.342789244 return 180.0 / math.pi * ( 2.0 * math.atan(math.exp(degrees * math.pi / 180.0)) - math.pi / 2.0 ) class GeoServerAdminClient: def __init__(self, client: httpx.Client | None = None) -> None: if not settings.GEOSERVER_USERNAME or not settings.GEOSERVER_PASSWORD: raise GeoServerProvisioningError( "GEOSERVER_USERNAME and GEOSERVER_PASSWORD must be configured" ) self._owns_client = client is None self._client = client or httpx.Client( base_url=settings.GEOSERVER_URL.rstrip("/"), auth=(settings.GEOSERVER_USERNAME, settings.GEOSERVER_PASSWORD), timeout=30.0, ) def __enter__(self) -> "GeoServerAdminClient": return self def __exit__(self, *_args) -> None: if self._owns_client: self._client.close() @staticmethod def database_config() -> GeoServerDatabaseConfig: return GeoServerDatabaseConfig( host=settings.GEOSERVER_DB_HOST or settings.DB_HOST, port=settings.GEOSERVER_DB_PORT or settings.DB_PORT, user=settings.GEOSERVER_DB_USER or settings.DB_USER, password=settings.GEOSERVER_DB_PASSWORD or settings.DB_PASSWORD, ) def _request(self, method: str, path: str, **kwargs) -> httpx.Response: response = self._client.request(method, path, **kwargs) if response.is_error: raise GeoServerProvisioningError( f"GeoServer {method} {path} failed with HTTP {response.status_code}" ) return response def check_ready(self) -> None: self._request("GET", "/rest/about/version.json") def workspace_exists(self, workspace: str) -> bool: response = self._client.get(f"/rest/workspaces/{workspace}.json") if response.status_code == 404: return False if response.is_error: raise GeoServerProvisioningError( f"GeoServer workspace check failed with HTTP {response.status_code}" ) return True def create_project_workspace( self, *, workspace: str, database_name: str, map_bbox: tuple[float, float, float, float], ) -> tuple[str, ...]: if self.workspace_exists(workspace): raise ValueError(f"GeoServer workspace {workspace!r} already exists") self._request( "POST", "/rest/workspaces", json={"workspace": {"name": workspace}}, ) database = self.database_config() entries = { "dbtype": "postgis", "host": database.host, "port": database.port, "database": database_name, "schema": "gis", "user": database.user, "passwd": database.password, "namespace": workspace, "Expose primary keys": "true", "Estimated extends": "true", "validate connections": "true", "min connections": "1", "max connections": "10", "Connection timeout": "20", } self._request( "POST", f"/rest/workspaces/{workspace}/datastores", json={ "dataStore": { "name": workspace, "description": f"{workspace} GIS materialized views", "type": "PostGIS", "enabled": True, "connectionParameters": { "entry": [ {"@key": key, "$": value} for key, value in entries.items() ] }, } }, ) minx, miny, maxx, maxy = map_bbox geographic_bbox = { "minx": _longitude(minx), "miny": _latitude(miny), "maxx": _longitude(maxx), "maxy": _latitude(maxy), "crs": "EPSG:4326", } native_bbox = { "minx": minx, "miny": miny, "maxx": maxx, "maxy": maxy, "crs": "EPSG:3857", } for layer in PROJECT_LAYER_NAMES: self._request( "POST", f"/rest/workspaces/{workspace}/datastores/{workspace}/featuretypes", json={ "featureType": { "name": layer, "nativeName": layer, "title": layer, "srs": "EPSG:3857", "projectionPolicy": "FORCE_DECLARED", "enabled": True, "advertised": True, "nativeBoundingBox": native_bbox, "latLonBoundingBox": geographic_bbox, } }, ) self._set_client_cache(workspace, layer) return PROJECT_LAYER_NAMES def _set_client_cache(self, workspace: str, layer: str) -> None: path = f"/gwc/rest/layers/{workspace}:{layer}.xml" response: httpx.Response | None = None for _ in range(20): response = self._client.get(path) if response.status_code == 200: break if response.status_code != 404: raise GeoServerProvisioningError( f"GeoWebCache layer check failed with HTTP {response.status_code}" ) time.sleep(0.1) if response is None or response.status_code != 200: raise GeoServerProvisioningError( f"GeoWebCache layer {workspace}:{layer} was not registered" ) root = ElementTree.fromstring(response.content) expiry = root.find("expireClients") if expiry is None: expiry = ElementTree.SubElement(root, "expireClients") expiry.text = str(settings.GEOSERVER_CLIENT_CACHE_SECONDS) self._request( "PUT", path, content=ElementTree.tostring( root, encoding="utf-8", xml_declaration=True, ), headers={"Content-Type": "application/xml"}, ) def delete_workspace(self, workspace: str) -> None: response = self._client.delete( f"/rest/workspaces/{workspace}", params={"recurse": "true"}, ) if response.status_code == 404: return if response.is_error: raise GeoServerProvisioningError( f"GeoServer workspace cleanup failed with HTTP {response.status_code}" )