From ed53173a1d3e587383bb9055c2144042c24cddac Mon Sep 17 00:00:00 2001
From: Leon Mika <lmika@lmika.org>
Date: Mon, 31 Jul 2023 20:59:05 +1000
Subject: [PATCH] Added the "export -all" switch (#54)
Extended the "export" command with an "-all" flag. When included, all rows of the table matching the query will be exported to CSV.
---
cmd/dynamo-browse/main.go | 2 +-
internal/common/sliceutils/map.go | 8 ++
internal/dynamo-browse/controllers/export.go | 94 ++++++++++++-------
.../dynamo-browse/controllers/export_test.go | 64 ++++++++++++-
.../dynamo-browse/controllers/jobbuilder.go | 3 +
.../dynamo-browse/controllers/tableread.go | 2 +-
.../controllers/tableread_test.go | 6 +-
.../controllers/tablewrite_test.go | 2 +-
internal/dynamo-browse/models/models.go | 4 +
internal/dynamo-browse/ui/model.go | 9 +-
10 files changed, 150 insertions(+), 44 deletions(-)
diff --git a/cmd/dynamo-browse/main.go b/cmd/dynamo-browse/main.go
index 1e0b05d..170a2dc 100644
--- a/cmd/dynamo-browse/main.go
+++ b/cmd/dynamo-browse/main.go
@@ -108,7 +108,7 @@ func main() {
tableReadController := controllers.NewTableReadController(state, tableService, workspaceService, itemRendererService, jobsController, inputHistoryService, eventBus, *flagTable)
tableWriteController := controllers.NewTableWriteController(state, tableService, jobsController, tableReadController, settingStore)
columnsController := controllers.NewColumnsController(eventBus)
- exportController := controllers.NewExportController(state, columnsController)
+ exportController := controllers.NewExportController(state, tableService, jobsController, columnsController)
settingsController := controllers.NewSettingsController(settingStore, eventBus)
keyBindings := keybindings.Default()
scriptController := controllers.NewScriptController(scriptManagerService, tableReadController, settingsController, eventBus)
diff --git a/internal/common/sliceutils/map.go b/internal/common/sliceutils/map.go
index ae2a1be..0864a56 100644
--- a/internal/common/sliceutils/map.go
+++ b/internal/common/sliceutils/map.go
@@ -9,6 +9,14 @@ func All[T any](ts []T, predicate func(t T) bool) bool {
return true
}
+func Generate[U any](from, to int, fn func(t int) U) []U {
+ us := make([]U, to-from+1)
+ for i := from; i <= to; i++ {
+ us[i-from] = fn(i)
+ }
+ return us
+}
+
func Map[T, U any](ts []T, fn func(t T) U) []U {
us := make([]U, len(ts))
for i, t := range ts {
diff --git a/internal/dynamo-browse/controllers/export.go b/internal/dynamo-browse/controllers/export.go
index 3f70c93..15a86c6 100644
--- a/internal/dynamo-browse/controllers/export.go
+++ b/internal/dynamo-browse/controllers/export.go
@@ -1,57 +1,85 @@
package controllers
import (
+ "context"
"encoding/csv"
+ "fmt"
tea "github.com/charmbracelet/bubbletea"
"github.com/lmika/dynamo-browse/internal/common/ui/events"
"github.com/lmika/dynamo-browse/internal/dynamo-browse/models/attrutils"
+ "github.com/lmika/dynamo-browse/internal/dynamo-browse/services/jobs"
"github.com/pkg/errors"
"os"
)
type ExportController struct {
- state *State
- columns *ColumnsController
+ state *State
+ tableService TableReadService
+ jobController *JobsController
+ columns *ColumnsController
}
-func NewExportController(state *State, columns *ColumnsController) *ExportController {
- return &ExportController{state, columns}
+func NewExportController(state *State, tableService TableReadService, jobsController *JobsController, columns *ColumnsController) *ExportController {
+ return &ExportController{state, tableService, jobsController, columns}
}
-func (c *ExportController) ExportCSV(filename string) tea.Msg {
+func (c *ExportController) ExportCSV(filename string, opts ExportOptions) tea.Msg {
resultSet := c.state.ResultSet()
if resultSet == nil {
return events.Error(errors.New("no result set"))
}
- f, err := os.Create(filename)
- if err != nil {
- return events.Error(errors.Wrapf(err, "cannot export to '%v'", filename))
- }
- defer f.Close()
-
- cw := csv.NewWriter(f)
- defer cw.Flush()
-
- columns := c.columns.Columns().VisibleColumns()
-
- colNames := make([]string, len(columns))
- for i, c := range columns {
- colNames[i] = c.Name
- }
- if err := cw.Write(colNames); err != nil {
- return events.Error(errors.Wrapf(err, "cannot export to '%v'", filename))
- }
-
- row := make([]string, len(columns))
- for _, item := range resultSet.Items() {
- for i, col := range columns {
- row[i], _ = attrutils.AttributeToString(col.Evaluator.EvaluateForItem(item))
+ return NewJob(c.jobController, fmt.Sprintf("Exporting to %v…", filename), func(ctx context.Context) (int, error) {
+ f, err := os.Create(filename)
+ if err != nil {
+ return 0, errors.Wrapf(err, "cannot export to '%v'", filename)
}
- if err := cw.Write(row); err != nil {
- return events.Error(errors.Wrapf(err, "cannot export to '%v'", filename))
- }
- }
+ defer f.Close()
- return nil
+ cw := csv.NewWriter(f)
+ defer cw.Flush()
+
+ columns := c.columns.Columns().VisibleColumns()
+
+ colNames := make([]string, len(columns))
+ for i, c := range columns {
+ colNames[i] = c.Name
+ }
+ if err := cw.Write(colNames); err != nil {
+ return 0, errors.Wrapf(err, "cannot export to '%v'", filename)
+ }
+
+ totalRows := 0
+ row := make([]string, len(columns))
+ for {
+ for _, item := range resultSet.Items() {
+ for i, col := range columns {
+ row[i], _ = attrutils.AttributeToString(col.Evaluator.EvaluateForItem(item))
+ }
+ if err := cw.Write(row); err != nil {
+ return 0, errors.Wrapf(err, "cannot export to '%v'", filename)
+ }
+ }
+ totalRows += len(resultSet.Items())
+
+ if !opts.AllResults || !resultSet.HasNextPage() {
+ break
+ }
+
+ jobs.PostUpdate(ctx, fmt.Sprintf("exported %d items", totalRows))
+ resultSet, err = c.tableService.NextPage(ctx, resultSet)
+ if err != nil {
+ return 0, errors.Wrapf(err, "cannot get next page while exporting to '%v'", filename)
+ }
+ }
+
+ return totalRows, nil
+ }).OnDone(func(rows int) tea.Msg {
+ return events.StatusMsg(applyToN("Exported ", rows, "item", "items", " to "+filename))
+ }).Submit()
+}
+
+type ExportOptions struct {
+ // AllResults returns all results from the table
+ AllResults bool
}
diff --git a/internal/dynamo-browse/controllers/export_test.go b/internal/dynamo-browse/controllers/export_test.go
index f70a387..ce62fc3 100644
--- a/internal/dynamo-browse/controllers/export_test.go
+++ b/internal/dynamo-browse/controllers/export_test.go
@@ -1,6 +1,9 @@
package controllers_test
import (
+ "fmt"
+ "github.com/lmika/dynamo-browse/internal/common/sliceutils"
+ "github.com/lmika/dynamo-browse/internal/dynamo-browse/controllers"
"github.com/stretchr/testify/assert"
"os"
"strings"
@@ -14,7 +17,7 @@ func TestExportController_ExportCSV(t *testing.T) {
tempFile := tempFile(t)
invokeCommand(t, srv.readController.Init())
- invokeCommand(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
bts, err := os.ReadFile(tempFile)
assert.NoError(t, err)
@@ -27,13 +30,66 @@ func TestExportController_ExportCSV(t *testing.T) {
}, ""))
})
+ t.Run("should export all pages of the results", func(t *testing.T) {
+ pageLimits := []int{5, 10, 50}
+
+ for _, pageLimit := range pageLimits {
+ t.Run(fmt.Sprintf("page size %d", pageLimit), func(t *testing.T) {
+ t.Run("all results", func(t *testing.T) {
+ srv := newService(t, serviceConfig{tableName: "count-to-30", defaultLimit: 5})
+
+ tempFile := tempFile(t)
+
+ expected := append([]string{
+ "pk,sk,num\n",
+ }, sliceutils.Generate(1, 30, func(i int) string {
+ return fmt.Sprintf("NUM,NUM#%02d,%d\n", i, i)
+ })...)
+
+ invokeCommand(t, srv.readController.Init())
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{
+ AllResults: true,
+ }))
+
+ bts, err := os.ReadFile(tempFile)
+ assert.NoError(t, err)
+
+ assert.Equal(t, strings.Join(expected, ""), string(bts))
+ })
+
+ t.Run("with query", func(t *testing.T) {
+ srv := newService(t, serviceConfig{tableName: "count-to-30", defaultLimit: 5})
+
+ tempFile := tempFile(t)
+
+ expected := append([]string{
+ "pk,sk,num\n",
+ }, sliceutils.Generate(1, 15, func(i int) string {
+ return fmt.Sprintf("NUM,NUM#%02d,%d\n", i, i)
+ })...)
+
+ invokeCommand(t, srv.readController.Init())
+ invokeCommandWithPrompt(t, srv.readController.PromptForQuery(), "num<=15")
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{
+ AllResults: true,
+ }))
+
+ bts, err := os.ReadFile(tempFile)
+ assert.NoError(t, err)
+
+ assert.Equal(t, strings.Join(expected, ""), string(bts))
+ })
+ })
+ }
+ })
+
t.Run("should return error if result set is not set", func(t *testing.T) {
srv := newService(t, serviceConfig{tableName: "non-existant-table"})
tempFile := tempFile(t)
invokeCommandExpectingError(t, srv.readController.Init())
- invokeCommandExpectingError(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommandExpectingError(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
})
t.Run("should honour new columns in CSV file", func(t *testing.T) {
@@ -48,7 +104,7 @@ func TestExportController_ExportCSV(t *testing.T) {
invokeCommandWithPrompt(t, srv.columnsController.AddColumn(1), "address.street")
invokeCommand(t, srv.columnsController.ShiftColumnLeft(1))
- invokeCommand(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
bts, err := os.ReadFile(tempFile)
assert.NoError(t, err)
@@ -71,7 +127,7 @@ func TestExportController_ExportCSV(t *testing.T) {
invokeCommand(t, srv.columnsController.ToggleVisible(1))
invokeCommand(t, srv.columnsController.ToggleVisible(2))
- invokeCommand(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
bts, err := os.ReadFile(tempFile)
assert.NoError(t, err)
diff --git a/internal/dynamo-browse/controllers/jobbuilder.go b/internal/dynamo-browse/controllers/jobbuilder.go
index cbe7004..2128c44 100644
--- a/internal/dynamo-browse/controllers/jobbuilder.go
+++ b/internal/dynamo-browse/controllers/jobbuilder.go
@@ -51,6 +51,9 @@ func (jb JobBuilder[T]) executeJob(ctx context.Context) tea.Msg {
if jb.onEither != nil {
return jb.onEither(res, err)
} else if err == nil {
+ if jb.onDone == nil {
+ return nil
+ }
return jb.onDone(res)
} else {
if jb.onErr != nil {
diff --git a/internal/dynamo-browse/controllers/tableread.go b/internal/dynamo-browse/controllers/tableread.go
index a39bc1c..5297d1a 100644
--- a/internal/dynamo-browse/controllers/tableread.go
+++ b/internal/dynamo-browse/controllers/tableread.go
@@ -377,7 +377,7 @@ func (c *TableReadController) NextPage() tea.Msg {
resultSet := c.state.ResultSet()
if resultSet == nil {
return events.StatusMsg("Result-set is nil")
- } else if resultSet.LastEvaluatedKey == nil {
+ } else if !resultSet.HasNextPage() {
return events.StatusMsg("No more results")
}
currentFilter := c.state.filter
diff --git a/internal/dynamo-browse/controllers/tableread_test.go b/internal/dynamo-browse/controllers/tableread_test.go
index c042983..18bf075 100644
--- a/internal/dynamo-browse/controllers/tableread_test.go
+++ b/internal/dynamo-browse/controllers/tableread_test.go
@@ -89,7 +89,7 @@ func TestTableReadController_Query(t *testing.T) {
invokeCommand(t, srv.readController.Init())
invokeCommandWithPrompts(t, srv.readController.PromptForQuery(), `pk ^= "abc"`)
- invokeCommand(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
bts, err := os.ReadFile(tempFile)
assert.NoError(t, err)
@@ -107,7 +107,7 @@ func TestTableReadController_Query(t *testing.T) {
invokeCommand(t, srv.readController.Init())
invokeCommandWithPrompts(t, srv.readController.PromptForQuery(), `alpha = "This is some value"`)
- invokeCommand(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommand(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
bts, err := os.ReadFile(tempFile)
assert.NoError(t, err)
@@ -124,7 +124,7 @@ func TestTableReadController_Query(t *testing.T) {
tempFile := tempFile(t)
invokeCommandExpectingError(t, srv.readController.Init())
- invokeCommandExpectingError(t, srv.exportController.ExportCSV(tempFile))
+ invokeCommandExpectingError(t, srv.exportController.ExportCSV(tempFile, controllers.ExportOptions{}))
})
}
diff --git a/internal/dynamo-browse/controllers/tablewrite_test.go b/internal/dynamo-browse/controllers/tablewrite_test.go
index a28f182..a4f25ec 100644
--- a/internal/dynamo-browse/controllers/tablewrite_test.go
+++ b/internal/dynamo-browse/controllers/tablewrite_test.go
@@ -621,7 +621,7 @@ func newService(t *testing.T, cfg serviceConfig) *services {
writeController := controllers.NewTableWriteController(state, service, jobsController, readController, settingStore)
settingsController := controllers.NewSettingsController(settingStore, eventBus)
columnsController := controllers.NewColumnsController(eventBus)
- exportController := controllers.NewExportController(state, columnsController)
+ exportController := controllers.NewExportController(state, service, jobsController, columnsController)
scriptController := controllers.NewScriptController(scriptService, readController, settingsController, eventBus)
commandController := commandctrl.NewCommandController(inputHistoryService)
diff --git a/internal/dynamo-browse/models/models.go b/internal/dynamo-browse/models/models.go
index 1e9bf16..f94aa72 100644
--- a/internal/dynamo-browse/models/models.go
+++ b/internal/dynamo-browse/models/models.go
@@ -135,3 +135,7 @@ func (rs *ResultSet) RefreshColumns() {
rs.columns = columns
}
+
+func (rs *ResultSet) HasNextPage() bool {
+ return rs.LastEvaluatedKey != nil
+}
diff --git a/internal/dynamo-browse/ui/model.go b/internal/dynamo-browse/ui/model.go
index b3aa53a..35aa9cd 100644
--- a/internal/dynamo-browse/ui/model.go
+++ b/internal/dynamo-browse/ui/model.go
@@ -102,7 +102,14 @@ func NewModel(
if len(args) == 0 {
return events.Error(errors.New("expected filename"))
}
- return exportController.ExportCSV(args[0])
+
+ opts := controllers.ExportOptions{}
+ if len(args) == 2 && args[0] == "-all" {
+ opts.AllResults = true
+ args = args[1:]
+ }
+
+ return exportController.ExportCSV(args[0], opts)
},
"mark": func(ctx commandctrl.ExecContext, args []string) tea.Msg {
var markOp = controllers.MarkOpMark