From 12163483dd2dee57f4512cec83eea5f780220287 Mon Sep 17 00:00:00 2001
From: David Dworken <david@daviddworken.com>
Date: Sun, 23 Oct 2022 21:42:22 -0700
Subject: [PATCH] Add data types for custom columns

---
 client/data/data.go      | 40 +++++++++++++++++++++++++++++++---------
 client/data/data_test.go | 20 ++++++++++++++++++++
 client/hctx/hctx.go      |  7 +++++++
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/client/data/data.go b/client/data/data.go
index 6afc495..8fafc91 100644
--- a/client/data/data.go
+++ b/client/data/data.go
@@ -6,6 +6,7 @@ import (
 	"crypto/hmac"
 	"crypto/rand"
 	"crypto/sha256"
+	"database/sql/driver"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -25,15 +26,36 @@ const (
 )
 
 type HistoryEntry struct {
-	LocalUsername           string    `json:"local_username" gorm:"uniqueIndex:compositeindex"`
-	Hostname                string    `json:"hostname" gorm:"uniqueIndex:compositeindex"`
-	Command                 string    `json:"command" gorm:"uniqueIndex:compositeindex"`
-	CurrentWorkingDirectory string    `json:"current_working_directory" gorm:"uniqueIndex:compositeindex"`
-	HomeDirectory           string    `json:"home_directory" gorm:"uniqueIndex:compositeindex"`
-	ExitCode                int       `json:"exit_code" gorm:"uniqueIndex:compositeindex"`
-	StartTime               time.Time `json:"start_time" gorm:"uniqueIndex:compositeindex"`
-	EndTime                 time.Time `json:"end_time" gorm:"uniqueIndex:compositeindex"`
-	DeviceId                string    `json:"device_id" gorm:"uniqueIndex:compositeindex"`
+	LocalUsername           string        `json:"local_username" gorm:"uniqueIndex:compositeindex"`
+	Hostname                string        `json:"hostname" gorm:"uniqueIndex:compositeindex"`
+	Command                 string        `json:"command" gorm:"uniqueIndex:compositeindex"`
+	CurrentWorkingDirectory string        `json:"current_working_directory" gorm:"uniqueIndex:compositeindex"`
+	HomeDirectory           string        `json:"home_directory" gorm:"uniqueIndex:compositeindex"`
+	ExitCode                int           `json:"exit_code" gorm:"uniqueIndex:compositeindex"`
+	StartTime               time.Time     `json:"start_time" gorm:"uniqueIndex:compositeindex"`
+	EndTime                 time.Time     `json:"end_time" gorm:"uniqueIndex:compositeindex"`
+	DeviceId                string        `json:"device_id" gorm:"uniqueIndex:compositeindex"`
+	CustomColumns           CustomColumns `json:"custom_columns"`
+}
+
+type CustomColumns []CustomColumns
+
+type CustomColumn struct {
+	Name string `json:"name"`
+	Val  string `json:"value"`
+}
+
+func (c *CustomColumns) Scan(value interface{}) error {
+	bytes, ok := value.([]byte)
+	if !ok {
+		return fmt.Errorf("failed to unmarshal CustomColumns value %#v", value)
+	}
+
+	return json.Unmarshal(bytes, c)
+}
+
+func (c CustomColumns) Value() (driver.Value, error) {
+	return json.Marshal(c)
 }
 
 func (h *HistoryEntry) GoString() string {
diff --git a/client/data/data_test.go b/client/data/data_test.go
index 71ceb48..de0d9e2 100644
--- a/client/data/data_test.go
+++ b/client/data/data_test.go
@@ -70,3 +70,23 @@ func TestParseTimeGenerously(t *testing.T) {
 		t.Fatalf("parsed time incorrectly: %d", ts.Unix())
 	}
 }
+
+func TestCustomColumnSerialization(t *testing.T) {
+	// cc1 := CustomColumn{
+	// 	Name: "name1",
+	// 	Val:  "val1",
+	// }
+	// cc2 := CustomColumn{
+	// 	Name: "name2",
+	// 	Val:  "val2",
+	// }
+	var ccs CustomColumns = make(CustomColumns, 0)
+	v, err := ccs.Value()
+	if err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+	val := string(v.([]uint8))
+	if val != "[]" {
+		t.Fatalf("unexpected val for empty CustomColumns: %#v", val)
+	}
+}
diff --git a/client/hctx/hctx.go b/client/hctx/hctx.go
index a34f960..11d004c 100644
--- a/client/hctx/hctx.go
+++ b/client/hctx/hctx.go
@@ -160,6 +160,13 @@ type ClientConfig struct {
 	ControlRSearchEnabled bool `json:"enable_control_r_search"`
 	// The set of columns that the user wants to be displayed
 	DisplayedColumns []string `json:"displayed_columns"`
+	// Custom columns
+	CustomColumns []CustomColumnDefinition `json:"custom_columns"`
+}
+
+type CustomColumnDefinition struct {
+	ColumnName    string `json:"column_name"`
+	ColumnCommand string `json:"column_command"`
 }
 
 func GetConfigContents() ([]byte, error) {