From ab309084c5a61d7bb48c4a7355fc73d4ae38901d Mon Sep 17 00:00:00 2001
From: Leon Mika <lmika@lmika.org>
Date: Thu, 23 Feb 2023 21:31:01 +1100
Subject: [PATCH] Refreshing will update the columns if they haven't been
 rearranged

---
 internal/dynamo-browse/controllers/columns.go |  6 +++
 .../dynamo-browse/models/columns/columns.go   | 41 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/internal/dynamo-browse/controllers/columns.go b/internal/dynamo-browse/controllers/columns.go
index 53f3386..b58f97b 100644
--- a/internal/dynamo-browse/controllers/columns.go
+++ b/internal/dynamo-browse/controllers/columns.go
@@ -38,6 +38,7 @@ func (cc *ColumnsController) ShiftColumnLeft(idx int) tea.Msg {
 
 	col := cc.colModel.Columns[idx-1]
 	cc.colModel.Columns[idx-1], cc.colModel.Columns[idx] = cc.colModel.Columns[idx], col
+	cc.colModel.WasRearranged = true
 
 	return ColumnsUpdated{}
 }
@@ -49,6 +50,7 @@ func (cc *ColumnsController) ShiftColumnRight(idx int) tea.Msg {
 
 	col := cc.colModel.Columns[idx+1]
 	cc.colModel.Columns[idx+1], cc.colModel.Columns[idx] = cc.colModel.Columns[idx], col
+	cc.colModel.WasRearranged = true
 
 	return ColumnsUpdated{}
 }
@@ -63,6 +65,8 @@ func (cc *ColumnsController) onNewResultSet(rs *models.ResultSet, op resultSetUp
 
 	if cc.colModel == nil || (op == resultSetUpdateInit || op == resultSetUpdateQuery) {
 		cc.colModel = columns.NewColumnsFromResultSet(rs)
+	} else {
+		cc.colModel.AddMissingColumns(rs)
 	}
 }
 
@@ -89,6 +93,7 @@ func (cc *ColumnsController) AddColumn(afterIndex int) tea.Msg {
 
 			cc.colModel.Columns = newCols
 		}
+		cc.colModel.WasRearranged = true
 
 		return tea.Batch(
 			events.SetTeaMessage(ColumnsUpdated{}),
@@ -106,6 +111,7 @@ func (cc *ColumnsController) DeleteColumn(afterIndex int) tea.Msg {
 	newCols = append(newCols, cc.colModel.Columns[:afterIndex]...)
 	newCols = append(newCols, cc.colModel.Columns[afterIndex+1:]...)
 	cc.colModel.Columns = newCols
+	cc.colModel.WasRearranged = true
 
 	return ColumnsUpdated{}
 }
diff --git a/internal/dynamo-browse/models/columns/columns.go b/internal/dynamo-browse/models/columns/columns.go
index eb87711..c24fe1d 100644
--- a/internal/dynamo-browse/models/columns/columns.go
+++ b/internal/dynamo-browse/models/columns/columns.go
@@ -7,8 +7,9 @@ import (
 )
 
 type Columns struct {
-	TableInfo *models.TableInfo
-	Columns   []Column
+	TableInfo     *models.TableInfo
+	WasRearranged bool
+	Columns       []Column
 }
 
 func NewColumnsFromResultSet(rs *models.ResultSet) *Columns {
@@ -28,6 +29,42 @@ func NewColumnsFromResultSet(rs *models.ResultSet) *Columns {
 	}
 }
 
+func (cols *Columns) AddMissingColumns(rs *models.ResultSet) {
+	existingColumns := make(map[string]Column)
+	for _, col := range cols.Columns {
+		existingColumns[col.Name] = col
+	}
+
+	rsCols := rs.Columns()
+	var newCols []Column
+
+	if cols.WasRearranged {
+		newCols = append([]Column{}, cols.Columns...)
+		for _, c := range rsCols {
+			if _, hasCol := existingColumns[c]; !hasCol {
+				newCols = append(newCols, Column{
+					Name:      c,
+					Evaluator: SimpleFieldValueEvaluator(c),
+				})
+			}
+		}
+	} else {
+		newCols = make([]Column, len(rsCols))
+		for i, c := range rsCols {
+			if existingCol, hasCol := existingColumns[c]; hasCol {
+				newCols[i] = existingCol
+			} else {
+				newCols[i] = Column{
+					Name:      c,
+					Evaluator: SimpleFieldValueEvaluator(c),
+				}
+			}
+		}
+	}
+
+	cols.Columns = newCols
+}
+
 func (cols *Columns) VisibleColumns() []Column {
 	if cols == nil {
 		return []Column{}