sqs-browse: fixed assumption regarding table keys

internal/dynamo-browse/controllers/tableread.go

@@ -3,18 +3,20 @@ package controllers
 import (
 	"context"
 	"github.com/lmika/awstools/internal/common/ui/uimodels"
+	"github.com/lmika/awstools/internal/dynamo-browse/models"
 	"github.com/lmika/awstools/internal/dynamo-browse/services/tables"
+	"github.com/pkg/errors"
 )
 
 type TableReadController struct {
 	tableService *tables.Service
-	tableName string
+	tableName    string
 }
 
 func NewTableReadController(tableService *tables.Service, tableName string) *TableReadController {
 	return &TableReadController{
 		tableService: tableService,
-		tableName: tableName,
+		tableName:    tableName,
 	}
 }
 
@@ -24,14 +26,19 @@ func (c *TableReadController) Scan() uimodels.Operation {
 	})
 }
 
-func (c *TableReadController) doScan(ctx context.Context, quiet bool) error {
+func (c *TableReadController) doScan(ctx context.Context, quiet bool) (err error) {
 	uiCtx := uimodels.Ctx(ctx)
 
 	if !quiet {
 		uiCtx.Message("Scanning...")
 	}
 
-	resultSet, err := c.tableService.Scan(ctx, c.tableName)
+	tableInfo, err := c.tableInfo(ctx)
+	if err != nil {
+		return err
+	}
+
+	resultSet, err := c.tableService.Scan(ctx, tableInfo)
 	if err != nil {
 		return err
 	}
@@ -41,4 +48,18 @@ func (c *TableReadController) doScan(ctx context.Context, quiet bool) error {
 	}
 	uiCtx.Send(NewResultSet{resultSet})
 	return nil
-}
+}
+
+// tableInfo returns the table info from the state if a result set exists.  If not, it fetches the
+// table information from the service.
+func (c *TableReadController) tableInfo(ctx context.Context) (*models.TableInfo, error) {
+	if existingResultSet := CurrentState(ctx).ResultSet; existingResultSet != nil {
+		return existingResultSet.TableInfo, nil
+	}
+
+	tableInfo, err := c.tableService.Describe(ctx, c.tableName)
+	if err != nil {
+		return nil, errors.Wrapf(err, "cannot describe %v", c.tableName)
+	}
+	return tableInfo, nil
+}

internal/dynamo-browse/controllers/tablewrite.go

@@ -69,8 +69,13 @@ func (c *TableWriteController) Duplicate() uimodels.Operation {
 					return errors.New("operation aborted")
 				}
 
+				tableInfo, err := c.tableReadControllers.tableInfo(ctx)
+				if err != nil {
+					return err
+				}
+
 				// Delete the item
-				if err := c.tableService.Put(ctx, c.tableName, newItem); err != nil {
+				if err := c.tableService.Put(ctx, tableInfo, newItem); err != nil {
 					return err
 				}
 
@@ -105,8 +110,13 @@ func (c *TableWriteController) Delete() uimodels.Operation {
 				return errors.New("operation aborted")
 			}
 
+			tableInfo, err := c.tableReadControllers.tableInfo(ctx)
+			if err != nil {
+				return err
+			}
+
 			// Delete the item
-			if err := c.tableService.Delete(ctx, c.tableName, state.SelectedItem); err != nil {
+			if err := c.tableService.Delete(ctx, tableInfo, state.SelectedItem); err != nil {
 				return err
 			}
 

internal/dynamo-browse/models/models.go

@@ -3,9 +3,9 @@ package models
 import "github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
 
 type ResultSet struct {
-	Table   string
+	TableInfo *TableInfo
-	Columns []string
+	Columns   []string
-	Items   []Item
+	Items     []Item
 }
 
 type Item map[string]types.AttributeValue
@@ -21,3 +21,12 @@ func (i Item) Clone() Item {
 
 	return newItem
 }
+
+func (i Item) KeyValue(info *TableInfo) map[string]types.AttributeValue {
+	itemKey := make(map[string]types.AttributeValue)
+	itemKey[info.Keys.PartitionKey] = i[info.Keys.PartitionKey]
+	if info.Keys.SortKey != "" {
+		itemKey[info.Keys.SortKey] = i[info.Keys.SortKey]
+	}
+	return itemKey
+}

internal/dynamo-browse/models/sorted.go

@@ -5,13 +5,13 @@ import "sort"
 // sortedItems is a collection of items that is sorted.
 // Items are sorted based on the PK, and SK in ascending order
 type sortedItems struct {
-	pk, sk string
+	tableInfo *TableInfo
-	items  []Item
+	items     []Item
 }
 
 // Sort sorts the items in place
-func Sort(items []Item, pk, sk string) {
+func Sort(items []Item, tableInfo *TableInfo) {
-	si := sortedItems{items: items, pk: pk, sk: sk}
+	si := sortedItems{items: items, tableInfo: tableInfo}
 	sort.Sort(&si)
 }
 
@@ -21,7 +21,7 @@ func (si *sortedItems) Len() int {
 
 func (si *sortedItems) Less(i, j int) bool {
 	// Compare primary keys
-	pv1, pv2 := si.items[i][si.pk], si.items[j][si.pk]
+	pv1, pv2 := si.items[i][si.tableInfo.Keys.PartitionKey], si.items[j][si.tableInfo.Keys.PartitionKey]
 	pc, ok := compareScalarAttributes(pv1, pv2)
 	if !ok {
 		return i < j
@@ -34,16 +34,18 @@ func (si *sortedItems) Less(i, j int) bool {
 	}
 
 	// Partition keys are equal, compare sort key
-	sv1, sv2 := si.items[i][si.sk], si.items[j][si.sk]
+	if sortKey := si.tableInfo.Keys.SortKey; sortKey != "" {
-	sc, ok := compareScalarAttributes(sv1, sv2)
+		sv1, sv2 := si.items[i][sortKey], si.items[j][sortKey]
-	if !ok {
+		sc, ok := compareScalarAttributes(sv1, sv2)
-		return i < j
+		if !ok {
-	}
+			return i < j
+		}
 
-	if sc < 0 {
+		if sc < 0 {
-		return true
+			return true
-	} else if sc > 0 {
+		} else if sc > 0 {
-		return false
+			return false
+		}
 	}
 
 	// This should never happen, but just in case

internal/dynamo-browse/models/tableinfo.go

@@ -0,0 +1,12 @@
+package models
+
+type TableInfo struct {
+	Name              string
+	Keys              KeyAttribute
+	DefinedAttributes []string
+}
+
+type KeyAttribute struct {
+	PartitionKey string
+	SortKey      string
+}

internal/dynamo-browse/providers/dynamo/provider.go

@@ -13,6 +13,32 @@ type Provider struct {
 	client *dynamodb.Client
 }
 
+func (p *Provider) DescribeTable(ctx context.Context, tableName string) (*models.TableInfo, error) {
+	out, err := p.client.DescribeTable(ctx, &dynamodb.DescribeTableInput{
+		TableName: aws.String(tableName),
+	})
+	if err != nil {
+		return nil, errors.Wrapf(err, "cannot describe table %v", tableName)
+	}
+
+	var tableInfo models.TableInfo
+	tableInfo.Name = aws.ToString(out.Table.TableName)
+
+	for _, keySchema := range out.Table.KeySchema {
+		if keySchema.KeyType == types.KeyTypeHash {
+			tableInfo.Keys.PartitionKey = aws.ToString(keySchema.AttributeName)
+		} else if keySchema.KeyType == types.KeyTypeRange {
+			tableInfo.Keys.SortKey = aws.ToString(keySchema.AttributeName)
+		}
+	}
+
+	for _, definedAttribute := range out.Table.AttributeDefinitions {
+		tableInfo.DefinedAttributes = append(tableInfo.DefinedAttributes, aws.ToString(definedAttribute.AttributeName))
+	}
+
+	return &tableInfo, nil
+}
+
 func (p *Provider) PutItem(ctx context.Context, name string, item models.Item) error {
 	_, err := p.client.PutItem(ctx, &dynamodb.PutItemInput{
 		TableName: aws.String(name),

internal/dynamo-browse/services/tables/iface.go

@@ -7,6 +7,7 @@ import (
 )
 
 type TableProvider interface {
+	DescribeTable(ctx context.Context, tableName string) (*models.TableInfo, error)
 	ScanItems(ctx context.Context, tableName string) ([]models.Item, error)
 	DeleteItem(ctx context.Context, tableName string, key map[string]types.AttributeValue) error
 	PutItem(ctx context.Context, name string, item models.Item) error

internal/dynamo-browse/services/tables/service.go

@@ -2,7 +2,6 @@ package tables
 
 import (
 	"context"
-	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
 	"github.com/lmika/awstools/internal/dynamo-browse/models"
 	"github.com/pkg/errors"
 	"sort"
@@ -18,24 +17,34 @@ func NewService(provider TableProvider) *Service {
 	}
 }
 
-func (s *Service) Scan(ctx context.Context, table string) (*models.ResultSet, error) {
+func (s *Service) Describe(ctx context.Context, table string) (*models.TableInfo, error) {
-	results, err := s.provider.ScanItems(ctx, table)
+	return s.provider.DescribeTable(ctx, table)
-	if err != nil {
+}
-		return nil, errors.Wrapf(err, "unable to scan table %v", table)
-	}
 
-	// TODO: need to get PKs and SKs from table
+func (s *Service) Scan(ctx context.Context, tableInfo *models.TableInfo) (*models.ResultSet, error) {
-	pk, sk := "pk", "sk"
+	results, err := s.provider.ScanItems(ctx, tableInfo.Name)
+	if err != nil {
+		return nil, errors.Wrapf(err, "unable to scan table %v", tableInfo.Name)
+	}
 
 	// Get the columns
 	seenColumns := make(map[string]int)
-	seenColumns[pk] = 0
+	seenColumns[tableInfo.Keys.PartitionKey] = 0
-	seenColumns[sk] = 1
+	if tableInfo.Keys.SortKey != "" {
+		seenColumns[tableInfo.Keys.SortKey] = 1
+	}
 
+	for _, definedAttribute := range tableInfo.DefinedAttributes {
+		if _, seen := seenColumns[definedAttribute]; !seen {
+			seenColumns[definedAttribute] = len(seenColumns)
+		}
+	}
+
+	otherColsRank := len(seenColumns)
 	for _, result := range results {
 		for k := range result {
 			if _, isSeen := seenColumns[k]; !isSeen {
-				seenColumns[k] = 2
+				seenColumns[k] = otherColsRank
 			}
 		}
 	}
@@ -51,23 +60,19 @@ func (s *Service) Scan(ctx context.Context, table string) (*models.ResultSet, er
 		return seenColumns[columns[i]] < seenColumns[columns[j]]
 	})
 
-	models.Sort(results, pk, sk)
+	models.Sort(results, tableInfo)
 
 	return &models.ResultSet{
-		Table:   table,
+		TableInfo:   tableInfo,
 		Columns: columns,
 		Items:   results,
 	}, nil
 }
 
-func (s *Service) Put(ctx context.Context, tableName string, item models.Item) error {
+func (s *Service) Put(ctx context.Context, tableInfo *models.TableInfo, item models.Item) error {
-	return s.provider.PutItem(ctx, tableName, item)
+	return s.provider.PutItem(ctx, tableInfo.Name, item)
 }
 
-func (s *Service) Delete(ctx context.Context, name string, item models.Item) error {
+func (s *Service) Delete(ctx context.Context, tableInfo *models.TableInfo, item models.Item) error {
-	// TODO: do not hardcode keys
+	return s.provider.DeleteItem(ctx, tableInfo.Name, item.KeyValue(tableInfo))
-	return s.provider.DeleteItem(ctx, name, map[string]types.AttributeValue{
-		"pk": item["pk"],
-		"sk": item["sk"],
-	})
 }

internal/dynamo-browse/services/tables/service_test.go

@@ -9,6 +9,28 @@ import (
 	"testing"
 )
 
+func TestService_Describe(t *testing.T) {
+	tableName := "service-describe-table"
+
+	client, cleanupFn := testdynamo.SetupTestTable(t, tableName, testData)
+	defer cleanupFn()
+	provider := dynamo.NewProvider(client)
+
+	t.Run("return details of the table", func(t *testing.T) {
+		ctx := context.Background()
+
+		service := tables.NewService(provider)
+		ti, err := service.Describe(ctx, tableName)
+		assert.NoError(t, err)
+
+		// Hash first, then range, then columns in alphabetic order
+		assert.Equal(t, ti.Name, tableName)
+		assert.Equal(t, "pk", ti.Keys.PartitionKey, "pk")
+		assert.Equal(t, "sk", ti.Keys.SortKey, "sk")
+		assert.Equal(t, []string{"pk", "sk"}, ti.DefinedAttributes)
+	})
+}
+
 func TestService_Scan(t *testing.T) {
 	tableName := "service-scan-test-table"
 
@@ -20,11 +42,14 @@ func TestService_Scan(t *testing.T) {
 		ctx := context.Background()
 
 		service := tables.NewService(provider)
-		rs, err := service.Scan(ctx, tableName)
+		ti, err := service.Describe(ctx, tableName)
+		assert.NoError(t, err)
+
+		rs, err := service.Scan(ctx, ti)
 		assert.NoError(t, err)
 
 		// Hash first, then range, then columns in alphabetic order
-		assert.Equal(t, rs.Table, tableName)
+		assert.Equal(t, rs.TableInfo, ti)
 		assert.Equal(t, rs.Columns, []string{"pk", "sk", "alpha", "beta", "gamma"})
 		assert.Equal(t, rs.Items[0], testdynamo.TestRecordAsItem(t, testData[1]))
 		assert.Equal(t, rs.Items[1], testdynamo.TestRecordAsItem(t, testData[0]))

internal/dynamo-browse/ui/model.go

@@ -71,6 +71,7 @@ func NewModel(dispatcher *dispatcher.Dispatcher, commandController *commandctrl.
 }
 
 func (m uiModel) Init() tea.Cmd {
+	m.invokeOperation(context.Background(), m.tableReadController.Scan())
 	return nil
 }
 
@@ -257,7 +258,7 @@ func (m uiModel) View() string {
 func (m uiModel) headerView() string {
 	var titleText string
 	if m.state.ResultSet != nil {
-		titleText = "Table: " + m.state.ResultSet.Table
+		titleText = "Table: " + m.state.ResultSet.TableInfo.Name
 	} else {
 		titleText = "No table"
 	}