view cmd/bottlenecks/main.go @ 589:26cdc43ff977

bottlenecks: Add a -dump flag to dump INSERT statements instead of writing directly to the database. Useful for a static setup.
author Sascha L. Teichmann <teichmann@intevation.de>
date Fri, 07 Sep 2018 12:01:47 +0200
parents d5626dd370a4
children d013fa5554cc
line wrap: on
line source

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"regexp"
	"strconv"
	"strings"

	"github.com/jackc/pgx"
	"github.com/jackc/pgx/stdlib"

	"gemma.intevation.de/gemma/pkg/soap/ifbn"
)

const insertSQL = `INSERT INTO waterway.bottlenecks (
  bottleneck_id,
  fk_g_fid,
  objnam,
  nobjnm,
  stretch,
  area,
  rb,
  lb,
  responsible_country,
  revisiting_time,
  limiting,
  date_info,
  source_organization
) VALUES(
  $1,
  isrs_fromText($2),
  $3,
  $4,
  isrsrange(isrs_fromText($5), isrs_fromText($6)),
  ST_MakePolygon(ST_ExteriorRing(ST_Buffer(ST_SetSRID(ST_Makepoint(13.05501, 47.80949), 4326), 0.01)))::Geography,
  $7,
  $8,
  $9,
  $10,
  $11,
  $12,
  $13
) ON CONFLICT (bottleneck_id) DO NOTHING`

const insertDumpSQL = `INSERT INTO waterway.bottlenecks (
  bottleneck_id,
  fk_g_fid,
  objnam,
  nobjnm,
  stretch,
  area,
  rb,
  lb,
  responsible_country,
  revisiting_time,
  limiting,
  date_info,
  source_organization
) VALUES(
  %s,
  isrs_fromText(%s),
  %s,
  %s,
  isrsrange(isrs_fromText(%s), isrs_fromText(%s)),
  ST_MakePolygon(ST_ExteriorRing(ST_Buffer(ST_SetSRID(ST_Makepoint(13.05501, 47.80949), 4326), 0.01)))::Geography,
  %s,
  %s,
  %s,
  %d,
  %s,
  %s::timestamp with time zone,
  %s
) ON CONFLICT (bottleneck_id) DO NOTHING;
`

var (
	url  = flag.String("url", "", "the IFBN service")
	dump = flag.Bool("dump", false, "dump SQL insert statements")

	insecure   = flag.Bool("insecure", false, "skip SSL verification")
	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)
}

var rblbRe = regexp.MustCompile(`(..)_(..)`)

func splitRBLB(s string) (string, string) {
	m := rblbRe.FindStringSubmatch(s)
	if len(m) == 0 {
		return "", ""
	}
	return m[1], m[2]
}

func revisitingTime(s string) int {
	v, err := strconv.Atoi(s)
	if err != nil {
		v = 0
	}
	return v
}

func quote(s string) string {
	return "'" + strings.Replace(s, "'", "'''", -1) + "'"
}

func dumpSQLStatements(bns []*ifbn.BottleNeckType) error {

	fmt.Println("BEGIN;")

	for i := range bns {
		bn := bns[i]
		rb, lb := splitRBLB(bn.Rb_lb)

		var limiting, country string

		if bn.Limiting_factor != nil {
			limiting = string(*bn.Limiting_factor)
		}

		if bn.Responsible_country != nil {
			country = string(*bn.Responsible_country)
		}

		if _, err := fmt.Printf(insertDumpSQL,
			quote(bn.Bottleneck_id),
			quote(bn.Fk_g_fid),
			quote(bn.OBJNAM),
			quote(bn.NOBJNM),
			quote(bn.From_ISRS), quote(bn.To_ISRS),
			quote(rb),
			quote(lb),
			quote(country),
			revisitingTime(bn.Revisiting_time),
			quote(limiting),
			quote(bn.Date_Info.Format("2006-01-02 15:04:05.999 MST")),
			quote(bn.Source),
		); err != nil {
			return err
		}
	}
	_, err := fmt.Println("END;")
	return err
}

func storeInDatabase(bns []*ifbn.BottleNeckType) error {
	return run(func(db *sql.DB) error {

		stmt, err := db.Prepare(insertSQL)
		if err != nil {
			return err
		}
		defer stmt.Close()

		tx, err := db.Begin()
		if err != nil {
			return err
		}

		st := tx.Stmt(stmt)

		for i := range bns {
			bn := bns[i]
			rb, lb := splitRBLB(bn.Rb_lb)

			var limiting, country string

			if bn.Limiting_factor != nil {
				limiting = string(*bn.Limiting_factor)
			}

			if bn.Responsible_country != nil {
				country = string(*bn.Responsible_country)
			}

			fmt.Printf("%s '%s' %s %s\n", bn.Fk_g_fid, bn.OBJNAM, bn.From_ISRS, bn.To_ISRS)

			if _, err := st.Exec(
				bn.Bottleneck_id,
				bn.Fk_g_fid,
				bn.OBJNAM,
				bn.NOBJNM,
				bn.From_ISRS, bn.To_ISRS,
				rb,
				lb,
				country,
				revisitingTime(bn.Revisiting_time),
				limiting,
				bn.Date_Info,
				bn.Source,
			); err != nil {
				tx.Rollback()
				return err
			}
		}

		return tx.Commit()
	})
}

func main() {

	flag.Parse()

	client := ifbn.NewIBottleneckService(*url, *insecure, nil)

	req := &ifbn.Export_bn_by_isrs{}

	resp, err := client.Export_bn_by_isrs(req)
	if err != nil {
		log.Fatalf("error: %v\n", err)
	}

	if resp.Export_bn_by_isrsResult == nil {
		return
	}
	bns := resp.Export_bn_by_isrsResult.BottleNeckType

	if *dump {
		err = dumpSQLStatements(bns)
	} else {
		err = storeInDatabase(bns)
	}

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