view cmd/tin2octree/main.go @ 720:aeaa2adf5a8b octree

octree: write SQL insert statements for octree indices.
author Sascha L. Teichmann <sascha.teichmann@intevation.de>
date Fri, 21 Sep 2018 17:36:00 +0200
parents 5fa2eaffac84
children a4fe07a21ba7
line wrap: on
line source

package main

import (
	"bytes"
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"time"

	"github.com/golang/snappy"
	"github.com/jackc/pgx"
	"github.com/jackc/pgx/stdlib"
)

var (
	bottleneck = flag.String("bottleneck", "", "bottleneck id")
	date       = flag.String("date", "", "date info")
	file       = flag.String("file", "", "save to file")
	insert     = flag.Bool("insert", false, "write as SQL insert statement")
	utm        = flag.Bool("utm", true, "fetch in matchin UTM zone")

	dbhost     = flag.String("dbhost", "localhost", "database host")
	dbport     = flag.Uint("dbport", 5432, "database port")
	dbname     = flag.String("dbname", "gemma", "database user")
	dbuser     = flag.String("dbuser", "scott", "database user")
	dbpassword = flag.String("dbpw", "tiger", "database password")
	dbssl      = flag.String("dbssl", "prefer", "database SSL mode")
)

func run(fn func(*sql.DB) error) error {

	// To ease SSL config ride a bit on parsing.
	cc, err := pgx.ParseConnectionString("sslmode=" + *dbssl)
	if err != nil {
		return err
	}

	// Do the rest manually to allow whitespace in user/password.
	cc.Host = *dbhost
	cc.Port = uint16(*dbport)
	cc.User = *dbuser
	cc.Password = *dbpassword
	cc.Database = *dbname

	db := stdlib.OpenDB(cc)
	defer db.Close()

	return fn(db)
}

const (
	tinSQL = `
SELECT ST_AsBinary(ST_DelaunayTriangles(point_cloud::geometry, 0, 2))
FROM waterway.sounding_results
WHERE bottleneck_id = $1 AND date_info = $2
`
	tinUTMSQL = `
SELECT ST_AsBinary(
  ST_DelaunayTriangles(
	  ST_Transform(point_cloud::geometry, $3::int), 0, 2))
FROM waterway.sounding_results
WHERE bottleneck_id = $1 AND date_info = $2
`
	centroidSQL = `
SELECT ST_X(ST_Centroid(point_cloud::geometry)), ST_Y(ST_Centroid(point_cloud::geometry))
FROM waterway.sounding_results
WHERE bottleneck_id = $1 AND date_info = $2
`
)

func utmZone(x, y float64) int {
	var pref int
	if y > 0 {
		pref = 32600
	} else {
		pref = 32700
	}
	zone := int(math.Floor((x+180)/6)) + 1
	return zone + pref
}

func main() {
	flag.Parse()

	if *bottleneck == "" || *date == "" {
		log.Fatalln("missing bottleneck or date option.")
	}

	dateInfo, err := time.Parse("2006-01-02", *date)
	if err != nil {
		log.Fatalf("error: %v\n", err)
	}

	var t tin

	if err := run(func(db *sql.DB) error {
		var utmZ int

		if *utm {
			var cx, cy float64
			err := db.QueryRow(centroidSQL, *bottleneck, dateInfo).Scan(&cx, &cy)
			switch {
			case err == sql.ErrNoRows:
				return nil
			case err != nil:
				return err
			}
			log.Printf("lat/lon: [%f, %f]\n", cx, cy)
			utmZ = utmZone(cx, cy)
			log.Printf("UTM zone: %d\n", utmZ)
		}

		start := time.Now()
		var err error
		if *utm {
			err = db.QueryRow(tinUTMSQL, *bottleneck, dateInfo, utmZ).Scan(&t)
		} else {
			err = db.QueryRow(tinSQL, *bottleneck, dateInfo).Scan(&t)
		}
		switch {
		case err == sql.ErrNoRows:
			return nil
		case err != nil:
			return err
		}
		log.Printf("query took: %s\n", time.Since(start))

		if *utm {
			t.epsg = uint32(utmZ)
		}

		return nil
	}); err != nil {
		log.Fatalf("error: %v\n", err)
	}

	tb := &treeBuilder{t: &t}
	tb.build()

	if *insert {
		var w io.Writer
		var f *os.File

		if *file != "" {
			if f, err = os.Create(*file); err != nil {
				log.Fatalf("error: %v\n", err)
			}
			w = f
		} else {
			w = os.Stdout
		}

		var buf bytes.Buffer
		if err := write(&buf, tb); err != nil {
			log.Fatalf("error: %v\n", err)
		}
		data := buf.String()
		h := sha1.New()
		buf.WriteTo(h)
		fmt.Fprintln(w, "BEGIN;")
		fmt.Fprintln(w, "INSERT INTO waterway.octrees")
		fmt.Fprintf(w, "SELECT sr.id, '%x',\n", h.Sum(nil))
		fmt.Fprint(w, "decode('")
		fmt.Fprintf(w, "%s", base64.StdEncoding.EncodeToString([]byte(data)))
		fmt.Fprintln(w, "', 'base64')")
		fmt.Fprintln(w, "FROM waterway.sounding_results sr")
		fmt.Fprintf(w,
			"WHERE sr.bottleneck_id = '%s' AND sr.date_info = '%s'::date;\n",
			*bottleneck, *date)
		fmt.Fprintln(w, "END;")

		if f != nil {
			if err := f.Close(); err != nil {
				log.Fatalf("error: %v\n", err)
			}
		}
		return
	}

	if *file != "" {
		f, err := os.Create(*file)
		if err != nil {
			log.Printf("error: %v\n", err)
		}
		err = write(f, tb)
		if err2 := f.Close(); err == nil {
			if err != nil {
				err = err2
			}
		}
		if err != nil {
			log.Fatalf("error: %v\n", err)
		}
	}
}

func write(w io.Writer, tb *treeBuilder) error {
	out := snappy.NewBufferedWriter(w)
	if err := tb.t.Serialize(out); err != nil {
		return err
	}
	if err := tb.Serialize(out); err != nil {
		return err
	}
	return out.Flush()
}