args.go

package main

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/alecthomas/kingpin"
	"github.com/xo/usql/text"
)

// CommandOrFile is a special type to deal with interspersed -c, -f,
// command-line options, to ensure proper order execution.
type CommandOrFile struct {
	Command bool
	Value   string
}

// Args are the command line arguments.
type Args struct {
	DSN               string
	CommandOrFiles    []CommandOrFile
	Out               string
	ForcePassword     bool
	NoPassword        bool
	NoRC              bool
	SingleTransaction bool
	Variables         []string
	PVariables        []string
}

func (args *Args) Next() (string, bool, error) {
	if len(args.CommandOrFiles) == 0 {
		return "", false, io.EOF
	}
	cmd := args.CommandOrFiles[0]
	args.CommandOrFiles = args.CommandOrFiles[1:]
	return cmd.Value, cmd.Command, nil
}

type commandOrFile struct {
	args    *Args
	command bool
}

func (c commandOrFile) Set(value string) error {
	c.args.CommandOrFiles = append(c.args.CommandOrFiles, CommandOrFile{
		Command: c.command,
		Value:   value,
	})
	return nil
}

func (c commandOrFile) String() string {
	return ""
}

func (c commandOrFile) IsCumulative() bool {
	return true
}

// for populating args.PVariables with user-specified options
type pset struct {
	args *Args
	vals []string
}

func (p pset) Set(value string) error {
	for i, v := range p.vals {
		if strings.ContainsRune(v, '%') {
			p.vals[i] = fmt.Sprintf(v, value)
		}
	}
	p.args.PVariables = append(p.args.PVariables, p.vals...)
	return nil
}

func (p pset) String() string {
	return ""
}

func (p pset) IsCumulative() bool {
	return true
}

func NewArgs() *Args {
	args := &Args{}
	// set usage template
	kingpin.UsageTemplate(text.UsageTemplate())
	kingpin.Arg("dsn", "database url").StringVar(&args.DSN)
	// command / file flags
	kingpin.Flag("command", "run only single command (SQL or internal) and exit").Short('c').SetValue(commandOrFile{args, true})
	kingpin.Flag("file", "execute commands from file and exit").Short('f').SetValue(commandOrFile{args, false})
	// general flags
	kingpin.Flag("no-password", "never prompt for password").Short('w').BoolVar(&args.NoPassword)
	kingpin.Flag("no-rc", "do not read start up file").Short('X').BoolVar(&args.NoRC)
	kingpin.Flag("out", "output file").Short('o').StringVar(&args.Out)
	kingpin.Flag("password", "force password prompt (should happen automatically)").Short('W').BoolVar(&args.ForcePassword)
	kingpin.Flag("single-transaction", "execute as a single transaction (if non-interactive)").Short('1').BoolVar(&args.SingleTransaction)
	kingpin.Flag("set", "set variable NAME to VALUE").Short('v').PlaceHolder(", --variable=NAME=VALUE").StringsVar(&args.Variables)
	// pset
	kingpin.Flag("pset", `set printing option VAR to ARG (see \pset command)`).Short('P').PlaceHolder("VAR[=ARG]").StringsVar(&args.PVariables)
	// pset flags
	kingpin.Flag("field-separator", `field separator for unaligned and CSV output (default "|" and ",")`).Short('F').SetValue(pset{args, []string{"fieldsep=%q", "csv_fieldsep=%q"}})
	kingpin.Flag("record-separator", `record separator for unaligned and CSV output (default \n)`).Short('R').SetValue(pset{args, []string{"recordsep=%q"}})
	kingpin.Flag("table-attr", "set HTML table tag attributes (e.g., width, border)").Short('T').SetValue(pset{args, []string{"tableattr=%q"}})
	type psetconfig struct {
		long  string
		short rune
		help  string
		vals  []string
	}
	pc := func(long string, r rune, help string, vals ...string) psetconfig {
		return psetconfig{long, r, help, vals}
	}
	for _, c := range []psetconfig{
		pc("no-align", 'A', "unaligned table output mode", "format=unaligned"),
		pc("html", 'H', "HTML table output mode", "format=html"),
		pc("tuples-only", 't', "print rows only", "tuples_only=on"),
		pc("expanded", 'x', "turn on expanded table output", "expanded=on"),
		pc("field-separator-zero", 'z', "set field separator for unaligned and CSV output to zero byte", "fieldsep_zero=on"),
		pc("record-separator-zero", '0', "set record separator for unaligned and CSV output to zero byte", "recordsep_zero=on"),
		pc("json", 'J', "JSON output mode", "format=json"),
		pc("csv", 'C', "CSV output mode", "format=csv"),
		pc("vertical", 'G', "vertical output mode", "format=vertical"),
	} {
		// make copy of values for the callback closure (see https://stackoverflow.com/q/26692844)
		vals := make([]string, len(c.vals))
		copy(vals, c.vals)
		kingpin.Flag(c.long, c.help).Short(c.short).PlaceHolder("TEXT").PreAction(func(*kingpin.ParseContext) error {
			args.PVariables = append(args.PVariables, vals...)
			return nil
		}).Bool()
	}
	kingpin.Flag("quiet", "run quietly (no messages, only query output)").Short('q').PreAction(func(*kingpin.ParseContext) error {
		args.Variables = append(args.Variables, "QUIET=on")
		return nil
	}).Bool()
	// add --set as a hidden alias for --variable
	kingpin.Flag("variable", "set variable NAME to VALUE").Hidden().StringsVar(&args.Variables)
	// add --version flag
	kingpin.Flag("version", "display version and exit").PreAction(func(*kingpin.ParseContext) error {
		fmt.Fprintln(os.Stdout, text.CommandName, text.CommandVersion)
		os.Exit(0)
		return nil
	}).Bool()
	// hide help flag
	kingpin.HelpFlag.Short('h').Hidden()
	// parse
	kingpin.Parse()
	return args
}