MapComplete/scripts/osm2pgsql/tilecountServer.ts

import { Client } from "pg"
import { Tiles } from "../../src/Models/TileRange"
import { Server } from "../server"

/**
 * Just the OSM2PGSL default database
 */
interface PoiDatabaseMeta {
    attributes
    current_timestamp
    db_format
    flat_node_file
    import_timestamp
    output
    prefix
    replication_base_url
    replication_sequence_number
    replication_timestamp
    style
    updatable
    version
}

/**
 * Connects with a Postgis database, gives back how much items there are within the given BBOX
 */
class OsmPoiDatabase {
    private static readonly prefixes: ReadonlyArray<string> = ["pois", "lines", "polygons"]
    private readonly _client: Client
    private isConnected = false
    private supportedLayers: Set<string> = undefined
    private supportedLayersDate: Date = undefined
    private metaCache: PoiDatabaseMeta = undefined
    private metaCacheDate: Date = undefined

    constructor(connectionString: string) {
        this._client = new Client(connectionString)
    }

    async getCount(
        layer: string,
        bbox: [[number, number], [number, number]] = undefined
    ): Promise<{ count: number; lat: number; lon: number }> {
        if (!this.isConnected) {
            await this._client.connect()
            this.isConnected = true
        }

        let total: number = 0
        let latSum = 0
        let lonSum = 0
        for (const prefix of OsmPoiDatabase.prefixes) {
            let query =
                "SELECT COUNT(*), ST_AsText(ST_Centroid(ST_Collect(geom))) FROM " +
                prefix +
                "_" +
                layer

            if (bbox) {
                query += ` WHERE ST_MakeEnvelope (${bbox[0][0]}, ${bbox[0][1]}, ${bbox[1][0]}, ${bbox[1][1]}, 4326) ~ geom`
            }
            const result = await this._client.query(query)
            const count = Number(result.rows[0].count)
            let point = result.rows[0].st_astext
            if (count === 0) {
                continue
            }
            total += count
            if (!point) {
                continue
            }
            point = point.substring(6, point.length - 1)
            const [lon, lat] = point.split(" ")
            latSum += lat * count
            lonSum += lon * count
        }

        return { count: total, lat: latSum / total, lon: lonSum / total }
    }

    disconnect() {
        this._client.end()
    }

    async getLayers(): Promise<Set<string>> {
        if (
            this.supportedLayers !== undefined &&
            new Date().getTime() - this.supportedLayersDate.getTime() < 1000 * 60 * 60 * 24
        ) {
            return this.supportedLayers
        }
        const q =
            "SELECT table_name \n" +
            "FROM information_schema.tables \n" +
            "WHERE table_schema = 'public' AND table_name LIKE 'lines_%';"
        const result = await this._client.query(q)
        const layers = result.rows.map((r) => r.table_name.substring("lines_".length))
        this.supportedLayers = new Set(layers)
        this.supportedLayersDate = new Date()
        return this.supportedLayers
    }

    async getMeta(): Promise<PoiDatabaseMeta> {
        const now = new Date()
        if (this.metaCache !== undefined) {
            const diffSec = (this.metaCacheDate.getTime() - now.getTime()) / 1000
            if (diffSec < 120) {
                return this.metaCache
            }
        }
        const result = await this._client.query("SELECT * FROM public.osm2pgsql_properties")
        const meta = {}
        for (const { property, value } of result.rows) {
            meta[property] = value
        }
        this.metaCacheDate = now
        this.metaCache = <any>meta
        return this.metaCache
    }
}

class CachedSqlCount {
    private readonly _cache: Record<
        string,
        Record<
            number,
            {
                date: Date
                entry: { count: number; lat: number; lon: number }
            }
        >
    > = {}

    private readonly _poiDatabase: OsmPoiDatabase
    private readonly _maxAge: number

    constructor(poiDatabase: OsmPoiDatabase, maxAge: number) {
        this._poiDatabase = poiDatabase
        this._maxAge = maxAge
    }

    public async getCount(
        layer: string,
        tileId: number
    ): Promise<{ count: number; lat: number; lon: number }> {
        const cachedEntry = this._cache[layer]?.[tileId]
        if (cachedEntry) {
            const age = (new Date().getTime() - cachedEntry.date.getTime()) / 1000
            if (age < this._maxAge) {
                return cachedEntry.entry
            }
        }
        const bbox = Tiles.tile_bounds_lon_lat(...Tiles.tile_from_index(tileId))
        const count = await this._poiDatabase.getCount(layer, bbox)
        if (!this._cache[layer]) {
            this._cache[layer] = {}
        }
        this._cache[layer][tileId] = { entry: count, date: new Date() }
        return count
    }
}

const connectionString = "postgresql://user:password@localhost:5444/osm-poi"
const tcs = new OsmPoiDatabase(connectionString)
const withCache = new CachedSqlCount(tcs, 60 * 60 * 24)
new Server(2345, { ignorePathPrefix: ["summary"] }, [
    {
        mustMatch: "status.json",
        mimetype: "application/json",
        handle: async (path: string) => {
            const layers = await tcs.getLayers()
            const meta = await tcs.getMeta()
            return JSON.stringify({ meta, layers: Array.from(layers) })
        },
    },
    {
        mustMatch: /[a-zA-Z0-9+_-]+\/[0-9]+\/[0-9]+\/[0-9]+\.json/,
        mimetype: "application/json", // "application/vnd.geo+json",
        async handle(path) {
            const [layers, z, x, y] = path.split(".")[0].split("/")

            let sum = 0
            let properties: Record<string, number> = {}
            const availableLayers = await tcs.getLayers()
            let latSum = 0
            let lonSum = 0
            for (const layer of layers.split("+")) {
                if (!availableLayers.has(layer)) {
                    continue
                }
                const count = await withCache.getCount(
                    layer,
                    Tiles.tile_index(Number(z), Number(x), Number(y))
                )

                properties[layer] = count.count
                if (count.count !== 0) {
                    latSum += count.lat * count.count
                    lonSum += count.lon * count.count
                    sum += count.count
                }
            }

            properties["lon"] = lonSum / sum
            properties["lat"] = latSum / sum

            return JSON.stringify({ ...properties, total: sum })
        },
    },
])
console.log(
    ">>>",
    await tcs.getCount("drinking_water", [
        [3.194358020772171, 51.228073636083394],
        [3.2839964396059145, 51.172701162680994],
    ])
)