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

csv.*csvHandler.loadTotalRows   A

Complexity

Conditions 5

Size

Total Lines 21
Code Lines 16

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 5
eloc 16
dl 0
loc 21
c 0
b 0
f 0
rs 9.1333
nop 0
1
package csv
2
3
import (
4
	"adrianolaselva.github.io/csvql/pkg/filehandler"
5
	"adrianolaselva.github.io/csvql/pkg/storage"
6
	"bytes"
7
	"database/sql"
8
	"encoding/csv"
9
	"fmt"
10
	"github.com/schollz/progressbar/v3"
11
	"io"
12
	"os"
13
)
14
15
type csvHandler struct {
16
	bar       *progressbar.ProgressBar
17
	storage   storage.Storage
18
	file      *os.File
19
	fileInput string
20
	lines     int
21
	delimiter rune
22
}
23
24
func NewCsvHandler(fileInput string, delimiter rune, bar *progressbar.ProgressBar, storage storage.Storage) filehandler.FileHandler {
25
	return &csvHandler{fileInput: fileInput, delimiter: delimiter, storage: storage, bar: bar}
26
}
27
28
// Import import data
29
func (c *csvHandler) Import() error {
30
	if err := c.openFile(); err != nil {
31
		return err
32
	}
33
34
	if err := c.loadTotalRows(); err != nil {
35
		return err
36
	}
37
38
	if err := c.loadDataFromFile(); err != nil {
39
		return err
40
	}
41
42
	return nil
43
}
44
45
// Query execute statements
46
func (c *csvHandler) Query(cmd string) (*sql.Rows, error) {
47
	return c.storage.Query(cmd)
48
}
49
50
// Lines return total lines
51
func (c *csvHandler) Lines() int {
52
	return c.lines
53
}
54
55
// Close execute in defer
56
func (c *csvHandler) Close() error {
57
	defer func(storage storage.Storage) {
58
		_ = storage.Close()
59
	}(c.storage)
60
61
	defer func(file *os.File) {
62
		_ = file.Close()
63
	}(c.file)
64
65
	return nil
66
}
67
68
func (c *csvHandler) loadDataFromFile() error {
69
	c.bar.ChangeMax(c.lines)
70
71
	r := csv.NewReader(c.file)
72
	r.Comma = c.delimiter
73
74
	headers, err := r.Read()
75
	if err != nil {
76
		return fmt.Errorf("failed to load headers: %s", err)
77
	}
78
79
	if err := c.storage.SetColumns(headers).BuildStructure(); err != nil {
80
		return fmt.Errorf("failed to load headers and build structure: %s", err)
81
	}
82
83
	line := 1
84
	for {
85
		line++
86
		records, err := r.Read()
87
		if err == io.EOF {
88
			break
89
		}
90
91
		var values []any
92
		for _, r := range records {
93
			values = append(values, r)
94
		}
95
96
		_ = c.bar.Add(1)
97
		if err := c.storage.InsertRow(values); err != nil {
98
			return fmt.Errorf("failed to process row: %s", err)
99
		}
100
	}
101
102
	return nil
103
}
104
105
// openFile open file
106
func (c *csvHandler) openFile() error {
107
	f, err := os.Open(c.fileInput)
108
	if err != nil {
109
		return fmt.Errorf("failed to open file: %s", err)
110
	}
111
112
	c.file = f
113
114
	return nil
115
}
116
117
// loadTotalRows load total rows in file
118
func (c *csvHandler) loadTotalRows() error {
119
	r, err := os.Open(c.fileInput)
120
	if err != nil {
121
		return err
122
	}
123
	defer r.Close()
124
125
	buf := make([]byte, 32*1024)
126
	c.lines = 0
127
	lineSep := []byte{'\n'}
128
129
	for {
130
		r, err := r.Read(buf)
131
		c.lines += bytes.Count(buf[:r], lineSep)
132
133
		switch {
134
		case err == io.EOF:
135
			return nil
136
137
		case err != nil:
138
			return fmt.Errorf("failed to totalize rows: %s", err)
139
		}
140
	}
141
}
142