pbf/navsea_geometry_enable.py

from __future__ import annotations

import os
from dataclasses import dataclass

import pymysql


@dataclass(frozen=True)
class DbConfig:
    host: str = os.getenv("NAVSEA_DB_HOST", "localhost")
    port: int = int(os.getenv("NAVSEA_DB_PORT", "3306"))
    user: str = os.getenv("NAVSEA_DB_USER", "root")
    password: str = os.getenv("NAVSEA_DB_PASSWORD", "2chi9ks2")
    database: str = os.getenv("NAVSEA_DB_NAME", "pbf_analysis")
    unix_socket: str | None = os.getenv("NAVSEA_DB_SOCKET", "/tmp/mysql.sock")
    source_table: str = os.getenv("NAVSEA_GEOMETRY_SOURCE_TABLE", "feature_geometry_stage")


class GeometryEnablement:
    def __init__(self, config: DbConfig) -> None:
        self.config = config

    def connect(self):
        kwargs = {
            "host": self.config.host,
            "port": self.config.port,
            "user": self.config.user,
            "password": self.config.password,
            "database": self.config.database,
            "charset": "utf8mb4",
            "autocommit": False,
        }
        if self.config.unix_socket and self.config.host in {"localhost", "127.0.0.1"}:
            kwargs["unix_socket"] = self.config.unix_socket
        return pymysql.connect(**kwargs)

    def run(self) -> None:
        with self.connect() as conn:
            with conn.cursor() as cur:
                source_mode = self.detect_source_mode(cur)
                self.create_geometry_table(cur, source_mode)
                self.create_spatial_view(cur)
            conn.commit()
        print("Built navsea_feature_geometry and navsea_detection_objects_spatial")

    def detect_source_mode(self, cur) -> str:
        cur.execute("SHOW TABLES LIKE %s", (self.config.source_table,))
        if cur.fetchone() is None:
            raise RuntimeError(
                "geometry staging table missing: "
                f"{self.config.source_table}. "
                "Expected one of: "
                "(feature_id, geometry_wkt), "
                "(feature_id, geometry_wkb), "
                "(feature_id, lon, lat), "
                "or (feature_id, geometry GEOMETRY)."
            )

        cur.execute(
            """
            SELECT COLUMN_NAME, DATA_TYPE
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s
            """,
            (self.config.database, self.config.source_table),
        )
        columns = {name: data_type for name, data_type in cur.fetchall()}

        if "feature_id" not in columns:
            raise RuntimeError(f"{self.config.source_table} missing required column: feature_id")

        if "geometry" in columns and columns["geometry"] == "geometry":
            return "geometry"
        if "geometry_wkt" in columns:
            return "geometry_wkt"
        if "geometry_wkb" in columns:
            return "geometry_wkb"
        if {"lon", "lat"}.issubset(columns):
            return "lonlat"

        raise RuntimeError(
            f"{self.config.source_table} does not expose a supported geometry source. "
            "Supported layouts: geometry GEOMETRY, geometry_wkt, geometry_wkb, or lon/lat."
        )

    def create_geometry_table(self, cur, source_mode: str) -> None:
        cur.execute("DROP TABLE IF EXISTS navsea_feature_geometry")
        cur.execute(
            """
            CREATE TABLE navsea_feature_geometry (
                feature_id BIGINT NOT NULL,
                geometry GEOMETRY NOT NULL,
                geometry_source VARCHAR(32) NOT NULL,
                PRIMARY KEY (feature_id),
                SPATIAL KEY idx_geometry (geometry)
            ) ENGINE=MyISAM DEFAULT CHARSET=utf8mb4
            """
        )

        source_table = self.config.source_table
        if source_mode == "geometry":
            cur.execute(
                f"""
                INSERT INTO navsea_feature_geometry (feature_id, geometry, geometry_source)
                SELECT feature_id, geometry, 'geometry'
                FROM {source_table}
                """
            )
        elif source_mode == "geometry_wkt":
            cur.execute(
                f"""
                INSERT INTO navsea_feature_geometry (feature_id, geometry, geometry_source)
                SELECT feature_id, GeomFromText(geometry_wkt), 'geometry_wkt'
                FROM {source_table}
                WHERE geometry_wkt IS NOT NULL AND geometry_wkt <> ''
                """
            )
        elif source_mode == "geometry_wkb":
            cur.execute(
                f"""
                INSERT INTO navsea_feature_geometry (feature_id, geometry, geometry_source)
                SELECT feature_id, GeomFromWKB(geometry_wkb), 'geometry_wkb'
                FROM {source_table}
                WHERE geometry_wkb IS NOT NULL
                """
            )
        elif source_mode == "lonlat":
            cur.execute(
                f"""
                INSERT INTO navsea_feature_geometry (feature_id, geometry, geometry_source)
                SELECT feature_id, Point(lon, lat), 'lonlat'
                FROM {source_table}
                WHERE lon IS NOT NULL AND lat IS NOT NULL
                """
            )
        else:
            raise RuntimeError(f"unsupported source mode: {source_mode}")

    @staticmethod
    def create_spatial_view(cur) -> None:
        cur.execute("DROP VIEW IF EXISTS navsea_detection_objects_spatial")
        cur.execute(
            """
            CREATE VIEW navsea_detection_objects_spatial AS
            SELECT
                d.feature_id,
                g.geometry,
                d.z,
                d.x,
                d.y,
                d.geom_type,
                d.source_layer,
                d.class_name,
                d.shape_name,
                d.layer_name,
                d.canonical_object_type,
                d.object_family,
                d.detection_class,
                d.capability,
                d.detection_key
            FROM navsea_detection_objects d
            JOIN navsea_feature_geometry g
                ON d.feature_id = g.feature_id
            """
        )


def main() -> None:
    GeometryEnablement(DbConfig()).run()


if __name__ == "__main__":
    main()