-
Notifications
You must be signed in to change notification settings - Fork 3
/
store.go
150 lines (122 loc) · 3.3 KB
/
store.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package csv
import (
"encoding/csv"
"fmt"
"os"
"path/filepath"
"reflect"
)
// Store allows to store data in csv format
type Store struct {
path string
f *os.File
reader *csv.Reader
writer *csv.Writer
}
// NewStore creates a Store which data is kept in file path
func NewStore(path string) *Store {
return &Store{
path: path,
}
}
// Open the store file
func (s *Store) Open(createIfNotExist bool) (err error) {
s.f, err = openFile(s.path, createIfNotExist)
if err != nil {
return
}
// Set csv reader and writer
s.reader = csv.NewReader(s.f)
s.writer = csv.NewWriter(s.f)
return nil
}
// Reader returns CSV reader
// You MUST call Open() before calling Reader
func (s *Store) Reader() *csv.Reader {
return s.reader
}
// ReadAll reads all CSV records from disk
func (s *Store) ReadAll() ([][]string, error) {
err := s.Open(false)
if err != nil {
return nil, err
}
defer s.Close()
return s.reader.ReadAll()
}
var unmarshallerType = reflect.TypeOf((*Unmarshaller)(nil)).Elem()
// ReadAllStructs reads all CSV records from disk and set it into v
// It assumes v to be a pointer to a slice of object implementing the Unmarshaller interface
func (s *Store) ReadAllStructs(v interface{}) error {
rv := reflect.ValueOf(v)
if rv.Kind() != reflect.Ptr || rv.IsNil() || rv.Elem().Kind() != reflect.Slice || rv.Elem().Type().Elem().Kind() != reflect.Ptr || !rv.Elem().Type().Elem().Implements(unmarshallerType) {
return fmt.Errorf("invalid type %T, expects a slice of %v", v, unmarshallerType)
}
records, err := s.ReadAll()
if err != nil {
return err
}
sv := reflect.MakeSlice(rv.Type().Elem(), 0, len(records))
for _, rec := range records {
recV := reflect.New(rv.Type().Elem().Elem().Elem())
outputs := recV.MethodByName("UnmarshalCSV").Call([]reflect.Value{reflect.ValueOf(rec)})
if !outputs[0].IsNil() {
return outputs[0].Interface().(error)
}
sv = reflect.Append(sv, recV)
}
rv.Elem().Set(sv)
return nil
}
// Writer returns CSV writer
// You MUST call Open(true) before calling Writer
func (s *Store) Writer() *csv.Writer {
return s.writer
}
// WriteAll writes all CSV records to disk
func (s *Store) WriteAll(records [][]string) error {
err := s.Open(true)
if err != nil {
return err
}
defer s.Close()
return s.writer.WriteAll(records)
}
var marshallerType = reflect.TypeOf((*Marshaller)(nil)).Elem()
// WriteAllStructs writes all values to disk
// It expects all values to implement Marshaller interface
func (s *Store) WriteAllStructs(values []interface{}) error {
var records [][]string
for _, v := range values {
marshaller, ok := v.(Marshaller)
if !ok {
return fmt.Errorf("invalid value type %T does not implement %T", v, marshallerType)
}
rec, err := marshaller.MarshalCSV()
if err != nil {
return err
}
records = append(records, rec)
}
return s.WriteAll(records)
}
func (s *Store) Close() error {
return s.f.Close()
}
func openFile(path string, createIfNotExist bool) (*os.File, error) {
// check if file exist
_, err := os.Stat(path)
if err != nil {
if !createIfNotExist {
return nil, fmt.Errorf("file %v does not exist", path)
}
// file does not exist so we create it
err = os.MkdirAll(filepath.Dir(path), 0o700)
if err == nil {
return os.Create(path)
}
} else {
return os.OpenFile(path, os.O_APPEND|os.O_RDWR, 0o644)
}
return nil, err
}