postgres-seeder/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
)

func main() {
	var (
		pairs     = make(map[string]([]string))
		databases = make(map[string]interface{})
		l         = log.New(os.Stderr, "postgres-seeder: ", log.Lmsgprefix|log.LUTC)
		err       error

		hbastr, initstr string
		hba, init       *os.File
		hbabuf, initbuf strings.Builder
	)

	flag.StringVar(&hbastr, "hba", "pg_hba.conf", "Host-Based authentication file output")
	flag.StringVar(&initstr, "init", "init.sql", "SQL command list file output")
	flag.Parse()

	if hba, err = os.Create(hbastr); err != nil {
		l.Fatalf("Cannot create or truncate HBA file %s: %s\n", hbastr, err)
	}
	defer hba.Close()

	if init, err = os.Create(initstr); err != nil {
		l.Fatalf("Cannot create or truncate init file %s: %s\n", initstr, err)
	}
	defer init.Close()

	for _, v := range flag.Args() {
		tok := strings.Split(v, "/")
		if len(tok) != 2 {
			l.Fatalln("Incorrect token count")
		}
		var (
			user = tok[0]
			db   = tok[1]
			v    []string
		)

		// Groups
		if user != db {
			v, _ = pairs[user]
			pairs[user] = append(v, db)
		}

		databases[db] = nil
	}

	hbabuf.WriteString("local\tall\tpostgres\ttrust\n")

	for d := range databases {
		initbuf.WriteString(fmt.Sprintf("CREATE ROLE %s LOGIN;\n", d))
		initbuf.WriteString(fmt.Sprintf("CREATE DATABASE %s OWNER %s;\n", d, d))
		hbabuf.WriteString(fmt.Sprintf("hostssl\t%s\t%s\tall\tcert\tclientcert=verify-full\n", d, d))
	}

	for u, v := range pairs {
		initbuf.WriteString(fmt.Sprintf("CREATE ROLE %s LOGIN;\n", u))
		for _, d := range v {
			initbuf.WriteString(fmt.Sprintf("GRANT %s TO %s;\n", d, u))
			hbabuf.WriteString(fmt.Sprintf("hostssl\t%s\t%s\tall\tcert\tclientcert=verify-full\n", u, d))
		}
	}

	if _, err = hba.WriteString(hbabuf.String()); err != nil {
		l.Fatalf("Failed to write to HBA: %s\n", err)
	}
	if _, err = init.WriteString(initbuf.String()); err != nil {
		l.Fatalf("Failed to write to init: %s\n", err)
	}
}