TJWaterServerBinary/app/infra/geoserver/client.py

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}"
            )