Julien Riou
c43270a1d4
- Add debug messages - Use mutex to compute statistics - Add UUID to identify database connections - Fix connection handling for MySQL Signed-off-by: Julien Riou <julien@riou.xyz>
111 lines
2.7 KiB
Go
111 lines
2.7 KiB
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()
|
|
}
|