Passed
Pull Request — main (#3)
by Adriano
02:14
created

csvql.*csvql.executeQueryAndExport   A

Complexity

Conditions 2

Size

Total Lines 12
Code Lines 9

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 2
eloc 9
dl 0
loc 12
c 0
b 0
f 0
rs 9.95
nop 1
1
package csvql
2
3
import (
4
	"adrianolaselva.github.io/csvql/pkg/exportdata/jsonl"
5
	"adrianolaselva.github.io/csvql/pkg/filehandler"
6
	csvHandler "adrianolaselva.github.io/csvql/pkg/filehandler/csv"
7
	"adrianolaselva.github.io/csvql/pkg/storage"
8
	"adrianolaselva.github.io/csvql/pkg/storage/sqlite"
9
	"database/sql"
10
	"fmt"
11
	"github.com/chzyer/readline"
12
	"github.com/fatih/color"
13
	_ "github.com/mattn/go-sqlite3"
14
	"github.com/rodaine/table"
15
	"github.com/schollz/progressbar/v3"
16
	"io"
17
	"os"
18
	"strings"
19
)
20
21
const (
22
	cliPrompt          = "csvql> "
23
	cliInterruptPrompt = "^C"
24
	cliEOFPrompt       = "exit"
25
)
26
27
type Csvql interface {
28
	Run() error
29
	Close() error
30
}
31
32
type csvql struct {
33
	storage     storage.Storage
34
	bar         *progressbar.ProgressBar
35
	params      CsvqlParams
36
	fileHandler filehandler.FileHandler
37
}
38
39
func New(params CsvqlParams) (Csvql, error) {
40
	sqLiteStorage, err := sqlite.NewSqLiteStorage(params.DataSourceName)
41
	if err != nil {
42
		return nil, err
43
	}
44
45
	bar := progressbar.NewOptions(0,
46
		progressbar.OptionSetWriter(os.Stdout),
47
		progressbar.OptionEnableColorCodes(true),
48
		progressbar.OptionShowBytes(true),
49
		progressbar.OptionFullWidth(),
50
		progressbar.OptionSetDescription("[cyan][1/1][reset] loading data..."),
51
		progressbar.OptionSetTheme(progressbar.Theme{
52
			Saucer:        "[green]=[reset]",
53
			SaucerHead:    "[green]>[reset]",
54
			SaucerPadding: " ",
55
			BarStart:      "[",
56
			BarEnd:        "]",
57
		}))
58
59
	impData := csvHandler.NewCsvHandler(params.FileInput, rune(params.Delimiter[0]), bar, sqLiteStorage)
60
61
	return &csvql{params: params, bar: bar, fileHandler: impData, storage: sqLiteStorage}, nil
62
}
63
64
func (c *csvql) Run() error {
65
	defer c.bar.Clear()
66
67
	if err := c.fileHandler.Import(); err != nil {
68
		return fmt.Errorf("failed to import data %s", err)
69
	}
70
	defer c.fileHandler.Close()
71
72
	return c.execute()
73
}
74
75
func (c *csvql) execute() error {
76
	if c.params.Query != "" && c.params.Export == "" {
77
		return c.executeQuery(c.params.Query)
78
	}
79
80
	if c.params.Query != "" && c.params.Export != "" {
81
		return c.executeQueryAndExport(c.params.Query)
82
	}
83
84
	if err := c.initializePrompt(); err != nil {
85
		return err
86
	}
87
88
	return nil
89
}
90
91
func (c *csvql) Close() error {
92
	defer func(fileHandler filehandler.FileHandler) {
93
		_ = fileHandler.Close()
94
	}(c.fileHandler)
95
96
	return nil
97
}
98
99
func (c *csvql) initializePrompt() error {
100
	l, err := readline.NewEx(&readline.Config{
101
		Prompt:          cliPrompt,
102
		InterruptPrompt: cliInterruptPrompt,
103
		EOFPrompt:       cliEOFPrompt,
104
		AutoComplete: readline.SegmentFunc(func(i [][]rune, i2 int) [][]rune {
105
			return nil
106
		}),
107
	})
108
	if err != nil {
109
		return err
110
	}
111
	defer l.Close()
112
113
	l.CaptureExitSignal()
114
115
	for {
116
		line, err := l.Readline()
117
		if err == readline.ErrInterrupt {
118
			if len(line) == 0 {
119
				break
120
			}
121
122
			continue
123
		}
124
125
		if err == io.EOF {
126
			break
127
		}
128
129
		line = strings.TrimSpace(line)
130
		if err := c.executeQuery(line); err != nil {
131
			fmt.Fprintf(os.Stderr, "%s\n", err.Error())
132
		}
133
	}
134
135
	return nil
136
}
137
138
func (c *csvql) executeQueryAndExport(line string) error {
139
	c.bar.Reset()
140
	c.bar.ChangeMax(c.fileHandler.Lines())
141
	defer c.bar.Finish()
142
143
	rows, err := c.storage.Query(line)
144
	if err != nil {
145
		return err
146
	}
147
	defer rows.Close()
148
149
	return jsonl.NewJsonlExport(rows, c.params.Export, c.bar).Export()
150
}
151
152
func (c *csvql) executeQuery(line string) error {
153
	rows, err := c.storage.Query(line)
154
	if err != nil {
155
		return err
156
	}
157
	defer rows.Close()
158
159
	return c.printResult(rows)
160
}
161
162
func (c *csvql) printResult(rows *sql.Rows) error {
163
	columns, err := rows.Columns()
164
	if err != nil {
165
		return err
166
	}
167
168
	cols := make([]interface{}, 0)
169
	for _, c := range columns {
170
		cols = append(cols, c)
171
	}
172
173
	tbl := table.New(cols...).
174
		WithHeaderFormatter(color.New(color.FgGreen, color.Underline).SprintfFunc()).
175
		WithFirstColumnFormatter(color.New(color.FgYellow).SprintfFunc()).
176
		WithWriter(os.Stdout)
177
178
	for rows.Next() {
179
		values := make([]interface{}, len(columns))
180
		pointers := make([]interface{}, len(columns))
181
		for i := range values {
182
			pointers[i] = &values[i]
183
		}
184
185
		if err := rows.Scan(pointers...); err != nil {
186
			return err
187
		}
188
189
		tbl.AddRow(values...)
190
	}
191
192
	_ = c.bar.Clear()
193
	tbl.Print()
194
195
	return nil
196
}
197