From 642e864de0bc386d8c43a717e1d275f496e50a3f Mon Sep 17 00:00:00 2001
From: Marco Mayer <marco.mayer@spring-media.de>
Date: Thu, 4 Aug 2022 16:42:32 +0200
Subject: [PATCH] implement custom value functions

---
 pkg/csvexport/csv.go      | 38 +++++++++++++++++++++++++++++++-
 pkg/csvexport/csv_test.go | 46 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/pkg/csvexport/csv.go b/pkg/csvexport/csv.go
index e1d503a..0c79713 100644
--- a/pkg/csvexport/csv.go
+++ b/pkg/csvexport/csv.go
@@ -35,6 +35,22 @@ type ScanOption struct {
 
 type Columns []Column
 
+func (cols Columns) ValueFunc(colName string) (ValueFunc, string, bool) {
+	for i, c := range cols {
+		if c.Name != colName {
+			continue
+		}
+
+		if cols[i].ValueFunc != nil {
+			return cols[i].ValueFunc, cols[i].ValueFuncCol, true
+		}
+
+		return nil, "", false
+	}
+
+	return nil, "", false
+}
+
 type Column struct {
 	Name string
 	// If TargetName is set, Name gets renamed to TargetName in header
@@ -42,8 +58,13 @@ type Column struct {
 	// If OverwriteValue is set to true, all values of the column are set to OverwriteWithValue.
 	OverwriteValue     bool
 	OverwriteWithValue interface{}
+	// Function to process value in col and use that as result
+	ValueFunc    ValueFunc
+	ValueFuncCol string
 }
 
+type ValueFunc func(v interface{}) (string, error)
+
 type CSVExporter struct {
 	cols Columns
 }
@@ -96,7 +117,7 @@ func DynamoToCSV(db Storage, ctx context.Context, scanOpt ScanOption, opts ...Op
 					header = append(header, headerName)
 				}
 
-				if len(keyOrder) == 0 {
+				if csvExp.cols == nil {
 					for k := range attr {
 						keyOrder = append(keyOrder, k)
 					}
@@ -119,6 +140,21 @@ func DynamoToCSV(db Storage, ctx context.Context, scanOpt ScanOption, opts ...Op
 					value = csvExp.cols[i].OverwriteWithValue
 				}
 
+				// Custom function?
+				valueFn, valueFnCol, ok := csvExp.cols.ValueFunc(k)
+
+				if ok {
+					if valueFnCol != "" {
+						value = attr[valueFnCol]
+					}
+					newVal, err := valueFn(value)
+					if err != nil {
+						return nil, fmt.Errorf("failed to process custom valueFunction on column %s: %w", valueFnCol, err)
+					}
+					record = append(record, newVal)
+					continue
+				}
+
 				switch val := value.(type) {
 				case float64:
 					// protect exponential notation layout
diff --git a/pkg/csvexport/csv_test.go b/pkg/csvexport/csv_test.go
index 7f81a0e..cd1fd11 100644
--- a/pkg/csvexport/csv_test.go
+++ b/pkg/csvexport/csv_test.go
@@ -2,6 +2,7 @@ package csvexport_test
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/service/dynamodb"
@@ -128,3 +129,48 @@ func TestDynamoToCSVWithColsTargetName(t *testing.T) {
 `
 	assert.Equal(t, expectedCSV, string(b))
 }
+
+func TestDynamoToCSVWithColsValueFunc(t *testing.T) {
+	t.Parallel()
+	db := mockScan{resp: dynamoMockResp}
+
+	valueFn := func(val interface{}) (string, error) {
+		v, ok := val.(string)
+		if !ok {
+			return "", fmt.Errorf("failed to cast to string")
+		}
+
+		if v == "Meldungen1" {
+			return "Meldungen123", nil
+		}
+
+		return v, nil
+	}
+
+	valueFnNew := func(val interface{}) (string, error) {
+		v, ok := val.(string)
+		if !ok {
+			return "", fmt.Errorf("failed to cast to string")
+		}
+
+		if v == "Meldungen2" {
+			return "Meldungen999", nil
+		}
+
+		return v, nil
+	}
+	cols := csvexport.Columns{
+		csvexport.Column{Name: "ArticleLastUpdated"},
+		csvexport.Column{Name: "Block", ValueFunc: valueFn},
+		csvexport.Column{Name: "NewBlock", ValueFunc: valueFnNew, ValueFuncCol: "Block"},
+	}
+	b, err := csvexport.DynamoToCSV(db, context.Background(), csvexport.ScanOption{}, csvexport.WithColumns(cols))
+
+	assert.Nil(t, err)
+
+	expectedCSV := `ArticleLastUpdated,Block,NewBlock
+2022-04-27T08:36:48.386Z,Meldungen123,Meldungen1
+2022-04-28T08:36:48.386Z,Meldungen2,Meldungen999
+`
+	assert.Equal(t, expectedCSV, string(b))
+}