view pkg/controllers/shapestretches.go @ 5520:05db984d3db1

Improve performance of bottleneck area calculation Avoid buffer calculations by replacing them with simple distance comparisons and calculate the boundary of the result geometry only once per iteration. In some edge cases with very large numbers of iterations, this reduced the runtime of a bottleneck import by a factor of more than twenty.
author Tom Gottfried <tom@intevation.de>
date Thu, 21 Oct 2021 19:50:39 +0200
parents 5f47eeea988d
children 31973f6f5cca
line wrap: on
line source

// This is Free Software under GNU Affero General Public License v >= 3.0
// without warranty, see README.md and license for details.
//
// SPDX-License-Identifier: AGPL-3.0-or-later
// License-Filename: LICENSES/AGPL-3.0.txt
//
// Copyright (C) 2019 by via donau
//   – Österreichische Wasserstraßen-Gesellschaft mbH
// Software engineering by Intevation GmbH
//
// Author(s):
//  * Sascha L. Teichmann <sascha.teichmann@intevation.de>

package controllers

import (
	"archive/zip"
	"database/sql"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/gorilla/mux"
	"github.com/jonas-p/go-shp"

	"gemma.intevation.de/gemma/pkg/config"
	"gemma.intevation.de/gemma/pkg/log"
	"gemma.intevation.de/gemma/pkg/middleware"
	"gemma.intevation.de/gemma/pkg/wkb"
)

const (
	selectStretchSQL = `
SELECT
  id,
  isrs_asText(lower(stretch)),
  isrs_asText(upper(stretch)),
  ST_AsBinary(area::geometry),
  objnam,
  nobjnam,
  date_info,
  source_organization
FROM users.stretches WHERE
  staging_done AND
  name = $1`

	selectStretchCountriesSQL = `
SELECT country FROM users.stretch_countries
WHERE stretch_id = $1
ORDER BY country`
)

func flattenPoints(mp wkb.MultiPolygonGeom) []shp.Point {

	var n int
	for _, p := range mp {
		for _, r := range p {
			n += len(r)
		}
	}

	ps := make([]shp.Point, n)
	var i int
	for _, p := range mp {
		for _, r := range p {
			for _, v := range r {
				ps[i].X = v.X
				ps[i].Y = v.Y
				i++
			}
		}
	}

	return ps
}

func asShpPolygon(mp wkb.MultiPolygonGeom) *shp.Polygon {

	points := flattenPoints(mp)

	var numParts int
	for _, pl := range mp {
		numParts += len(pl)
	}

	p := &shp.Polygon{}
	p.NumParts = int32(numParts)
	p.NumPoints = int32(len(points))
	p.Parts = make([]int32, numParts)

	var marker int32
	var i int
	for _, pl := range mp {
		for _, r := range pl {
			p.Parts[i] = marker
			i++
			marker += int32(len(r))
		}
	}

	p.Points = points
	p.Box = shp.BBoxFromPoints(points)
	return p
}

func stretchShapeDownload(rw http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	name := vars["name"]

	conn := middleware.GetDBConn(req)
	ctx := req.Context()

	tx, err := conn.BeginTx(ctx, &sql.TxOptions{ReadOnly: true})
	if err != nil {
		http.Error(
			rw, fmt.Sprintf("DB error: %v.", err),
			http.StatusInternalServerError)
		return
	}
	defer tx.Rollback()

	var (
		id           int64
		lower, upper string
		data         []byte
		objnam       string
		nobjnam      sql.NullString
		dateInfo     time.Time
		source       string
	)

	switch err := tx.QueryRowContext(ctx, selectStretchSQL, name).Scan(
		&id,
		&lower, &upper,
		&data,
		&objnam, &nobjnam,
		&dateInfo,
		&source,
	); {
	case err == sql.ErrNoRows:
		http.NotFound(rw, req)
		return
	case err != nil:
		http.Error(
			rw, fmt.Sprintf("DB error: %v.", err),
			http.StatusInternalServerError)
		return
	}

	var countries []string

	if err := func() error {
		rows, err := tx.Query(selectStretchCountriesSQL, id)
		if err != nil {
			return err
		}
		defer rows.Close()

		for rows.Next() {
			var country string
			if err := rows.Scan(&country); err != nil {
				return err
			}
			countries = append(countries, country)
		}

		return rows.Err()
	}(); err != nil {
		http.Error(
			rw, fmt.Sprintf("DB error: %v.", err),
			http.StatusInternalServerError)
		return
	}

	var geom wkb.MultiPolygonGeom

	if err := geom.FromWKB(data); err != nil {
		http.Error(
			rw, fmt.Sprintf("Decoding WKB error: %v.", err),
			http.StatusInternalServerError)
		return
	}

	tmp := config.TmpDir()

	dir, err := ioutil.TempDir(tmp, "stretch-download")
	if err != nil {
		http.Error(
			rw, fmt.Sprintf("Cannot create temp dir: %v.", err),
			http.StatusInternalServerError)
		return
	}
	defer os.RemoveAll(dir)

	shpDir := filepath.Join(dir, "stretch")

	if err := os.Mkdir(shpDir, os.ModePerm); err != nil {
		http.Error(
			rw, fmt.Sprintf("Cannot create temp dir: %v.", err),
			http.StatusInternalServerError)
		return
	}

	filename := filepath.Join(shpDir, "stretch")

	if err := func() error {
		writer, err := shp.Create(filename, shp.POLYGON)
		if err != nil {
			return err
		}
		defer writer.Close()

		fields := []shp.Field{
			shp.StringField("NAME", 50),
			shp.StringField("OBJNAM", 50),
			shp.StringField("NOBJNAM", 50),
			shp.StringField("LOWER", 25),
			shp.StringField("UPPER", 25),
			shp.StringField("SOURCE", 50),
			shp.StringField("DATEINFO", 25),
			shp.StringField("COUNTRIES", 50),
		}

		writer.SetFields(fields)

		p := asShpPolygon(geom)

		writer.Write(p)

		write := func(field int, v interface{}) {
			if err == nil {
				err = writer.WriteAttribute(0, field, v)
			}
		}

		write(0, name)
		write(1, objnam)
		if nobjnam.Valid {
			write(2, nobjnam.String)
		}
		write(3, lower)
		write(4, upper)
		write(5, source)
		write(6, dateInfo.Format(time.RFC3339))
		if len(countries) > 0 {
			write(7, strings.Join(countries, ","))
		}

		return err
	}(); err != nil {
		http.Error(
			rw, fmt.Sprintf("creating shapefile failed: %v.", err),
			http.StatusInternalServerError)
		return
	}

	entries, err := func() ([]os.FileInfo, error) {
		f, err := os.Open(shpDir)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		return f.Readdir(-1)
	}()
	if err != nil {
		http.Error(
			rw, fmt.Sprintf("cannot read directory: %v.", err),
			http.StatusInternalServerError)
		return
	}

	rw.Header().Set("Content-Type", "application/zip")
	rw.Header().Set("Content-Disposition", `attachment; filename="stretch.zip"`)

	zipfile := zip.NewWriter(rw)
	defer zipfile.Close()

	now := time.Now()

	base := filepath.Base(shpDir)
	for _, info := range entries {
		if !info.Mode().IsRegular() {
			continue
		}
		if err := func() error {
			srcFile, err := os.Open(filepath.Join(shpDir, info.Name()))
			if err != nil {
				return err
			}
			defer srcFile.Close()
			header := &zip.FileHeader{
				Name:     base + "/" + info.Name(),
				Method:   zip.Deflate,
				Modified: now,
			}
			out, err := zipfile.CreateHeader(header)
			if err != nil {
				return err
			}
			_, err = io.Copy(out, srcFile)
			return err
		}(); err != nil {
			log.Errorf("cannot write file: %v\n", err)
			return
		}
	}
}