From 1897bef08d1da5f1af6b2101b8d46eea82a20885 Mon Sep 17 00:00:00 2001
From: Dan Sosedoff <dan.sosedoff@gmail.com>
Date: Wed, 21 Dec 2022 15:02:40 -0600
Subject: [PATCH] Results struct cleanup (#627)

* Add results format test
* Init results slice before formatting data
* Add extra test for nil check in CSV function
---
 pkg/client/client.go      |  9 ---------
 pkg/client/result.go      | 32 +++++++++++++++++++-------------
 pkg/client/result_test.go | 33 +++++++++++++++++++++++++++------
 3 files changed, 46 insertions(+), 28 deletions(-)

diff --git a/pkg/client/client.go b/pkg/client/client.go
index 008733c..f1d7307 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -34,15 +34,6 @@ type Client struct {
 	ConnectionString string           `json:"connection_string"`
 }
 
-// Struct to hold table rows browsing options
-type RowsOptions struct {
-	Where      string // Custom filter
-	Offset     int    // Number of rows to skip
-	Limit      int    // Number of rows to fetch
-	SortColumn string // Column to sort by
-	SortOrder  string // Sort direction (ASC, DESC)
-}
-
 func getSchemaAndTable(str string) (string, string) {
 	chunks := strings.Split(str, ".")
 	if len(chunks) == 1 {
diff --git a/pkg/client/result.go b/pkg/client/result.go
index f446a02..c0e3de6 100644
--- a/pkg/client/result.go
+++ b/pkg/client/result.go
@@ -22,8 +22,18 @@ const (
 )
 
 type (
+	// Row represents a single row of data
 	Row []interface{}
 
+	// RowsOptions contains a list of parameters for table browsing requests
+	RowsOptions struct {
+		Where      string // Custom filter
+		Offset     int    // Number of rows to skip
+		Limit      int    // Number of rows to fetch
+		SortColumn string // Column to sort by
+		SortOrder  string // Sort direction (ASC, DESC)
+	}
+
 	Pagination struct {
 		Rows    int64 `json:"rows_count"`
 		Page    int64 `json:"page"`
@@ -94,16 +104,15 @@ func (res *Result) PostProcess() {
 }
 
 func (res *Result) Format() []map[string]interface{} {
-	var items []map[string]interface{}
+	items := make([]map[string]interface{}, len(res.Rows))
 
-	for _, row := range res.Rows {
+	for rowIdx, row := range res.Rows {
 		item := make(map[string]interface{})
-
 		for i, c := range res.Columns {
 			item[c] = row[i]
 		}
 
-		items = append(items, item)
+		items[rowIdx] = item
 	}
 
 	return items
@@ -121,20 +130,17 @@ func (res *Result) CSV() []byte {
 		record := make([]string, len(res.Columns))
 
 		for i, item := range row {
-			if item != nil {
-				switch v := item.(type) {
-				case time.Time:
-					record[i] = v.Format("2006-01-02 15:04:05")
-				default:
-					record[i] = fmt.Sprintf("%v", item)
-				}
-			} else {
+			switch v := item.(type) {
+			case time.Time:
+				record[i] = v.Format("2006-01-02 15:04:05")
+			case nil:
 				record[i] = ""
+			default:
+				record[i] = fmt.Sprintf("%v", item)
 			}
 		}
 
 		err := writer.Write(record)
-
 		if err != nil {
 			fmt.Println(err)
 			break
diff --git a/pkg/client/result_test.go b/pkg/client/result_test.go
index 7919908..4dfa1b0 100644
--- a/pkg/client/result_test.go
+++ b/pkg/client/result_test.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"encoding/json"
+	"strings"
 	"testing"
 	"time"
 
@@ -51,17 +52,20 @@ func TestPostProcess(t *testing.T) {
 
 func TestCSV(t *testing.T) {
 	result := Result{
-		Columns: []string{"id", "name", "email"},
+		Columns: []string{"id", "name", "email", "extra"},
 		Rows: []Row{
-			{1, "John", "john@example.com"},
-			{2, "Bob", "bob@example.com"},
+			{1, "John", "john@example.com", "data"},
+			{2, "Bob", "bob@example.com", nil},
 		},
 	}
 
-	expected := "id,name,email\n1,John,john@example.com\n2,Bob,bob@example.com\n"
-	output := string(result.CSV())
+	expected := strings.Join([]string{
+		"id,name,email,extra",
+		"1,John,john@example.com,data",
+		"2,Bob,bob@example.com,",
+	}, "\n") + "\n"
 
-	assert.Equal(t, expected, output)
+	assert.Equal(t, expected, string(result.CSV()))
 }
 
 func TestJSON(t *testing.T) {
@@ -116,3 +120,20 @@ func TestJSON(t *testing.T) {
 		assert.Equal(t, `[{"value":"2022-01-01T00:00:00Z"},{"value":"9022-01-01T00:00:00Z"},{"value":"ERR: INVALID_DATE"}]`, string(result.JSON()))
 	})
 }
+
+func TestResultFormat(t *testing.T) {
+	result := Result{
+		Columns: []string{"col1", "col2", "col3", "col4"},
+		Rows: []Row{
+			{"1", "2", "3", nil},
+			{"4", "5", "6", nil},
+		},
+	}
+
+	expected := []map[string]interface{}{
+		{"col1": "1", "col2": "2", "col3": "3", "col4": nil},
+		{"col1": "4", "col2": "5", "col3": "6", "col4": nil},
+	}
+
+	assert.Equal(t, expected, result.Format())
+}