benchito/config.go

package main

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"

	"gopkg.in/yaml.v2"
)

// Config to store all configurations (from command line, from file, etc)
type Config struct {
	Driver         string        `yaml:"driver"`
	Connections    int           `yaml:"connections"`
	Query          string        `yaml:"query"`
	Duration       time.Duration `yaml:"duration"`
	Reconnect      bool          `yaml:"reconnect"`
	DSN            string        `yaml:"dsn"`
	Host           string        `yaml:"host"`
	Port           int           `yaml:"port"`
	User           string        `yaml:"user"`
	Password       string        `yaml:"password"`
	Database       string        `yaml:"database"`
	TLS            string        `yaml:"tls"`
	ConnectTimeout int           `yaml:"connect_timeout"`
}

// NewConfig creates a Config struct
func NewConfig() *Config {
	return &Config{}
}

// Read YaML configuration file from disk
func (c *Config) Read(file string) error {
	file, err := filepath.Abs(file)
	if err != nil {
		return err
	}

	yamlFile, err := ioutil.ReadFile(file)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(yamlFile, &c)
	if err != nil {
		return err
	}

	return nil
}

// ParseDSN detects the database driver then creates the DSN accordingly
func (c *Config) ParseDSN() {
	if c.DSN == "" {
		switch c.Driver {
		case "postgres":
			c.DSN = c.parsePostgresDSN()
		case "mysql":
			c.DSN = c.parseMysqlDSN()
		}
	}
}

func (c *Config) parsePostgresDSN() string {
	var parameters []string
	if c.Host != "" {
		parameters = append(parameters, fmt.Sprintf("host=%s", c.Host))
	}
	if c.Port != 0 {
		parameters = append(parameters, fmt.Sprintf("port=%d", c.Port))
	}
	if c.User != "" {
		parameters = append(parameters, fmt.Sprintf("user=%s", c.User))
	}
	if c.Password != "" {
		parameters = append(parameters, fmt.Sprintf("password=%s", c.Password))
	}
	if c.Database != "" {
		parameters = append(parameters, fmt.Sprintf("database=%s", c.Database))
	}
	if c.ConnectTimeout != 0 {
		parameters = append(parameters, fmt.Sprintf("connect_timeout=%d", c.ConnectTimeout))
	}
	if AppName != "" {
		parameters = append(parameters, fmt.Sprintf("application_name=%s", AppName))
	}
	if c.TLS != "" {
		parameters = append(parameters, fmt.Sprintf("sslmode=%s", c.TLS))
	}
	return strings.Join(parameters, " ")
}

func (c *Config) parseMysqlDSN() (dsn string) {
	config := mysql.NewConfig()
	config.Timeout = 1 * time.Second
	config.Addr = c.Host
	if c.Port != 0 {
		config.Net = "tcp"
		config.Addr += fmt.Sprintf(":%d", c.Port)
	}
	config.User = c.User
	config.Passwd = c.Password
	config.DBName = c.Database
	config.TLSConfig = c.TLS
	return config.FormatDSN()
}