From 2fe05fd98d5bea033c47feae965369a7573870a4 Mon Sep 17 00:00:00 2001
From: ma91n <mano.junki@gmail.com>
Date: Thu, 29 Feb 2024 16:55:28 +0900
Subject: [PATCH] Add --limit option in dump command

---
 .gitignore       |  2 ++
 cli/dump.go      | 12 +++++-------
 cli/dump_test.go |  2 +-
 cli/main.go      | 11 ++++++-----
 4 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/.gitignore b/.gitignore
index 931ecec..c223828 100644
--- a/.gitignore
+++ b/.gitignore
@@ -30,3 +30,5 @@
 testdata/csv/*.csv
 
 dist/
+
+*.xlsx
diff --git a/cli/dump.go b/cli/dump.go
index a902d2b..58d429b 100644
--- a/cli/dump.go
+++ b/cli/dump.go
@@ -18,8 +18,6 @@ import (
 
 const DefaultColumnCnt = 32
 
-var DumpWithDataRowLimit = 10
-
 var query = `SELECT tab.relname        AS table_name
 				 , tabdesc.description AS table_description
 				 , col.column_name
@@ -90,7 +88,7 @@ type ColumnDef struct {
 	DefaultValue string
 }
 
-func Dump(dbSource, targetFile string, tableNameArg, systemColumnArg string) error {
+func Dump(dbSource, targetFile, tableNameArg, systemColumnArg string, maxDumpSize int) error {
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
 	defer stop()
 
@@ -222,7 +220,7 @@ func Dump(dbSource, targetFile string, tableNameArg, systemColumnArg string) err
 			_ = f.SetCellStyle(sheetName, axisComment, axisName, style)
 		}
 
-		records, err := selectExistsRecords(ctx, conn, tableDef)
+		records, err := selectExistsRecords(ctx, conn, tableDef, maxDumpSize)
 		if err != nil {
 			return fmt.Errorf("select exists records: %w", err)
 		}
@@ -326,14 +324,14 @@ func Str(a any) string {
 	return a.(string)
 }
 
-func selectExistsRecords(ctx context.Context, conn *pgxpool.Pool, tableDef TableDef) ([][]any, error) {
-	rows, err := conn.Query(ctx, fmt.Sprintf(`select * from %s limit %d`, tableDef.Name, DumpWithDataRowLimit))
+func selectExistsRecords(ctx context.Context, conn *pgxpool.Pool, tableDef TableDef, maxDumpSize int) ([][]any, error) {
+	rows, err := conn.Query(ctx, fmt.Sprintf(`select * from %s limit %d`, tableDef.Name, maxDumpSize))
 	if err != nil {
 		return nil, fmt.Errorf("db access: %w", err)
 	}
 	defer rows.Close()
 
-	dataRecords := make([][]any, 0, DumpWithDataRowLimit)
+	dataRecords := make([][]any, 0, maxDumpSize)
 
 	for rows.Next() {
 		g := make([]any, len(tableDef.Columns))
diff --git a/cli/dump_test.go b/cli/dump_test.go
index 83610fc..c81ea4c 100644
--- a/cli/dump_test.go
+++ b/cli/dump_test.go
@@ -42,7 +42,7 @@ func Test_dump(t *testing.T) {
 	for _, tt := range tests {
 
 		t.Run(tt.name, func(t *testing.T) {
-			if err := Dump(tt.args.dbSource, tt.args.targetFile, tt.args.tableNameArg, tt.args.systemColumnArg); (err != nil) != tt.wantErr {
+			if err := Dump(tt.args.dbSource, tt.args.targetFile, tt.args.tableNameArg, tt.args.systemColumnArg, 10); (err != nil) != tt.wantErr {
 				t.Errorf("dump() error = %v, wantErr %v", err, tt.wantErr)
 			}
 		})
diff --git a/cli/main.go b/cli/main.go
index 6042248..b1d6b8e 100644
--- a/cli/main.go
+++ b/cli/main.go
@@ -17,10 +17,11 @@ var (
 	app    = kingpin.New("exceltesting", "Excel file driven testing helper tool").Version(Version)
 	source = app.Flag("source", `Database source (e.g. postgres://user:pass@host/dbname?sslmode=disable). EXCELTESTING_CONNECTION envvar is acceptable.`).Short('c').Envar("EXCELTESTING_CONNECTION").String()
 
-	dumpCommand = app.Command("dump", "Generate excel template file from database")
-	dumpFile    = dumpCommand.Arg("file", "Target excel file path (e.g. dump.xlsx)").Required().NoEnvar().String()
-	table       = dumpCommand.Flag("table", "Dump target table names (e.g. table1,table2,table3)").NoEnvar().String()
-	systemcolum = dumpCommand.Flag("systemcolum", "Specific system columns for cell style (e.g. created_at,updated_at,revision)").NoEnvar().String()
+	dumpCommand        = app.Command("dump", "Generate excel template file from database")
+	dumpFile           = dumpCommand.Arg("file", "Target excel file path (e.g. dump.xlsx)").Required().NoEnvar().String()
+	table              = dumpCommand.Flag("table", "Dump target table names (e.g. table1,table2,table3)").NoEnvar().String()
+	systemcolum        = dumpCommand.Flag("systemcolum", "Specific system columns for cell style (e.g. created_at,updated_at,revision)").NoEnvar().String()
+	maxDumpRecordLimit = dumpCommand.Flag("limit", "Max dump record limit size (e.g. created_at,updated_at,revision)").NoEnvar().Default("500").Int()
 
 	loadCommand                     = app.Command("load", "Load from excel file to database")
 	loadFile                        = loadCommand.Arg("file", "Target excel file path (e.g. input.xlsx)").Required().NoEnvar().ExistingFile()
@@ -40,7 +41,7 @@ func Main() {
 	var err error
 	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
 	case dumpCommand.FullCommand():
-		err = Dump(*source, *dumpFile, *table, *systemcolum)
+		err = Dump(*source, *dumpFile, *table, *systemcolum, *maxDumpRecordLimit)
 	case loadCommand.FullCommand():
 		req := exceltesting.LoadRequest{
 			TargetBookPath:                  *loadFile,