Merge pull request #1 from someengineering/lloesche/csvformat

Add CSV format support

format/format.go

@@ -1,7 +1,11 @@
 package format
 
 import (
+	"bytes"
+	"encoding/csv"
 	"encoding/json"
+	"fmt"
+	"strings"
 
 	"gopkg.in/yaml.v2"
 )
@@ -22,3 +26,45 @@ func ToYAML(data interface{}) (string, error) {
 	}
 	return string(bytes), nil
 }
+
+func ToCSV(data interface{}, headers []string) (string, error) {
+	jsonObj, ok := data.(map[string]interface{})
+	if !ok {
+		return "", fmt.Errorf("data is not a JSON object")
+	}
+
+	var csvBuffer bytes.Buffer
+	writer := csv.NewWriter(&csvBuffer)
+
+	record := make([]string, len(headers))
+	for i, header := range headers {
+		header = strings.TrimPrefix(header, "/")
+		var value interface{} = jsonObj
+
+		for _, key := range strings.Split(header, ".") {
+			if tempMap, ok := value.(map[string]interface{}); ok {
+				value, ok = tempMap[key]
+				if !ok {
+					value = ""
+					break
+				}
+			} else {
+				break
+			}
+		}
+
+		record[i] = fmt.Sprintf("%v", value)
+	}
+
+	if err := writer.Write(record); err != nil {
+		return "", fmt.Errorf("writing record to CSV failed: %w", err)
+	}
+
+	writer.Flush()
+
+	if err := writer.Error(); err != nil {
+		return "", fmt.Errorf("CSV writing failed: %w", err)
+	}
+
+	return csvBuffer.String(), nil
+}

format/format_test.go

@@ -39,3 +39,62 @@ func TestToYAML(t *testing.T) {
 		t.Errorf("Expected YAML string to contain %s, got %s", expectedSubstring, yamlStr)
 	}
 }
+
+func TestToCSV(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    map[string]interface{}
+		headers []string
+		want    string
+	}{
+		{
+			name: "Simple fields",
+			data: map[string]interface{}{
+				"name": "Example",
+				"id":   "123",
+			},
+			headers: []string{"/name", "/id"},
+			want:    "Example,123\n",
+		},
+		{
+			name: "Fields with commas and quotes",
+			data: map[string]interface{}{
+				"description": `Product "A", the best one`,
+				"notes":       "It's, literally, \"awesome\".",
+			},
+			headers: []string{"/description", "/notes"},
+			want:    "\"Product \"\"A\"\", the best one\",\"It's, literally, \"\"awesome\"\".\"\n",
+		},
+		{
+			name: "Nested fields",
+			data: map[string]interface{}{
+				"reported": map[string]interface{}{
+					"location": "Warehouse, 42",
+					"status":   "In-stock",
+				},
+			},
+			headers: []string{"/reported.location", "/reported.status"},
+			want:    "\"Warehouse, 42\",In-stock\n",
+		},
+		{
+			name: "Missing and empty fields",
+			data: map[string]interface{}{
+				"reported": map[string]interface{}{
+					"location": "Remote",
+				},
+			},
+			headers: []string{"/reported.location", "/reported.quantity"},
+			want:    "Remote,\n",
+		},
+	}
+
+	for _, tt := range tests {
+		got, err := ToCSV(tt.data, tt.headers)
+		if err != nil {
+			t.Errorf("TestToCSV %s failed with error: %v", tt.name, err)
+		}
+		if got != tt.want {
+			t.Errorf("TestToCSV %s expected %q, got %q", tt.name, tt.want, got)
+		}
+	}
+}

main.go

@@ -27,7 +27,8 @@ func main() {
 	searchStrPtr := flag.String("search", "", "Search string")
 	usernamePtr := flag.String("username", "", "Username (env FIX_USERNAME)")
 	passwordPtr := flag.String("password", "", "Password (env FIX_PASSWORD)")
-	formatPtr := flag.String("format", "json", "Output format: json or yaml")
+	formatPtr := flag.String("format", "json", "Output format: json, yaml or csv")
+	csvHeadersPtr := flag.String("csv-headers", "id,name,kind,/ancestors.cloud.reported.id,/ancestors.account.reported.id,/ancestors.region.reported.id", "CSV headers (comma-separated, relative to /reported by default)")
 	withEdgesPtr := flag.Bool("with-edges", false, "Include edges in search results")
 	help := flag.Bool("help", false, "Display help information")
 	flag.Parse()
@@ -37,36 +38,44 @@ func main() {
 	}
 
 	withEdges := *withEdgesPtr
-	outputFormat := *formatPtr
 
+	invalidArgs := false
 	username, password, err := utils.SanitizeCredentials(utils.GetEnvOrDefault("FIX_USERNAME", *usernamePtr), utils.GetEnvOrDefault("FIX_PASSWORD", *passwordPtr))
 	if err != nil {
 		fmt.Println("Invalid username or password:", err)
-		os.Exit(1)
+		invalidArgs = true
 	}
 	searchStr, err := utils.SanitizeSearchString(*searchStrPtr)
 	if err != nil {
 		fmt.Println("Invalid search string:", err)
-		os.Exit(1)
+		invalidArgs = true
 	}
 	apiEndpoint, err := utils.SanitizeAPIEndpoint(utils.GetEnvOrDefault("FIX_ENDPOINT", *apiEndpointPtr))
 	if err != nil {
 		fmt.Println("Invalid API endpoint:", err)
-		os.Exit(1)
+		invalidArgs = true
 	}
 	fixToken, err := utils.SanitizeToken(utils.GetEnvOrDefault("FIX_TOKEN", *fixTokenPtr))
 	if err != nil {
 		fmt.Println("Invalid token:", err)
-		os.Exit(1)
+		invalidArgs = true
 	}
 	workspaceID, err := utils.SanitizeWorkspaceId(utils.GetEnvOrDefault("FIX_WORKSPACE", *workspacePtr))
 	if err != nil {
 		fmt.Println("Invalid workspace ID:", err)
-		os.Exit(1)
+		invalidArgs = true
 	}
-
-	if outputFormat != "json" && outputFormat != "yaml" {
-		fmt.Println("Invalid output format")
+	csvHeaders, err := utils.SanitizeCSVHeaders(*csvHeadersPtr)
+	if err != nil {
+		fmt.Println("Invalid CSV headers:", err)
+		invalidArgs = true
+	}
+	outputFormat, err := utils.SanitizeOutputFormat(*formatPtr)
+	if err != nil {
+		fmt.Println("Invalid output format:", err)
+		invalidArgs = true
+	}
+	if invalidArgs {
 		os.Exit(1)
 	}
 
@@ -96,6 +105,8 @@ func main() {
 				firstResult = false
 			}
 			output, err = format.ToYAML(result)
+		case "csv":
+			output, err = format.ToCSV(result, csvHeaders)
 		default:
 			output, err = format.ToJSON(result)
 		}

utils/utils.go

@@ -72,3 +72,37 @@ func SanitizeWorkspaceId(workspaceId string) (string, error) {
 
 	return workspaceId, nil
 }
+
+func SanitizeCSVHeaders(headers string) ([]string, error) {
+	if headers == "" {
+		return nil, fmt.Errorf("headers cannot be empty")
+	}
+
+	rawHeaders := strings.Split(headers, ",")
+	if len(rawHeaders) == 0 {
+		return nil, fmt.Errorf("at least one header must be specified")
+	}
+
+	csvHeaders := make([]string, len(rawHeaders))
+	for i, header := range rawHeaders {
+		trimmedHeader := strings.TrimSpace(header)
+		if trimmedHeader == "" {
+			return nil, fmt.Errorf("empty CSV header found")
+		}
+
+		if !strings.HasPrefix(trimmedHeader, "/") {
+			trimmedHeader = "/reported." + trimmedHeader
+		}
+		csvHeaders[i] = trimmedHeader
+	}
+	return csvHeaders, nil
+}
+
+func SanitizeOutputFormat(format string) (string, error) {
+	switch format {
+	case "json", "yaml", "csv":
+		return format, nil
+	default:
+		return "", fmt.Errorf("unsupported output format")
+	}
+}

utils/utils_test.go

@@ -2,6 +2,7 @@ package utils
 
 import (
 	"os"
+	"reflect"
 	"strings"
 	"testing"
 )
@@ -140,3 +141,121 @@ func TestSanitizeWorkspaceId(t *testing.T) {
 		}
 	}
 }
+
+func TestSanitizeOutputFormat(t *testing.T) {
+	tests := []struct {
+		name      string
+		format    string
+		want      string
+		wantError bool
+	}{
+		{
+			name:      "Valid format json",
+			format:    "json",
+			want:      "json",
+			wantError: false,
+		},
+		{
+			name:      "Valid format yaml",
+			format:    "yaml",
+			want:      "yaml",
+			wantError: false,
+		},
+		{
+			name:      "Valid format csv",
+			format:    "csv",
+			want:      "csv",
+			wantError: false,
+		},
+		{
+			name:      "Unsupported format",
+			format:    "xml",
+			want:      "",
+			wantError: true,
+		},
+		{
+			name:      "Empty format",
+			format:    "",
+			want:      "",
+			wantError: true,
+		},
+		{
+			name:      "Whitespace format",
+			format:    " ",
+			want:      "",
+			wantError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		got, err := SanitizeOutputFormat(tt.format)
+		if (err != nil) != tt.wantError {
+			t.Errorf("%s: SanitizeOutputFormat(%s) expected error: %v, got: %v", tt.name, tt.format, tt.wantError, err)
+		}
+		if got != tt.want {
+			t.Errorf("%s: SanitizeOutputFormat(%s) = %v, want %v", tt.name, tt.format, got, tt.want)
+		}
+	}
+}
+
+func TestSanitizeCSVHeaders(t *testing.T) {
+	tests := []struct {
+		name      string
+		headers   string
+		want      []string
+		wantError bool
+	}{
+		{
+			name:      "Non-empty headers without leading slash",
+			headers:   "id,name,kind",
+			want:      []string{"/reported.id", "/reported.name", "/reported.kind"},
+			wantError: false,
+		},
+		{
+			name:      "Headers with leading slash",
+			headers:   "/metadata.expires,/metadata.cleaned",
+			want:      []string{"/metadata.expires", "/metadata.cleaned"},
+			wantError: false,
+		},
+		{
+			name:      "Mixed headers",
+			headers:   "name,/metadata.expires,kind",
+			want:      []string{"/reported.name", "/metadata.expires", "/reported.kind"},
+			wantError: false,
+		},
+		{
+			name:      "Empty headers string",
+			headers:   "",
+			want:      nil,
+			wantError: true,
+		},
+		{
+			name:      "Only whitespace",
+			headers:   "  ",
+			want:      nil,
+			wantError: true,
+		},
+		{
+			name:      "Header with only commas",
+			headers:   ",,,",
+			want:      nil,
+			wantError: true,
+		},
+		{
+			name:      "Empty header among valid headers",
+			headers:   "id,,name",
+			want:      nil,
+			wantError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		got, err := SanitizeCSVHeaders(tt.headers)
+		if (err != nil) != tt.wantError {
+			t.Errorf("%s: SanitizeCSVHeaders() error = %v, wantError %v", tt.name, err, tt.wantError)
+		}
+		if !reflect.DeepEqual(got, tt.want) {
+			t.Errorf("%s: SanitizeCSVHeaders() = %v, want %v", tt.name, got, tt.want)
+		}
+	}
+}