sq/cli/options.go

package cli

import (
	"github.com/neilotoole/sq/drivers"
	"github.com/neilotoole/sq/drivers/csv"
	"github.com/neilotoole/sq/libsq/core/errz"
	"github.com/neilotoole/sq/libsq/core/options"
	"github.com/neilotoole/sq/libsq/driver"
	"github.com/neilotoole/sq/libsq/source"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// getFlagOptions builds options.Options from flags. In effect, a flag
// such as --ingest.header is mapped to an option.Opt of the same name.
//
// See also: getCmdOptions, applySourceOptions, applyCollectionOptions.
func getFlagOptions(flags *pflag.FlagSet, reg *options.Registry) (options.Options, error) {
	o := options.Options{}
	err := reg.Visit(func(opt options.Opt) error {
		key := opt.Key()
		f := flags.Lookup(key)
		if f == nil {
			return nil
		}

		if !f.Changed {
			return nil
		}

		if f.Value == nil {
			// This shouldn't happen
			return nil
		}

		o[key] = f.Value.String()
		return nil
	})
	if err != nil {
		return nil, err
	}

	o, err = reg.Process(o)
	if err != nil {
		return nil, errz.Wrap(err, "options from flags")
	}

	return o, nil
}

// getCmdOptions returns the options.Options generated by merging
// config options and flag options.
//
// See also: getFlagOptions, applySourceOptions, applyCollectionOptions.
func getCmdOptions(cmd *cobra.Command) (options.Options, error) {
	rc := RunContextFrom(cmd.Context())
	var configOpts options.Options
	if rc.Config != nil && rc.Config.Options != nil {
		configOpts = rc.Config.Options
	} else {
		configOpts = options.Options{}
	}

	flagOpts, err := getFlagOptions(cmd.Flags(), rc.OptionsRegistry)
	if err != nil {
		return nil, err
	}

	return options.Merge(configOpts, flagOpts), nil
}

// applySourceOptions merges options from config, src, and flags.
// The src.Options field may be replaced or mutated. It will always
// be non-nil (unless an error is returned).
//
// See also: getFlagOptions, getCmdOptions, applyCollectionOptions.
func applySourceOptions(cmd *cobra.Command, src *source.Source) error {
	rc := RunContextFrom(cmd.Context())

	defaultOpts := rc.Config.Options
	if defaultOpts == nil {
		defaultOpts = options.Options{}
	}

	flagOpts, err := getFlagOptions(cmd.Flags(), rc.OptionsRegistry)
	if err != nil {
		return err
	}

	srcOpts := src.Options
	if srcOpts == nil {
		srcOpts = options.Options{}
	}

	effectiveOpts := options.Merge(defaultOpts, srcOpts, flagOpts)
	src.Options = effectiveOpts
	return nil
}

// applyCollectionOptions invokes applySourceOptions for
// each source in coll. The sources may have their Source.Options field
// mutated.
//
// See also: getCmdOptions, getFlagOptions, applySourceOptions.
func applyCollectionOptions(cmd *cobra.Command, coll *source.Collection) error {
	return coll.Visit(func(src *source.Source) error {
		return applySourceOptions(cmd, src)
	})
}

// RegisterDefaultOpts registers the options.Opt instances
// that the CLI knows about. It panics if an opt is added twice.
func RegisterDefaultOpts(reg *options.Registry) {
	reg.Add(driver.OptConnMaxOpen)
	reg.Add(driver.OptConnMaxIdle)
	reg.Add(driver.OptConnMaxIdleTime)
	reg.Add(driver.OptConnMaxLifetime)
	reg.Add(drivers.OptIngestHeader)
	reg.Add(csv.OptDelim)
	reg.Add(csv.OptEmptyAsNull)
	reg.Add(OptOutputFormat)
	reg.Add(OptPrintHeader)
	reg.Add(OptPingTimeout)
	reg.Add(OptShellCompletionTimeout)
}