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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
|
package proto
import (
"strconv"
"strings"
"github.com/go-faster/errors"
)
var (
_ Column = (*ColEnum)(nil)
_ ColumnOf[string] = (*ColEnum)(nil)
_ Inferable = (*ColEnum)(nil)
_ Preparable = (*ColEnum)(nil)
)
// ColEnum is inference helper for enums.
//
// You can set Values and actual enum mapping will be inferred during query
// execution.
type ColEnum struct {
t ColumnType
base ColumnType
rawToStr map[int]string
strToRaw map[string]int
raw8 ColEnum8
raw16 ColEnum16
// Values of ColEnum.
Values []string
}
func (e *ColEnum) raw() Column {
if e.t.Base() == ColumnTypeEnum8 {
return &e.raw8
}
return &e.raw16
}
func (e ColEnum) Row(i int) string {
return e.Values[i]
}
// Append value to Enum8 column.
func (e *ColEnum) Append(v string) {
e.Values = append(e.Values, v)
}
func (e *ColEnum) AppendArr(vs []string) {
e.Values = append(e.Values, vs...)
}
func (e *ColEnum) parse(t ColumnType) error {
if e.rawToStr == nil {
e.rawToStr = map[int]string{}
}
if e.strToRaw == nil {
e.strToRaw = map[string]int{}
}
elements := t.Elem().String()
for _, elem := range strings.Split(elements, ",") {
def := strings.TrimSpace(elem)
// 'hello' = 1
parts := strings.SplitN(def, "=", 2)
if len(parts) != 2 {
return errors.Errorf("bad enum definition %q", def)
}
var (
left = strings.TrimSpace(parts[0]) // 'hello'
right = strings.TrimSpace(parts[1]) // 1
)
idx, err := strconv.Atoi(right)
if err != nil {
return errors.Errorf("bad right side of definition %q", right)
}
left = strings.TrimFunc(left, func(c rune) bool {
return c == '\''
})
e.strToRaw[left] = idx
e.rawToStr[idx] = left
}
return nil
}
func (e *ColEnum) Infer(t ColumnType) error {
if !strings.HasPrefix(t.Base().String(), "Enum") {
return errors.Errorf("invalid base %q to infer enum", t.Base())
}
if err := e.parse(t); err != nil {
return errors.Wrap(err, "parse type")
}
base := t.Base()
switch base {
case ColumnTypeEnum8, ColumnTypeEnum16:
e.base = base
default:
return errors.Errorf("invalid base %q", base)
}
e.t = t
return nil
}
func (e *ColEnum) Rows() int {
return len(e.Values)
}
func appendEnum[E Enum8 | Enum16](c []E, mapping map[int]string, values []string) ([]string, error) {
for _, v := range c {
s, ok := mapping[int(v)]
if !ok {
return nil, errors.Errorf("unknown enum value %d", v)
}
values = append(values, s)
}
return values, nil
}
func (e *ColEnum) DecodeColumn(r *Reader, rows int) error {
if err := e.raw().DecodeColumn(r, rows); err != nil {
return errors.Wrap(err, "raw")
}
var (
err error
v []string
)
switch e.base {
case ColumnTypeEnum8:
v, err = appendEnum[Enum8](e.raw8, e.rawToStr, e.Values[:0])
case ColumnTypeEnum16:
v, err = appendEnum[Enum16](e.raw16, e.rawToStr, e.Values[:0])
default:
return errors.Errorf("invalid enum base %q", e.base)
}
if err != nil {
return errors.Wrap(err, "map values")
}
e.Values = v
return nil
}
func (e *ColEnum) Reset() {
e.raw().Reset()
e.Values = e.Values[:0]
}
func (e *ColEnum) Prepare() error {
e.raw8 = e.raw8[:0]
e.raw16 = e.raw16[:0]
for _, v := range e.Values {
raw, ok := e.strToRaw[v]
if !ok {
return errors.Errorf("unknown enum value %q", v)
}
switch e.base {
case ColumnTypeEnum8:
e.raw8.Append(Enum8(raw))
case ColumnTypeEnum16:
e.raw16.Append(Enum16(raw))
default:
return errors.Errorf("invalid base %q", e.base)
}
}
return nil
}
func (e *ColEnum) EncodeColumn(b *Buffer) {
e.raw().EncodeColumn(b)
}
func (e *ColEnum) Type() ColumnType { return e.t }
|
