From ed53173a1d3e587383bb9055c2144042c24cddac Mon Sep 17 00:00:00 2001 From: Leon Mika 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