es/config.go

package main

import (
	"fmt"
	elastic "github.com/olivere/elastic/v7"
	toml "github.com/pelletier/go-toml"
	"io/ioutil"
	"os"
)

type ElasticConfig struct {
	Elastic ElasticSettings
	Servers map[string]ElasticServer
}

type ElasticSettings struct {
	DefaultServer string
}

type ElasticServer struct {
	URL  string
	User string
	Pass string
}

func (ec *ElasticConfig) GetClient(name string) *elastic.Client {
	if name == "" {
		if ec.Elastic.DefaultServer == "" {
			fmt.Fprintf(os.Stderr, "no server name provided\n")
			os.Exit(1)
		}
		name = ec.Elastic.DefaultServer
	}

	if _, ok := ec.Servers[name]; !ok {
		fmt.Fprintf(os.Stderr, "no configuration found for server \"%s\"\n", name)
		os.Exit(1)
	}

	server := ec.Servers[name]

	opts := []elastic.ClientOptionFunc{
		elastic.SetURL(server.URL),
		elastic.SetSniff(false),
		elastic.SetGzip(true),
	}
	if server.User != "" && server.Pass != "" {
		opts = append(opts, elastic.SetBasicAuth(
			server.User, server.Pass,
		))
	}

	es, err := elastic.NewClient(opts...)
	if err != nil {
		fmt.Fprintf(os.Stderr, "elastic connection error: %s\n", err.Error())
		os.Exit(1)
	}

	return es
}

func LoadConfig(cfgpath string) *ElasticConfig {
	cfg := ElasticConfig{}
	if cfgpath != "" {
		cfgfile, err := ioutil.ReadFile(cfgpath)
		if err != nil {
			fmt.Fprintf(os.Stderr, "read error: %s\n", err.Error())
			os.Exit(1)
		}

		err = toml.Unmarshal(cfgfile, &cfg)
		if err != nil {
			fmt.Fprintf(os.Stderr, "parse error: %s\n", err.Error())
			os.Exit(1)
		}

		return &cfg
	}

	paths := []string{}
	homedir, _ := os.UserHomeDir()
	if homedir != "" {
		paths = append(paths, homedir+"/.es.toml")
	}
	paths = append(paths, "/etc/es.toml")

	for _, path := range paths {
		cfgfile, err := ioutil.ReadFile(path)
		if err != nil {
			continue
		}

		err = toml.Unmarshal(cfgfile, &cfg)
		if err != nil {
			continue
		}

		return &cfg
	}

	fmt.Fprintf(os.Stderr, "no valid configuration found in:\n")
	for _, path := range paths {
		fmt.Fprintf(os.Stderr, "    %s\n", path)
	}
	os.Exit(1)
	return nil
}