Add types qField and rField

Signed-off-by: Utku Ozdemir <uoz@protonmail.com>

cmd/query_fields_parser/main.go

@@ -16,10 +16,11 @@ func main() {
 	)
 
 	kingpin.Parse()
-	fields, err := exporter.ParseAutoQFields(*nvidiaSmiCommand)
+	qFields, err := exporter.ParseAutoQFields(*nvidiaSmiCommand)
 	if err != nil {
 		fmt.Printf("error: %v\n", err)
 		os.Exit(1)
 	}
+	fields := exporter.QFieldSliceToStringSlice(qFields)
 	fmt.Printf("Fields:\n\n%s\n", strings.Join(fields, "\n"))
 }

internal/exporter/csv.go

@@ -4,46 +4,46 @@ import "strings"
 
 type table struct {
 	rows          []row
-	rFields       []string
-	qFieldToCells map[string][]cell
+	rFields       []rField
+	qFieldToCells map[qField][]cell
 }
 
 type row struct {
-	qFieldToCells map[string]cell
+	qFieldToCells map[qField]cell
 	cells         []cell
 }
 
 type cell struct {
-	qField   string
-	rField   string
+	qField   qField
+	rField   rField
 	rawValue string
 }
 
-func parseCSVIntoTable(queryResult string, qFields []string) table {
+func parseCSVIntoTable(queryResult string, qFields []qField) table {
 	lines := strings.Split(strings.TrimSpace(queryResult), "\n")
 	titlesLine := lines[0]
 	valuesLines := lines[1:]
-	returnedFieldNames := parseCSVLine(titlesLine)
+	rFields := toRFieldSlice(parseCSVLine(titlesLine))
 
 	numCols := len(qFields)
 	numRows := len(valuesLines)
 
 	rows := make([]row, numRows)
 
-	qFieldToCells := make(map[string][]cell)
+	qFieldToCells := make(map[qField][]cell)
 	for _, qField := range qFields {
 		qFieldToCells[qField] = make([]cell, numRows)
 	}
 
 	for rowIndex, valuesLine := range valuesLines {
-		qFieldToCell := make(map[string]cell, numCols)
+		qFieldToCell := make(map[qField]cell, numCols)
 		cells := make([]cell, numCols)
 		rawValues := parseCSVLine(valuesLine)
 		for colIndex, rawValue := range rawValues {
 			qField := qFields[colIndex]
 			gm := cell{
 				qField:   qField,
-				rField:   returnedFieldNames[colIndex],
+				rField:   rFields[colIndex],
 				rawValue: rawValue,
 			}
 			qFieldToCell[qField] = gm
@@ -62,7 +62,7 @@ func parseCSVIntoTable(queryResult string, qFields []string) table {
 
 	return table{
 		rows:          rows,
-		rFields:       returnedFieldNames,
+		rFields:       rFields,
 		qFieldToCells: qFieldToCells,
 	}
 }

internal/exporter/csv_test.go

@@ -14,29 +14,29 @@ Some Dummy GPU, 12.34 W
 )
 
 func TestParseCsvIntoTable(t *testing.T) {
-	parsed := parseCSVIntoTable(testCsv, []string{"name", "power.draw"})
+	parsed := parseCSVIntoTable(testCsv, []qField{"name", "power.draw"})
 	assert.Len(t, parsed.rows, 2)
-	assert.Equal(t, []string{"name", "power.draw [W]"}, parsed.rFields)
+	assert.Equal(t, []rField{"name", "power.draw [W]"}, parsed.rFields)
 
 	cell00 := cell{qField: "name", rField: "name", rawValue: "NVIDIA GeForce RTX 2080 SUPER"}
 	cell01 := cell{qField: "power.draw", rField: "power.draw [W]", rawValue: "30.14 W"}
 	cell10 := cell{qField: "name", rField: "name", rawValue: "Some Dummy GPU"}
 	cell11 := cell{qField: "power.draw", rField: "power.draw [W]", rawValue: "12.34 W"}
 
 	row0 := row{
-		qFieldToCells: map[string]cell{"name": cell00, "power.draw": cell01},
+		qFieldToCells: map[qField]cell{"name": cell00, "power.draw": cell01},
 		cells:         []cell{cell00, cell01},
 	}
 
 	row1 := row{
-		qFieldToCells: map[string]cell{"name": cell10, "power.draw": cell11},
+		qFieldToCells: map[qField]cell{"name": cell10, "power.draw": cell11},
 		cells:         []cell{cell10, cell11},
 	}
 
 	expected := table{
 		rows:    []row{row0, row1},
-		rFields: []string{"name", "power.draw [W]"},
-		qFieldToCells: map[string][]cell{
+		rFields: []rField{"name", "power.draw [W]"},
+		qFieldToCells: map[qField][]cell{
 			"name":       {cell00, cell10},
 			"power.draw": {cell01, cell11},
 		},

internal/exporter/exporter.go

@@ -13,13 +13,19 @@ import (
 	"sync"
 )
 
+// qField stands for query field - the field name before the query
+type qField string
+
+// rField stands for returned field - the field name as returned by the nvidia-smi
+type rField string
+
 const (
 	DefaultPrefix           = "nvidia_smi"
 	DefaultNvidiaSmiCommand = "nvidia-smi"
 )
 
 var (
-	variableLabels = []string{uuidQField, nameQField}
+	variableLabels = []qField{uuidQField, nameQField}
 	numericRegex   = regexp.MustCompile("[+-]?([0-9]*[.])?[0-9]+")
 
 	runCmd = func(cmd *exec.Cmd) error { return cmd.Run() }
@@ -30,8 +36,8 @@ var (
 type gpuExporter struct {
 	mutex                 sync.RWMutex
 	prefix                string
-	qFields               []string
-	qFieldToMetricInfoMap map[string]MetricInfo
+	qFields               []qField
+	qFieldToMetricInfoMap map[qField]MetricInfo
 	nvidiaSmiCommand      string
 	failedScrapesTotal    prometheus.Counter
 	logger                log.Logger
@@ -63,10 +69,10 @@ func New(prefix string, nvidiaSmiCommand string, qFieldsRaw string, logger log.L
 }
 
 func buildQFieldToRFieldMap(logger log.Logger, qFieldsRaw string,
-	nvidiaSmiCommand string) (map[string]string, error) {
+	nvidiaSmiCommand string) (map[qField]rField, error) {
 	qFieldsSeparated := strings.Split(qFieldsRaw, ",")
 
-	var parsedQFields []string
+	var parsedQFields []qField
 	if len(qFieldsSeparated) == 1 && qFieldsSeparated[0] == qFieldsAuto {
 		parsed, err := ParseAutoQFields(nvidiaSmiCommand)
 		if err != nil {
@@ -84,7 +90,7 @@ func buildQFieldToRFieldMap(logger log.Logger, qFieldsRaw string,
 		return nil, err
 	}
 
-	r := make(map[string]string, len(parsedQFields))
+	r := make(map[qField]rField, len(parsedQFields))
 	for i, qField := range parsedQFields {
 		r[qField] = t.rFields[i]
 	}
@@ -131,8 +137,8 @@ func (e *gpuExporter) Collect(ch chan<- prometheus.Metric) {
 	}
 }
 
-func scrape(qFields []string, nvidiaSmiCommand string) (*table, error) {
-	qFieldsJoined := strings.Join(qFields, ",")
+func scrape(qFields []qField, nvidiaSmiCommand string) (*table, error) {
+	qFieldsJoined := strings.Join(QFieldSliceToStringSlice(qFields), ",")
 
 	cmdAndArgs := strings.Fields(nvidiaSmiCommand)
 	cmdAndArgs = append(cmdAndArgs, fmt.Sprintf("--query-gpu=%s", qFieldsJoined))
@@ -195,38 +201,39 @@ func transformRawValue(rawValue string, valueMultiplier float64) (float64, error
 	}
 }
 
-func buildQFieldToMetricInfoMap(prefix string, qFieldtoRFieldMap map[string]string) map[string]MetricInfo {
-	result := make(map[string]MetricInfo)
+func buildQFieldToMetricInfoMap(prefix string, qFieldtoRFieldMap map[qField]rField) map[qField]MetricInfo {
+	result := make(map[qField]MetricInfo)
 	for qField, rField := range qFieldtoRFieldMap {
 		result[qField] = buildMetricInfo(prefix, rField)
 	}
 
 	return result
 }
 
-func buildMetricInfo(prefix string, rField string) MetricInfo {
+func buildMetricInfo(prefix string, rField rField) MetricInfo {
 	fqName, multiplier := buildFQNameAndMultiplier(prefix, rField)
-	desc := prometheus.NewDesc(fqName, "", variableLabels, nil) // todo: add help text
+	desc := prometheus.NewDesc(fqName, "", QFieldSliceToStringSlice(variableLabels), nil) // todo: add help text
 	return MetricInfo{
 		desc:            desc,
 		mType:           prometheus.GaugeValue,
 		valueMultiplier: multiplier,
 	}
 }
 
-func buildFQNameAndMultiplier(prefix string, rField string) (string, float64) {
-	suffixTransformed := rField
+func buildFQNameAndMultiplier(prefix string, rField rField) (string, float64) {
+	rFieldStr := string(rField)
+	suffixTransformed := rFieldStr
 	multiplier := 1.0
-	split := strings.Split(rField, " ")[0]
-	if strings.HasSuffix(rField, " [W]") {
+	split := strings.Split(rFieldStr, " ")[0]
+	if strings.HasSuffix(rFieldStr, " [W]") {
 		suffixTransformed = split + "_watts"
-	} else if strings.HasSuffix(rField, " [MHz]") {
+	} else if strings.HasSuffix(rFieldStr, " [MHz]") {
 		suffixTransformed = split + "_clock_hz"
 		multiplier = 1000000
-	} else if strings.HasSuffix(rField, " [MiB]") {
+	} else if strings.HasSuffix(rFieldStr, " [MiB]") {
 		suffixTransformed = split + "_bytes"
 		multiplier = 1048576
-	} else if strings.HasSuffix(rField, " [%]") {
+	} else if strings.HasSuffix(rFieldStr, " [%]") {
 		suffixTransformed = split + "_ratio"
 		multiplier = 0.01
 	}
@@ -237,8 +244,8 @@ func buildFQNameAndMultiplier(prefix string, rField string) (string, float64) {
 	return fqName, multiplier
 }
 
-func getKeys(m map[string]string) []string {
-	r := make([]string, len(m))
+func getKeys(m map[qField]rField) []qField {
+	r := make([]qField, len(m))
 	i := 0
 	for key := range m {
 		r[i] = key

internal/exporter/exporter_test.go

@@ -105,7 +105,7 @@ func TestBuildMetricInfo(t *testing.T) {
 }
 
 func TestBuildQFieldToMetricInfoMap(t *testing.T) {
-	m := buildQFieldToMetricInfoMap("prefix", map[string]string{"aaa": "AAA", "bbb": "BBB"})
+	m := buildQFieldToMetricInfoMap("prefix", map[qField]rField{"aaa": "AAA", "bbb": "BBB"})
 	assert.Len(t, m, 2)
 
 	metricInfo1 := m["aaa"]

internal/exporter/fields.go

@@ -8,16 +8,16 @@ import (
 )
 
 const (
-	uuidQField    = "uuid"
-	nameQField    = "name"
-	qFieldsAuto   = "AUTO"
-	DefaultQField = qFieldsAuto
+	uuidQField    qField = "uuid"
+	nameQField    qField = "name"
+	qFieldsAuto          = "AUTO"
+	DefaultQField        = qFieldsAuto
 )
 
 var (
 	fieldRegex = regexp.MustCompile(`(?m)\n\s*\n^"([^"]+)"`)
 
-	fallbackQFieldToRFieldMap = map[string]string{
+	fallbackQFieldToRFieldMap = map[qField]rField{
 		"timestamp":                         "timestamp",
 		"driver_version":                    "driver_version",
 		"count":                             "count",
@@ -136,7 +136,7 @@ var (
 	}
 )
 
-func ParseAutoQFields(nvidiaSmiCommand string) ([]string, error) {
+func ParseAutoQFields(nvidiaSmiCommand string) ([]qField, error) {
 	cmdAndArgs := strings.Fields(nvidiaSmiCommand)
 	cmdAndArgs = append(cmdAndArgs, "--help-query-gpu")
 	cmd := exec.Command(cmdAndArgs[0], cmdAndArgs[1:]...)
@@ -153,12 +153,36 @@ func ParseAutoQFields(nvidiaSmiCommand string) ([]string, error) {
 	return fields, nil
 }
 
-func extractQFields(text string) []string {
+func extractQFields(text string) []qField {
 	found := fieldRegex.FindAllStringSubmatch(text, -1)
 
-	var fields []string
+	var fields []qField
 	for _, ss := range found {
-		fields = append(fields, ss[1])
+		fields = append(fields, qField(ss[1]))
 	}
 	return fields
 }
+
+func toQFieldSlice(ss []string) []qField {
+	r := make([]qField, len(ss))
+	for i, s := range ss {
+		r[i] = qField(s)
+	}
+	return r
+}
+
+func toRFieldSlice(ss []string) []rField {
+	r := make([]rField, len(ss))
+	for i, s := range ss {
+		r[i] = rField(s)
+	}
+	return r
+}
+
+func QFieldSliceToStringSlice(qs []qField) []string {
+	r := make([]string, len(qs))
+	for i, q := range qs {
+		r[i] = string(q)
+	}
+	return r
+}

internal/exporter/fields_test.go

@@ -9,8 +9,8 @@ import (
 
 var (
 	//go:embed _fields-test.txt
-	fieldsTest     string
-	expectedFields = []string{
+	fieldsTest      string
+	expectedQFields = []qField{
 		"timestamp", "driver_version", "count", "name", "serial", "uuid", "pci.bus_id",
 		"pci.domain", "pci.bus", "pci.device", "pci.device_id", "pci.sub_device_id", "pcie.link.gen.current",
 		"pcie.link.gen.max", "pcie.link.width.current", "pcie.link.width.max", "index", "display_mode", "display_active",
@@ -53,7 +53,7 @@ var (
 
 func TestExtractQFields(t *testing.T) {
 	fields := extractQFields(fieldsTest)
-	assert.Equal(t, expectedFields, fields)
+	assert.Equal(t, expectedQFields, fields)
 }
 
 func TestParseAutoQFields(t *testing.T) {
@@ -74,5 +74,5 @@ func TestParseAutoQFields(t *testing.T) {
 	assert.Equal(t, capturedCmd.Args[1], "--help-query-gpu")
 
 	assert.NoError(t, err)
-	assert.Equal(t, expectedFields, fields)
+	assert.Equal(t, expectedQFields, fields)
 }