<|start_filename|>config_test.go<|end_filename|>
package main

import (
	"os"
	"reflect"
	"testing"
	"time"
)

func Test_loadConfig(t *testing.T) {
	type args struct {
		file string
	}
	tests := []struct {
		name    string
		args    args
		want    *Config
		wantErr bool
	}{
		{
			name: "defaults-all",
			args: args{file: "test-resources/config-test/defaults-all.yml"},
			want: &Config{Defaults: DefaultsData{
				DataSourceRef:     "my-ds",
				QueryInterval:     time.Minute * 15,
				QueryTimeout:      time.Minute * 5,
				QueryValueOnError: "-1",
			},
			},
			wantErr: false,
		},
		{
			name:    "empty-file",
			args:    args{file: "test-resources/config-test/empty.yml"},
			want:    newConfig(),
			wantErr: false,
		},
		{
			name:    "missing-driver-in-datasource",
			args:    args{file: "test-resources/config-test/datasource-missing-driver.yml"},
			wantErr: true,
		},
		{
			name:    "missing-properties-in-datasource",
			args:    args{file: "test-resources/config-test/datasource-missing-properties.yml"},
			wantErr: true,
		},
		{
			name:    "no-datasource",
			args:    args{file: "test-resources/config-test/datasource-no.yml"},
			want:    newConfig(),
			wantErr: false,
		},
		{
			name: "one-datasource",
			args: args{file: "test-resources/config-test/datasource-one.yml"},
			want: &Config{Defaults: createDefaultsData(),
				DataSources: map[string]DataSource{
					"mysql-test": {
						Driver: "mysql",
						Properties: map[string]interface{}{
							"host":     "localhost",
							"port":     3306,
							"user":     "root",
							"password": "<PASSWORD>",
							"database": "test",
						},
					},
				},
			},
			wantErr: false,
		},
		{
			name: "two-datasource",
			args: args{file: "test-resources/config-test/datasource-two.yml"},
			want: &Config{Defaults: createDefaultsData(),
				DataSources: map[string]DataSource{
					"mysql-test": {
						Driver: "mysql",
						Properties: map[string]interface{}{
							"host":     "localhost",
							"port":     3306,
							"user":     "root",
							"password": "<PASSWORD>",
							"database": "test",
						},
					},
					"mysql-test-2": {
						Driver: "mysql-2",
						Properties: map[string]interface{}{
							"host":     "localhost-2",
							"port":     3307,
							"user":     "root-2",
							"password": "<PASSWORD>",
							"database": "test-2",
						},
					},
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := loadConfig(tt.args.file)
			if (err != nil) != tt.wantErr {
				t.Errorf("loadConfig() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			//spew.Dump(got)
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("loadConfig() = %v, want %v", got, tt.want)
			}
		})
	}

}

func Test_loadQueryConfig(t *testing.T) {
	type args struct {
		queriesFile string
		config      *Config
	}
	c, err := loadConfig("test-resources/config-test/queries-config.yml")
	if err != nil {
		t.Errorf("Failed to load config file: %s", err)
		return
	}
	tests := []struct {
		name    string
		args    args
		want    QueryList
		wantErr bool
	}{
		{
			name: "queries-datasource",
			args: args{
				queriesFile: "test-resources/config-test/queries-datasource.yml",
				config:      c,
			},
			want: []*Query{
				{
					Name:          "query_ds_1",
					DataSourceRef: "my-ds-1",
					Driver:        "mysql",
					Connection: map[string]interface{}{
						"host":     "localhost",
						"port":     3306,
						"user":     "root",
						"password": "<PASSWORD>",
						"database": "test",
					},
					SQL:          "select 1 from dual\n",
					Params:       nil,
					Interval:     time.Second * 10,
					Timeout:      time.Second * 5,
					DataField:    "",
					ValueOnError: "0",
				},
				{
					Name:          "query_ds_2",
					DataSourceRef: "my-ds-2",
					Driver:        "postgresql",
					Connection: map[string]interface{}{
						"host":     "localhost",
						"port":     5432,
						"user":     "postgres",
						"password": "<PASSWORD>",
						"database": "test",
					},
					SQL:          "select 1 from dual\n",
					Params:       nil,
					Interval:     time.Minute * 15,
					Timeout:      time.Minute * 5,
					DataField:    "",
					ValueOnError: "-1",
				},
			},
			wantErr: false,
		},
		{
			name: "queries-submetrics",
			args: args{
				queriesFile: "test-resources/config-test/queries-submetrics.yml",
				config:      c,
			},
			want: []*Query{
				{
					Name:          "query_ds_1",
					SQL:           `select 1 as "sum", 2 as "count" from dual`,
					DataSourceRef: "my-ds-1",
					Driver:        "mysql",
					Connection: map[string]interface{}{
						"host":     "localhost",
						"port":     3306,
						"user":     "root",
						"password": "<PASSWORD>",
						"database": "test",
					},
					Interval: time.Minute * 15,
					Timeout:  time.Minute * 5,
					Params:   nil,
					SubMetrics: map[string]string{
						"count": "count",
						"sum":   "sum",
					},
					ValueOnError: "-1",
					DataField:    "",
				},
			},
			wantErr: false,
		},
		{
			name: "queries-compatibility",
			args: args{
				queriesFile: "test-resources/config-test/queries-compatibility.yml",
				config:      newConfig(),
			},
			want: []*Query{
				{
					Name:   "num_products",
					Driver: "postgresql",
					Connection: map[string]interface{}{
						"host":     "example.org",
						"port":     5432,
						"user":     "postgres",
						"password": "<PASSWORD>",
						"database": "products",
					},
					SQL:          "select 1 from dual\n",
					Params:       nil,
					Interval:     DefaultInterval,
					Timeout:      DefaultTimeout,
					DataField:    "",
					ValueOnError: "",
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := loadQueryConfig(tt.args.queriesFile, tt.args.config)
			if (err != nil) != tt.wantErr {
				t.Errorf("loadQueryConfig() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			// spew.Dump(got)
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("loadQueryConfig() = %v, want %v", got, tt.want)
			}
		})
	}

}

func Test_loadNonExistingFiles(t *testing.T) {
	file := "does-not-exist"
	_, err := loadConfig(file)
	if err == nil {
		t.Errorf("No errors even if config file [%s] does not exist!", file)
		return
	}

	_, err = loadQueryConfig(file, newConfig())
	if err == nil {
		t.Errorf("No errors even if query file [%s] does not exist!", file)
		return
	}

	_, err = loadQueriesInDir(file, newConfig(), false)
	if err == nil {
		t.Errorf("No errors even if query directory [%s] does not exist!", file)
		return
	}
}

func Test_allowBrokenQueryFileInDir(t *testing.T) {
	//config should allow loading queries from a directory that contains invalid files; bad files should not
	//stop good queries from running

	//load a directory having one good and one bad query
	q, err := loadQueriesInDir("test-resources/config-test/one-good-query", newConfig(), true)
	//expect no error and 1 query
	if err != nil {
		t.Fatal(err)
	}
	if len(q) != 1 {
		t.Fatal(len(q))
	}
}

func Test_expandEnvFile(t *testing.T) {
	os.Setenv("DEFAULT_DATA_SOURCE", "mysql")
	os.Setenv("DB_HOST", "localhost")
	os.Setenv("DB_PORT", "")
	os.Setenv("DB_USER", "user")
	os.Setenv("DB_PASS", "<PASSWORD>")
	// DB_NAME not set

	file := "test-resources/config-test/expand-env.yml"
	got, err := loadConfig(file)
	if err != nil {
		t.Fatal(err)
	}

	want := &Config{
		Defaults: createDefaultsData(),
		DataSources: map[string]DataSource{
			"mysql": {
				Driver: "mysql",
				Properties: map[string]interface{}{
					"host":     "localhost",
					"port":     3306,
					"user":     "user",
					"password": "<PASSWORD>",
					"database": nil,
				},
			},
		},
	}
	want.Defaults.DataSourceRef = "mysql"

	if !reflect.DeepEqual(got, want) {
		t.Errorf("loadConfig() = %v, want %v", got, want)
	}
}


<|start_filename|>set.go<|end_filename|>
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"strconv"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
)

type record map[string]interface{}
type records []record

type metricStatus int

const (
	registered metricStatus = iota
	unregistered
)

// QueryResult contains query results
type QueryResult struct {
	Query  *Query
	Result map[string]prometheus.Gauge // Internally we represent each facet with a JSON-encoded string for simplicity
}

// NewQueryResult initializes a new metrics collector.
func NewQueryResult(q *Query) *QueryResult {
	r := &QueryResult{
		Query:  q,
		Result: make(map[string]prometheus.Gauge),
	}

	return r
}

func (r *QueryResult) generateMetricName(suffix string) string {
	metricName := r.Query.Name
	if suffix != "" {
		metricName = fmt.Sprintf("%s_%s", r.Query.Name, suffix)
	}
	return metricName
}

func (r *QueryResult) generateMetricUniqueKey(facets map[string]interface{}, suffix string) string {
	jsonData, _ := json.Marshal(facets)
	return fmt.Sprintf("%s%s", r.generateMetricName(suffix), string(jsonData))
}

func (r *QueryResult) createMetric(facets map[string]interface{}, suffix string) (string, metricStatus) {
	metricName := r.generateMetricName(suffix)
	resultKey := r.generateMetricUniqueKey(facets, suffix)

	labels := prometheus.Labels{}
	for k, v := range facets {
		labels[k] = strings.ToLower(fmt.Sprintf("%v", v))
	}

	if _, ok := r.Result[resultKey]; ok {
		// A metric with this key is already created and assumed to be registered
		return resultKey, registered
	}

	fmt.Println("Creating", resultKey)
	r.Result[resultKey] = prometheus.NewGauge(prometheus.GaugeOpts{
		Name:        fmt.Sprintf("query_result_%s", metricName),
		Help:        "Result of an SQL query",
		ConstLabels: labels,
	})
	return resultKey, unregistered
}

func setValueForResult(r prometheus.Gauge, v interface{}) error {
	switch t := v.(type) {
	case nil:
		r.Set(math.NaN())
	case string:
		f, err := strconv.ParseFloat(t, 64)
		if err != nil {
			return err
		}
		r.Set(f)
	case int:
		r.Set(float64(t))
	case float64:
		r.Set(t)
	default:
		return fmt.Errorf("Unhandled type %s", t)
	}
	return nil
}

// SetMetrics set and register metrics
func (r *QueryResult) SetMetrics(recs records, valueOnError string) error {
	// Queries that return only one record should only have one column
	if len(recs) > 1 && len(recs[0]) == 1 {
		return errors.New("There is more than one row in the query result - with a single column")
	}

	if r.Query.DataField != "" && len(r.Query.SubMetrics) > 0 {
		return errors.New("sub-metrics are not compatible with data-field")
	}

	// We need to make sure not to default to a value on error before
	// it has been registered once before since re-registering might
	// not work if different labels are used.
	if len(recs) == 0 && valueOnError != "" {
		metricSet := false
		for k := range r.Result {
			if strings.HasPrefix(k, r.generateMetricName("")) {
				err := setValueForResult(r.Result[k], valueOnError)
				if err != nil {
					return err
				}
				metricSet = true
			}
		}
		if metricSet {
			return nil
		}
	}

	submetrics := map[string]string{}

	if len(r.Query.SubMetrics) > 0 {
		submetrics = r.Query.SubMetrics
	} else {
		submetrics = map[string]string{"": r.Query.DataField}
	}

	facetsWithResult := make(map[string]metricStatus, 0)
	for _, row := range recs {
		for suffix, datafield := range submetrics {
			facet := make(map[string]interface{})
			var (
				dataVal   interface{}
				dataFound bool
			)
			for k, v := range row {
				if len(row) > 1 && strings.ToLower(k) != datafield { // facet field, add to facets
					submetric := false
					for _, n := range submetrics {
						if strings.ToLower(k) == n {
							submetric = true
						}
					}
					// it is a facet field and not a submetric field
					if !submetric {
						facet[strings.ToLower(fmt.Sprintf("%v", k))] = v
					}
				} else { // this is the actual gauge data
					if dataFound {
						return errors.New("Data field not specified for multi-column query")
					}
					dataVal = v
					dataFound = true
				}
			}

			if !dataFound {
				return errors.New("Data field not found in result set")
			}

			key, status := r.createMetric(facet, suffix)
			err := setValueForResult(r.Result[key], dataVal)
			if err != nil {
				return err
			}
			facetsWithResult[key] = status
		}
	}
	r.registerMetrics(facetsWithResult)
	return nil
}

// RegisterMetrics registers and unregister gauges
func (r *QueryResult) registerMetrics(facetsWithResult map[string]metricStatus) {
	for key, m := range r.Result {
		status, ok := facetsWithResult[key]
		if !ok {
			fmt.Println("Unregistering metric", key)
			prometheus.Unregister(m)
			delete(r.Result, key)
			continue
		}
		if status == unregistered {
			defer func(key string, m prometheus.Gauge) {
				fmt.Println("Registering metric", key)
				prometheus.MustRegister(m)
			}(key, m)
		}
	}
}


<|start_filename|>worker.go<|end_filename|>
package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/jpillora/backoff"
	"golang.org/x/net/context"
)

// Backoff for fetching. It starts by waiting the minimum duration after a
// failed fetch, doubling it each time (with a bitter of jitter) up to max
// duration between requests.
var defaultBackoff = backoff.Backoff{
	Min:    1 * time.Second,
	Max:    5 * time.Minute,
	Jitter: true,
	Factor: 2,
}

// Worker is responsible for fetching data via SQL Agent
type Worker struct {
	query   *Query
	payload []byte
	client  *http.Client
	result  *QueryResult
	log     *log.Logger
	backoff backoff.Backoff
	ctx     context.Context
}

func (w *Worker) setQueryResultMetrics(recs records) {
	err := w.result.SetMetrics(recs, w.query.ValueOnError)
	if err != nil {
		w.log.Printf("Error setting metrics: %s", err)
		return
	}
}

func (w *Worker) queryResultError() {
	w.setQueryResultMetrics(nil)
}

func (w *Worker) fetchRecords(url string) error {
	var (
		t    time.Time
		err  error
		req  *http.Request
		resp *http.Response
	)

	for {
		t = time.Now()

		req, err = http.NewRequest("POST", url, bytes.NewBuffer(w.payload))

		if err != nil {
			panic(err)
		}
		req = req.WithContext(w.ctx)

		// Set the content-type of the request body and accept LD-JSON.
		req.Header.Set("content-type", "application/json")
		req.Header.Set("accept", "application/json")

		resp, err = w.client.Do(req)

		// No formal error, but a non-successful status code. Construct an error.
		if err == nil && resp.StatusCode != 200 {
			b, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			err = fmt.Errorf("%s: %s", resp.Status, string(b))
		}

		// No error, break to read the data.
		if err == nil {
			break
		}
		w.log.Print(err)

		w.queryResultError()

		// Backoff on an error.
		d := w.backoff.Duration()
		w.log.Printf("Backing off for %s", d)
		select {
		case <-time.After(d):
			continue
		case <-w.ctx.Done():
			return errors.New("Execution was canceled")
		}
	}

	w.backoff.Reset()

	w.log.Printf("Fetch took %s", time.Now().Sub(t))

	var recs []record

	defer resp.Body.Close()

	if err = json.NewDecoder(resp.Body).Decode(&recs); err != nil {
		return err
	}

	w.setQueryResultMetrics(recs)

	return nil
}

// Start fetching data from specified URL
func (w *Worker) Start(url string, wg *sync.WaitGroup) {
	tick := func() {
		err := w.fetchRecords(url)
		if err != nil {
			w.log.Printf("Error fetching records: %s", err)
			return
		}
	}

	tick()
	ticker := time.NewTicker(w.query.Interval)

	for {
		select {
		case <-w.ctx.Done():
			wg.Done()
			w.log.Printf("Stopping worker")
			return

		case <-ticker.C:
			tick()
		}
	}
}

// NewWorker creates a new worker for a query.
func NewWorker(ctx context.Context, q *Query) *Worker {
	// Encode the payload once for all subsequent requests.
	payload, err := json.Marshal(map[string]interface{}{
		"driver":     q.Driver,
		"connection": q.Connection,
		"sql":        q.SQL,
		"params":     q.Params,
	})

	if err != nil {
		panic(err)
	}

	return &Worker{
		query:   q,
		result:  NewQueryResult(q),
		payload: payload,
		backoff: defaultBackoff,
		log:     log.New(os.Stderr, fmt.Sprintf("[%s] ", q.Name), log.LstdFlags),
		client: &http.Client{
			Timeout: q.Timeout,
		},
		ctx: ctx,
	}
}
