configure labels export with EXPORT_CSV_LABELS_LIST and EXPORT_CSV_LABELS_ALL

Signed-off-by: r2k1 <yokree@gmail.com>

pkg/cmd/costmodel/costmodel.go

@@ -2,8 +2,8 @@ package costmodel
 
 import (
 	"context"
+	"fmt"
 	"net/http"
-	"os"
 	"time"
 
 	"github.com/julienschmidt/httprouter"
@@ -34,7 +34,10 @@ func Execute(opts *CostModelOpts) error {
 	log.Infof("Starting cost-model version %s", version.FriendlyVersion())
 	a := costmodel.Initialize()
 
-	StartExportWorker(context.Background(), a.Model)
+	err := StartExportWorker(context.Background(), a.Model)
+	if err != nil {
+		log.Errorf("couldn't start CSV export worker: %v", err)
+	}
 
 	rootMux := http.NewServeMux()
 	a.Router.GET("/healthz", Healthz)
@@ -48,18 +51,15 @@ func Execute(opts *CostModelOpts) error {
 	return http.ListenAndServe(":9003", errors.PanicHandlerMiddleware(handler))
 }
 
-func StartExportWorker(ctx context.Context, model costmodel.AllocationModel) {
+func StartExportWorker(ctx context.Context, model costmodel.AllocationModel) error {
 	// TODO: there should be a better way to load the configuration
-	exportPath := os.Getenv(env.ExportCSVFile)
+	exportPath := env.GetExportCSVFile()
 	if exportPath == "" {
-		log.Infof("%s is not set, skipping CSV exporter", env.ExportCSVFile)
-		return
+		return fmt.Errorf("%s is not set, skipping CSV exporter", env.ExportCSVFile)
 	}
-
 	fm, err := filemanager.NewFileManager(exportPath)
 	if err != nil {
-		log.Errorf("could not start CSV exporter: %v", err)
-		return
+		return fmt.Errorf("could not create file manager: %v", err)
 	}
 	go func() {
 		log.Info("Starting CSV exporter worker...")
@@ -71,7 +71,7 @@ func StartExportWorker(ctx context.Context, model costmodel.AllocationModel) {
 			case <-ctx.Done():
 				return
 			case <-time.After(nextRunAt.Sub(time.Now())):
-				err := costmodel.UpdateCSV(ctx, fm, model)
+				err := costmodel.UpdateCSV(ctx, fm, model, env.GetExportCSVLabelsAll(), env.GetExportCSVLabelsList())
 				if err != nil {
 					// it's background worker, log error and carry on, maybe next time it will work
 					log.Errorf("Error updating CSV: %s", err)
@@ -83,4 +83,5 @@ func StartExportWorker(ctx context.Context, model costmodel.AllocationModel) {
 			}
 		}
 	}()
+	return nil
 }

pkg/costmodel/csv_export.go

@@ -24,10 +24,12 @@ type AllocationModel interface {
 
 var errNoData = errors.New("no data")
 
-func UpdateCSV(ctx context.Context, fileManager filemanager.FileManager, model AllocationModel) error {
+func UpdateCSV(ctx context.Context, fileManager filemanager.FileManager, model AllocationModel, labelsAll bool, labels []string) error {
 	exporter := &csvExporter{
 		FileManager: fileManager,
 		Model:       model,
+		LabelsAll:   labelsAll,
+		Labels:      labels,
 	}
 	return exporter.Update(ctx)
 }
@@ -35,6 +37,8 @@ func UpdateCSV(ctx context.Context, fileManager filemanager.FileManager, model A
 type csvExporter struct {
 	FileManager filemanager.FileManager
 	Model       AllocationModel
+	Labels      []string // If not empty, create a column for each label prefixed with "Label_"
+	LabelsAll   bool     // if true, export all labels to a "Labels" column in JSON format
 }
 
 // Update updates CSV file in cloud storage with new allocation data
@@ -154,36 +158,173 @@ func (e *csvExporter) writeCSVToWriter(ctx context.Context, w io.Writer, dates [
 	fmtFloat := func(f float64) string {
 		return strconv.FormatFloat(f, 'f', -1, 64)
 	}
+
+	type rowData struct {
+		date  time.Time
+		alloc *kubecost.Allocation
+	}
+
+	type columnDef struct {
+		column string
+		value  func(data rowData) string
+	}
+
+	csvDef := []columnDef{
+		{
+			column: "Date",
+			value: func(data rowData) string {
+				return data.date.Format("2006-01-02")
+			},
+		},
+		{
+			column: "Namespace",
+			value: func(data rowData) string {
+				return data.alloc.Properties.Namespace
+			},
+		},
+		{
+			column: "ControllerKind",
+			value: func(data rowData) string {
+				return data.alloc.Properties.ControllerKind
+			},
+		},
+		{
+			column: "ControllerName",
+			value: func(data rowData) string {
+				return data.alloc.Properties.Controller
+			},
+		},
+		{
+			column: "Pod",
+			value: func(data rowData) string {
+				return data.alloc.Properties.Pod
+			},
+		},
+		{
+			column: "Container",
+			value: func(data rowData) string {
+				return data.alloc.Properties.Container
+			},
+		},
+		{
+			column: "CPUCoreUsageAverage",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.CPUCoreUsageAverage)
+			},
+		},
+		{
+			column: "CPUCoreRequestAverage",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.CPUCoreRequestAverage)
+			},
+		},
+		{
+			column: "RAMBytesUsageAverage",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.RAMBytesUsageAverage)
+			},
+		},
+		{
+			column: "RAMBytesRequestAverage",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.RAMBytesRequestAverage)
+			},
+		},
+		{
+			column: "NetworkReceiveBytes",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.NetworkReceiveBytes)
+			},
+		},
+		{
+			column: "NetworkTransferBytes",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.NetworkTransferBytes)
+			},
+		},
+		{
+			column: "GPUs",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.GPUs())
+			},
+		},
+		{
+			column: "PVBytes",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.PVBytes())
+			},
+		},
+		{
+			column: "CPUCost",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.CPUTotalCost())
+			},
+		},
+		{
+			column: "RAMCost",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.RAMTotalCost())
+			},
+		},
+		{
+			column: "NetworkCost",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.NetworkTotalCost())
+			},
+		},
+		{
+			column: "PVCost",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.PVTotalCost())
+			},
+		},
+		{
+			column: "GPUCost",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.GPUTotalCost())
+			},
+		},
+		{
+			column: "TotalCost",
+			value: func(data rowData) string {
+				return fmtFloat(data.alloc.TotalCost())
+			},
+		},
+	}
+	if e.LabelsAll {
+		csvDef = append(csvDef, columnDef{
+			column: "Labels",
+			value: func(data rowData) string {
+				return fmtLabelsCSV(data.alloc.Properties.Labels)
+			},
+		})
+	}
+	for i := range e.Labels {
+		label := e.Labels[i] // it's important to copy the label name, otherwise all closures will reference the same label
+		csvDef = append(csvDef, columnDef{
+			column: "Label_" + label,
+			value: func(data rowData) string {
+				value, _ := data.alloc.Properties.Labels[label]
+				return value
+			},
+		})
+	}
+	csvDef = append(csvDef)
+
+	header := make([]string, 0, len(csvDef))
+	for _, def := range csvDef {
+		header = append(header, def.column)
+	}
+
 	csvWriter := csv.NewWriter(w)
-	err := csvWriter.Write([]string{
-		"Date",
-		"Namespace",
-		"ControllerKind",
-		"ControllerName",
-		"Pod",
-		"Container",
-		"Labels",
-
-		"CPUCoreUsageAverage",
-		"CPUCoreRequestAverage",
-		"RAMBytesUsageAverage",
-		"RAMBytesRequestAverage",
-		"NetworkReceiveBytes",
-		"NetworkTransferBytes",
-		"GPUs",
-		"PVBytes",
-
-		"CPUCost",
-		"RAMCost",
-		"NetworkCost",
-		"PVCost",
-		"GPUCost",
-		"TotalCost",
-	})
+	lines := 0
+	err := csvWriter.Write(header)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to write header: %w", err)
 	}
-	lines := 0
+
+	log.Infof("writing CSV with header: %v", header)
+
 	for _, date := range dates {
 		start := time.Date(date.Year(), date.Month(), date.Day(), 0, 0, 0, 0, time.UTC)
 		end := start.AddDate(0, 0, 1)
@@ -196,35 +337,15 @@ func (e *csvExporter) writeCSVToWriter(ctx context.Context, w io.Writer, dates [
 			if err := ctx.Err(); err != nil {
 				return err
 			}
-
-			err := csvWriter.Write([]string{
-				date.Format("2006-01-02"),
-				alloc.Properties.Namespace,
-				alloc.Properties.ControllerKind,
-				alloc.Properties.Controller,
-				alloc.Properties.Pod,
-				alloc.Properties.Container,
-				fmtLabelsCSV(alloc.Properties.Labels),
-
-				fmtFloat(alloc.CPUCoreUsageAverage),
-				fmtFloat(alloc.CPUCoreRequestAverage),
-				fmtFloat(alloc.RAMBytesUsageAverage),
-				fmtFloat(alloc.RAMBytesRequestAverage),
-				fmtFloat(alloc.NetworkReceiveBytes),
-				fmtFloat(alloc.NetworkTransferBytes),
-				fmtFloat(alloc.GPUs()),
-				fmtFloat(alloc.PVBytes()),
-
-				fmtFloat(alloc.CPUTotalCost()),
-				fmtFloat(alloc.RAMTotalCost()),
-				fmtFloat(alloc.NetworkTotalCost()),
-				fmtFloat(alloc.PVCost()),
-				fmtFloat(alloc.GPUCost),
-				fmtFloat(alloc.TotalCost()),
-			})
+			row := make([]string, 0, len(csvDef))
+			for _, def := range csvDef {
+				row = append(row, def.value(rowData{date: date, alloc: alloc}))
+			}
+			err := csvWriter.Write(row)
 			if err != nil {
-				return err
+				return fmt.Errorf("failed to write csv row: %w", err)
 			}
+
 			lines++
 		}
 	}

pkg/costmodel/csv_export_test.go

@@ -48,6 +48,39 @@ func Test_UpdateCSV(t *testing.T) {
 									Cost:      2.0,
 								},
 							}, // 2 PVBytes, 2 PVCost
+							Properties: &kubecost.AllocationProperties{
+								Namespace:      "test-namespace",
+								Controller:     "test-controller-name",
+								ControllerKind: "test-controller-kind",
+								Pod:            "test-pod",
+								Container:      "test-container",
+							},
+						},
+					},
+				}, nil
+			},
+		}
+		err := UpdateCSV(context.TODO(), storage, model, false, nil)
+		require.NoError(t, err)
+		// uploaded a single file with the data
+		assert.Len(t, model.ComputeAllocationCalls(), 1)
+		assert.Equal(t, time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC), model.ComputeAllocationCalls()[0].Start)
+		assert.Equal(t, time.Date(2021, 1, 2, 0, 0, 0, 0, time.UTC), model.ComputeAllocationCalls()[0].End)
+		assert.Equal(t, `Date,Namespace,ControllerKind,ControllerName,Pod,Container,CPUCoreUsageAverage,CPUCoreRequestAverage,RAMBytesUsageAverage,RAMBytesRequestAverage,NetworkReceiveBytes,NetworkTransferBytes,GPUs,PVBytes,CPUCost,RAMCost,NetworkCost,PVCost,GPUCost,TotalCost
+2021-01-01,test-namespace,test-controller-kind,test-controller-name,test-pod,test-container,0.1,0.2,0.4,0.5,11,10,2,2,0.3,0.6,0.9,2,0.8,4.6000000000000005
+`, string(storage.Data))
+	})
+
+	t.Run("export labels", func(t *testing.T) {
+		storage := &filemanager.InMemoryFile{}
+		model := &AllocationModelMock{
+			DateRangeFunc: func() (time.Time, time.Time, error) {
+				return time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC), time.Date(2021, 1, 2, 0, 0, 0, 0, time.UTC), nil
+			},
+			ComputeAllocationFunc: func(start time.Time, end time.Time, resolution time.Duration) (*kubecost.AllocationSet, error) {
+				return &kubecost.AllocationSet{
+					Allocations: map[string]*kubecost.Allocation{
+						"test": {
 							Properties: &kubecost.AllocationProperties{
 								Namespace:      "test-namespace",
 								Controller:     "test-controller-name",
@@ -64,21 +97,19 @@ func Test_UpdateCSV(t *testing.T) {
 				}, nil
 			},
 		}
-		err := UpdateCSV(context.TODO(), storage, model)
+		err := UpdateCSV(context.TODO(), storage, model, true, []string{"test-label1", "test-label2"})
 		require.NoError(t, err)
 		// uploaded a single file with the data
 		assert.Len(t, model.ComputeAllocationCalls(), 1)
-		assert.Equal(t, time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC), model.ComputeAllocationCalls()[0].Start)
-		assert.Equal(t, time.Date(2021, 1, 2, 0, 0, 0, 0, time.UTC), model.ComputeAllocationCalls()[0].End)
-		assert.Equal(t, `Date,Namespace,ControllerKind,ControllerName,Pod,Container,Labels,CPUCoreUsageAverage,CPUCoreRequestAverage,RAMBytesUsageAverage,RAMBytesRequestAverage,NetworkReceiveBytes,NetworkTransferBytes,GPUs,PVBytes,CPUCost,RAMCost,NetworkCost,PVCost,GPUCost,TotalCost
-2021-01-01,test-namespace,test-controller-kind,test-controller-name,test-pod,test-container,"{""test-label1"":""test-value1"",""test-label2"":""test-value2""}",0.1,0.2,0.4,0.5,11,10,2,2,0.3,0.6,0.9,2,0.8,4.6000000000000005
+		assert.Equal(t, `Date,Namespace,ControllerKind,ControllerName,Pod,Container,CPUCoreUsageAverage,CPUCoreRequestAverage,RAMBytesUsageAverage,RAMBytesRequestAverage,NetworkReceiveBytes,NetworkTransferBytes,GPUs,PVBytes,CPUCost,RAMCost,NetworkCost,PVCost,GPUCost,TotalCost,Labels,Label_test-label1,Label_test-label2
+2021-01-01,test-namespace,test-controller-kind,test-controller-name,test-pod,test-container,0,0,0,0,0,0,0,0,0,0,0,0,0,0,"{""test-label1"":""test-value1"",""test-label2"":""test-value2""}",test-value1,test-value2
 `, string(storage.Data))
 	})
 
 	t.Run("merge new data with previous data (with different CSV structure)", func(t *testing.T) {
 		storage := &filemanager.InMemoryFile{
-			Data: []byte(`Date,Namespace,CPUCoreUsageAverage,CPUCoreRequestAverage,CPUCost,RAMBytesUsageAverage,RAMBytesRequestAverage,RAMCost
-2021-01-01,test-namespace,0.1,0.2,0.3,0.4,0.5,0.6
+			Data: []byte(`Date,Namespace,CPUCoreUsageAverage,CPUCoreRequestAverage,CPUCost,RAMBytesUsageAverage,RAMBytesRequestAverage,RAMCost,Label_app
+2021-01-01,test-namespace,0.1,0.2,0.3,0.4,0.5,0.6,app1
 `),
 		}
 		model := &AllocationModelMock{
@@ -98,7 +129,7 @@ func Test_UpdateCSV(t *testing.T) {
 				}, nil
 			},
 		}
-		err := UpdateCSV(context.TODO(), storage, model)
+		err := UpdateCSV(context.TODO(), storage, model, false, nil)
 		require.NoError(t, err)
 		// uploaded a single file with the data
 		assert.Len(t, model.ComputeAllocationCalls(), 1)
@@ -106,8 +137,8 @@ func Test_UpdateCSV(t *testing.T) {
 		// 2021-01-01 is already in the export file, so we only compute for 2021-01-02
 		assert.Equal(t, time.Date(2021, 1, 2, 0, 0, 0, 0, time.UTC), model.ComputeAllocationCalls()[0].Start)
 		assert.Equal(t, time.Date(2021, 1, 3, 0, 0, 0, 0, time.UTC), model.ComputeAllocationCalls()[0].End)
-		assert.Equal(t, `Date,Namespace,CPUCoreUsageAverage,CPUCoreRequestAverage,CPUCost,RAMBytesUsageAverage,RAMBytesRequestAverage,RAMCost,ControllerKind,ControllerName,Pod,Container,Labels,NetworkReceiveBytes,NetworkTransferBytes,GPUs,PVBytes,NetworkCost,PVCost,GPUCost,TotalCost
-2021-01-01,test-namespace,0.1,0.2,0.3,0.4,0.5,0.6,,,,,,,,,,,,,
+		assert.Equal(t, `Date,Namespace,CPUCoreUsageAverage,CPUCoreRequestAverage,CPUCost,RAMBytesUsageAverage,RAMBytesRequestAverage,RAMCost,Label_app,ControllerKind,ControllerName,Pod,Container,NetworkReceiveBytes,NetworkTransferBytes,GPUs,PVBytes,NetworkCost,PVCost,GPUCost,TotalCost
+2021-01-01,test-namespace,0.1,0.2,0.3,0.4,0.5,0.6,app1,,,,,,,,,,,,
 2021-01-02,test-namespace,0,0,1,0,0,0,,,,,,0,0,0,0,0,0,0,1
 `, string(storage.Data))
 	})
@@ -124,7 +155,7 @@ func Test_UpdateCSV(t *testing.T) {
 				return time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC), time.Date(2021, 1, 2, 0, 0, 0, 0, time.UTC), nil
 			},
 		}
-		err := UpdateCSV(context.TODO(), storage, model)
+		err := UpdateCSV(context.TODO(), storage, model, false, nil)
 		require.Equal(t, err, errNoData)
 		assert.Equal(t, string(storage.Data), data)
 		assert.Len(t, model.ComputeAllocationCalls(), 0)
@@ -142,7 +173,7 @@ func Test_UpdateCSV(t *testing.T) {
 			},
 		}
 		storage := &filemanager.InMemoryFile{}
-		err := UpdateCSV(context.TODO(), storage, model)
+		err := UpdateCSV(context.TODO(), storage, model, false, nil)
 		require.Equal(t, err, errNoData)
 	})
 }

pkg/env/costmodelenv.go

@@ -96,7 +96,9 @@ const (
 
 	regionOverrideList = "REGION_OVERRIDE_LIST"
 
-	ExportCSVFile = "EXPORT_CSV_FILE"
+	ExportCSVFile       = "EXPORT_CSV_FILE"
+	ExportCSVLabelsList = "EXPORT_CSV_LABELS_LIST"
+	ExportCSVLabelsAll  = "EXPORT_CSV_LABELS_ALL"
 )
 
 const DefaultConfigMountPath = "/var/configs"
@@ -107,6 +109,18 @@ func IsETLReadOnlyMode() bool {
 	return GetBool(ETLReadOnlyMode, false)
 }
 
+func GetExportCSVFile() string {
+	return Get(ExportCSVFile, "")
+}
+
+func GetExportCSVLabelsAll() bool {
+	return GetBool(ExportCSVLabelsAll, false)
+}
+
+func GetExportCSVLabelsList() []string {
+	return GetList(ExportCSVLabelsList, ",")
+}
+
 // GetKubecostConfigBucket returns a file location for a mounted bucket configuration which is used to store
 // a subset of kubecost configurations that require sharing via remote storage.
 func GetKubecostConfigBucket() string {

pkg/filemanager/filemanager.go

@@ -44,6 +44,8 @@ func NewFileManager(path string) (FileManager, error) {
 		return NewGCSStorageFile(path)
 	case strings.Contains(path, "blob.core.windows.net"):
 		return NewAzureBlobFile(path)
+	case path == "":
+		return nil, errors.New("empty path")
 	default:
 		return NewSystemFile(path), nil
 	}