Better s3 async bucket loading

README.md

@@ -36,7 +36,7 @@ go build
 
 # Usage
 
-Make sure you have AWS credentials set. Press `?` to toggle help$.
+Make sure you have AWS credentials set. Press `?` to toggle help.
 
 To switch services press `s`. You can launch directly to a particular service like
 

data/s3.go

@@ -6,7 +6,6 @@ import (
 	"log"
 	"strconv"
 	"strings"
-	"sync"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
@@ -27,52 +26,31 @@ func NewS3Client(ctx context.Context, s3 *s3.Client) *S3Client {
 	}
 }
 
-func (c *S3Client) GetBuckets(nextToken *string) ([]table.Row, *string, error) {
+func (c *S3Client) GetBuckets() ([]table.Row, error) {
 	input := s3.ListBucketsInput{}
 	output, err := c.s3.ListBuckets(c.ctx, &input)
 	if err != nil {
 		log.Fatalf("unable to get databases: %v", err)
 	}
 
-	var wg sync.WaitGroup
-	regionsCh := make(chan struct {
-		int
-		string
-	}, len(output.Buckets))
-
 	var rows []table.Row
-	for i, b := range output.Buckets {
-		wg.Add(1)
-		go func(rowIndex int, bucket string) {
-			defer wg.Done()
-
-			region, err := manager.GetBucketRegion(c.ctx, c.s3, bucket)
-			regionsCh <- struct {
-				int
-				string
-			}{rowIndex, region}
-
-			if err != nil {
-				log.Fatalf("Error getting region for bucket %s, %v", bucket, err)
-			}
-		}(i, aws.ToString(b.Name))
-
+	for _, b := range output.Buckets {
 		rows = append(rows, table.Row{
 			aws.ToString(b.Name),
-			"...",
+			LOADING_ALIAS,
 			formatTime(b.CreationDate),
 		})
 	}
 
-	go func() {
-		wg.Wait()
-		close(regionsCh)
-	}()
-	for item := range regionsCh {
-		rows[item.int][1] = item.string
-	}
+	return rows, nil
+}
 
-	return rows, nil, nil
+func (c *S3Client) GetBucketRegion(bucket string) (string, error) {
+	region, err := manager.GetBucketRegion(c.ctx, c.s3, bucket)
+	if err != nil {
+		log.Fatalf("Error getting region for bucket %s, %v", bucket, err)
+	}
+	return region, nil
 }
 
 func (c *S3Client) GetBucketPolicy(bucket string, region string) (string, error) {

data/utils.go

@@ -9,6 +9,8 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/cloudwatch/types"
 )
 
+const LOADING_ALIAS = "..."
+
 func statisticOfDatapoint(datapoint types.Datapoint, statistic types.Statistic) *float64 {
 	switch statistic {
 	case types.StatisticAverage:

ui/components/page/page.go

@@ -29,6 +29,13 @@ type BatchedNewRowsMsg struct {
 	Msgs []NewRowsMsg
 }
 
+type UpdateRowMsg struct {
+	Page            string
+	PaneId          int
+	Row             table.Row
+	PrimaryKeyIndex int
+}
+
 type ChangePageMsg struct {
 	NewPage     string // Id of page to switch to
 	FetchData   bool   // If true, clears and (re)fetches data on new page
@@ -64,6 +71,7 @@ type Page interface {
 	SetPageContext(context interface{})
 	GetSpec() PageSpec
 
+	GetPaneAt(index int) pane.Pane
 	GetCurrentPaneId() int
 
 	Inspect(client data.Client) tea.Cmd
@@ -219,6 +227,10 @@ func (m *Model) Update(client data.Client, msg tea.Msg) (tea.Cmd, bool) {
 	return tea.Batch(cmds...), false
 }
 
+func (m *Model) GetPaneAt(index int) pane.Pane {
+	return m.Panes[index]
+}
+
 func (m *Model) GetCurrentPaneId() int {
 	return m.Tabs.CurrentTabId
 }

ui/components/table/table.go

@@ -174,6 +174,22 @@ func (m *Model) GetMarshalledRow() map[string]string {
 	return rowMap
 }
 
+func (m *Model) MarhsalRow(row Row) map[string]string {
+	rowMap := make(map[string]string)
+	for i, col := range m.Columns {
+		rowMap[col.Title] = row[i]
+	}
+	return rowMap
+}
+
+func (m *Model) UnmarhsalRow(row map[string]string) Row {
+	rowArr := make(Row, len(m.Columns))
+	for i, col := range m.Columns {
+		rowArr[i] = row[col.Title]
+	}
+	return rowArr
+}
+
 func (m *Model) nextCol() int {
 	m.Columns[m.currColumnId].isSelected = utils.BoolPtr(false)
 	m.currColumnId = (m.currColumnId + 1) % len(m.Columns)
@@ -220,6 +236,19 @@ func (m *Model) AppendRows(rows []Row) {
 	}
 }
 
+// Updates a row using a given column as the primary key
+func (m *Model) UpdateRow(primaryKeyIndex int, newRow Row) {
+	newRows := make([]Row, 0, len(m.rows))
+	for _, r := range m.rows {
+		if r[primaryKeyIndex] == newRow[primaryKeyIndex] {
+			newRows = append(newRows, newRow)
+		} else {
+			newRows = append(newRows, r)
+		}
+	}
+	m.SetRows(newRows)
+}
+
 func (m *Model) ClearRows() {
 	m.rows = make([]Row, 0)
 	m.filterRows()

ui/pages/s3/s3.go

@@ -22,25 +22,53 @@ func NewS3Page(ctx *context.ProgramContext) *S3PageModel {
 
 func (m *S3PageModel) FetchData(client data.Client) tea.Cmd {
 	return tea.Batch(
-		m.fetchBuckets(client, nil),
+		m.fetchBuckets(client),
 	)
 }
 
-func (m *S3PageModel) fetchBuckets(client data.Client, nextToken *string) tea.Cmd {
+func (m *S3PageModel) fetchBuckets(client data.Client) tea.Cmd {
 	return func() tea.Msg {
-		rows, nextToken, _ := client.S3.GetBuckets(nextToken)
-		msg := page.NewRowsMsg{
-			Page:   m.Spec.Name,
-			PaneId: m.GetPaneId("Buckets"),
-			Rows:   rows,
+		var nextCmds []tea.Cmd
+
+		rows, _ := client.S3.GetBuckets()
+		for _, row := range rows {
+			nextCmds = append(nextCmds, m.fetchBucketRegion(client, row))
 		}
-		if nextToken != nil {
-			msg.NextCmd = m.fetchBuckets(client, nextToken)
+
+		msg := page.NewRowsMsg{
+			Page:    m.Spec.Name,
+			PaneId:  m.GetPaneId("Buckets"),
+			Rows:    rows,
+			NextCmd: tea.Batch(nextCmds...),
 		}
 		return msg
 	}
 }
 
+func (m *S3PageModel) fetchBucketRegion(client data.Client, row table.Row) tea.Cmd {
+	table, ok := m.CurrentPane().(*table.Model)
+	if !ok {
+		log.Fatal("This pane is not a table")
+	}
+
+	return func() tea.Msg {
+		region, err := client.S3.GetBucketRegion(row[0])
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		marshalledRow := table.MarhsalRow(row)
+		marshalledRow["Region"] = region
+
+		return page.UpdateRowMsg{
+			Page:            m.Spec.Name,
+			PaneId:          m.GetPaneId("Buckets"),
+			Row:             table.UnmarhsalRow(marshalledRow),
+			PrimaryKeyIndex: 0,
+		}
+	}
+}
+
 func (m *S3PageModel) Inspect(client data.Client) tea.Cmd {
 	if m.Tabs.CurrentTabId != m.GetPaneId("Buckets") {
 		return nil
@@ -50,7 +78,16 @@ func (m *S3PageModel) Inspect(client data.Client) tea.Cmd {
 	if !ok {
 		log.Fatal("This pane is not a table")
 	}
+
 	row := table.GetMarshalledRow()
+
+	// If the region hasn't been found yet, we need to wait.
+	// An alternative idea might be to just fetch the region in GetObjects if it's missing at that
+	// point
+	if row["Region"] == data.LOADING_ALIAS {
+		return nil
+	}
+
 	changePageCmd := func() tea.Msg {
 		return page.ChangePageMsg{
 			NewPage:   "s3/objects",

ui/ui.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/danielcmessias/sawsy/ui/components/help"
 	"github.com/danielcmessias/sawsy/ui/components/page"
+	"github.com/danielcmessias/sawsy/ui/components/table"
 	"github.com/danielcmessias/sawsy/ui/context"
 	"github.com/danielcmessias/sawsy/ui/pages/glue"
 	"github.com/danielcmessias/sawsy/ui/pages/iam"
@@ -139,6 +140,9 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case page.NewRowsMsg:
 		cmds = append(cmds, m.parseNewRowsMsg(msg))
 
+	case page.UpdateRowMsg:
+		m.parseUpdateRowMsg(msg)
+
 	case page.BatchedNewRowsMsg:
 		for _, _msg := range msg.Msgs {
 			cmds = append(cmds, m.parseNewRowsMsg(_msg))
@@ -204,6 +208,14 @@ func (m *Model) parseNewRowsMsg(msg page.NewRowsMsg) tea.Cmd {
 	return tea.Batch(cmds...)
 }
 
+func (m *Model) parseUpdateRowMsg(msg page.UpdateRowMsg) {
+	table, ok := m.pages[msg.Page].GetPaneAt(msg.PaneId).(*table.Model)
+	if !ok {
+		log.Fatal("This pane is not a table")
+	}
+	table.UpdateRow(msg.PrimaryKeyIndex, msg.Row)
+}
+
 func (m *Model) changePage(pageName string, context interface{}, fetchData bool) tea.Cmd {
 	_, ok := m.pages[pageName]
 	if !ok {