adrianolaselva /
csvql
| Total Lines | 126 |
| Duplicated Lines | 0 % |
| Changes | 0 | ||
| 1 | package csv |
||
| 2 | |||
| 3 | import ( |
||
| 4 | "adrianolaselva.github.io/csvql/pkg/exportdata" |
||
| 5 | "database/sql" |
||
| 6 | "encoding/csv" |
||
| 7 | "fmt" |
||
| 8 | "github.com/schollz/progressbar/v3" |
||
| 9 | "os" |
||
| 10 | "path/filepath" |
||
| 11 | ) |
||
| 12 | |||
| 13 | const ( |
||
| 14 | fileModeDefault os.FileMode = 0644 |
||
| 15 | ) |
||
| 16 | |||
| 17 | type csvExport struct { |
||
| 18 | rows *sql.Rows |
||
| 19 | bar *progressbar.ProgressBar |
||
| 20 | file *os.File |
||
| 21 | exportPath string |
||
| 22 | columns []string |
||
| 23 | } |
||
| 24 | |||
| 25 | func NewCsvExport(rows *sql.Rows, exportPath string, bar *progressbar.ProgressBar) exportdata.Export { |
||
| 26 | return &csvExport{rows: rows, exportPath: exportPath, bar: bar} |
||
| 27 | } |
||
| 28 | |||
| 29 | // Export rows in file |
||
| 30 | func (c *csvExport) Export() error { |
||
| 31 | if err := c.loadColumns(); err != nil { |
||
| 32 | return fmt.Errorf("failed to load columns: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 33 | } |
||
| 34 | |||
| 35 | if err := c.openFile(); err != nil { |
||
| 36 | return fmt.Errorf("failed to open file: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 37 | } |
||
| 38 | |||
| 39 | w := csv.NewWriter(c.file) |
||
| 40 | defer w.Flush() |
||
| 41 | |||
| 42 | if err := w.Write(c.columns); err != nil { |
||
| 43 | return fmt.Errorf("failed to write headers: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 44 | } |
||
| 45 | |||
| 46 | for c.rows.Next() { |
||
| 47 | _ = c.bar.Add(1) |
||
| 48 | if err := c.readAndAppendFile(w); err != nil { |
||
| 49 | return fmt.Errorf("failed to read and append line in file: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 50 | } |
||
| 51 | } |
||
| 52 | |||
| 53 | return nil |
||
| 54 | } |
||
| 55 | |||
| 56 | // readAndAppendFile read line and append in file |
||
| 57 | func (c *csvExport) readAndAppendFile(w *csv.Writer) error { |
||
| 58 | values := make([]interface{}, len(c.columns)) |
||
| 59 | pointers := make([]interface{}, len(c.columns)) |
||
| 60 | for i := range values { |
||
| 61 | pointers[i] = &values[i] |
||
| 62 | } |
||
| 63 | |||
| 64 | if err := c.rows.Scan(pointers...); err != nil { |
||
| 65 | return fmt.Errorf("failed to load row: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 66 | } |
||
| 67 | |||
| 68 | if err := w.Write(c.convertToStringArray(values)); err != nil { |
||
| 69 | return fmt.Errorf("failed to write row: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 70 | } |
||
| 71 | |||
| 72 | return nil |
||
| 73 | } |
||
| 74 | |||
| 75 | // convertToStringArray convert string array to string array |
||
| 76 | func (c *csvExport) convertToStringArray(records []interface{}) []string { |
||
| 77 | values := make([]string, 0, len(records)) |
||
| 78 | for _, r := range records { |
||
| 79 | values = append(values, r.(string)) |
||
| 80 | } |
||
| 81 | |||
| 82 | return values |
||
| 83 | } |
||
| 84 | |||
| 85 | // Close execute in defer |
||
| 86 | func (c *csvExport) Close() error { |
||
| 87 | defer func(file *os.File) { |
||
| 88 | _ = file.Close() |
||
| 89 | }(c.file) |
||
| 90 | |||
| 91 | return nil |
||
| 92 | } |
||
| 93 | |||
| 94 | // openFile open file |
||
| 95 | func (c *csvExport) openFile() error { |
||
| 96 | if _, err := os.Stat(c.exportPath); !os.IsNotExist(err) { |
||
| 97 | err := os.Remove(c.exportPath) |
||
| 98 | if err != nil { |
||
| 99 | return fmt.Errorf("failed to remove file: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 100 | } |
||
| 101 | } |
||
| 102 | |||
| 103 | if err := os.MkdirAll(filepath.Dir(c.exportPath), os.ModePerm); err != nil { |
||
| 104 | return fmt.Errorf("failed to create path: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 105 | } |
||
| 106 | |||
| 107 | file, err := os.OpenFile(c.exportPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, fileModeDefault) |
||
| 108 | if err != nil { |
||
| 109 | return fmt.Errorf("failed to open file %s: %w", c.exportPath, err) |
||
|
1 ignored issue
–
show
|
|||
| 110 | } |
||
| 111 | |||
| 112 | c.file = file |
||
| 113 | |||
| 114 | return nil |
||
| 115 | } |
||
| 116 | |||
| 117 | // loadColumns load columns |
||
| 118 | func (c *csvExport) loadColumns() error { |
||
| 119 | columns, err := c.rows.Columns() |
||
| 120 | if err != nil { |
||
| 121 | return fmt.Errorf("failed to load columns: %w", err) |
||
|
1 ignored issue
–
show
|
|||
| 122 | } |
||
| 123 | |||
| 124 | c.columns = columns |
||
| 125 | |||
| 126 | return nil |
||
| 127 | } |
||
| 128 |
Loading history...