Skip to content

Commit f8507cb

Browse files
author
vl4deee11
committed
added nullable enums
1 parent c9d6b03 commit f8507cb

File tree

2 files changed

+140
-71
lines changed

2 files changed

+140
-71
lines changed

‎clickhouse_nullable_array_test.go‎

Lines changed: 95 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -31,51 +31,55 @@ func Test_NullableArray(t *testing.T) {
3131
arr_string Array(Nullable(String)),
3232
arr_arr_string Array(Array(Nullable(String))),
3333
arr_date Array(Nullable(Date)),
34-
arr_datetime Array(Nullable(DateTime))
34+
arr_datetime Array(Nullable(DateTime)),
35+
arr_enum8_str Array(Nullable(Enum8('a8' = 1, 'b8' = 2))),
36+
arr_enum8_int Array(Nullable(Enum8('a8' = 1, 'b8' = 2))),
37+
arr_enum16_str Array(Nullable(Enum16('a16' = 1, 'b16' = 2))),
38+
arr_enum16_int Array(Nullable(Enum16('a16' = 1, 'b16' = 2)))
3539
) Engine = Memory;
3640
`
3741
dml = `
38-
INSERT INTO clickhouse_test_nullable_array (
39-
arr_decimal,
40-
arr_int8,
41-
arr_int16,
42-
arr_int32,
43-
arr_int64,
44-
45-
arr_uint8,
46-
arr_uint16,
47-
arr_uint32,
48-
arr_uint64,
49-
50-
arr_float32,
51-
arr_float64,
52-
53-
arr_ipv6,
54-
arr_ipv4,
55-
56-
arr_string,
57-
arr_arr_string,
58-
59-
arr_date,
60-
arr_datetime
61-
) VALUES (
62-
?,
63-
?,
64-
?,
65-
?,
66-
?,
67-
?,
68-
?,
69-
?,
70-
?,
71-
?,
72-
?,
73-
?,
74-
?,
75-
?,
76-
?,
77-
?
78-
)
42+
INSERT INTO clickhouse_test_nullable_array (arr_decimal,
43+
arr_int8,
44+
arr_int16,
45+
arr_int32,
46+
arr_int64,
47+
arr_uint8,
48+
arr_uint16,
49+
arr_uint32,
50+
arr_uint64,
51+
arr_float32,
52+
arr_float64,
53+
arr_ipv6,
54+
arr_ipv4,
55+
arr_string,
56+
arr_arr_string,
57+
arr_date,
58+
arr_datetime,
59+
arr_enum8_str,
60+
arr_enum8_int,
61+
arr_enum16_str,
62+
arr_enum16_int)
63+
VALUES (?,
64+
?,
65+
?,
66+
?,
67+
?,
68+
?,
69+
?,
70+
?,
71+
?,
72+
?,
73+
?,
74+
?,
75+
?,
76+
?,
77+
?,
78+
?,
79+
?,
80+
?,
81+
?,
82+
?);
7983
`
8084
query = `
8185
SELECT
@@ -107,16 +111,28 @@ func Test_NullableArray(t *testing.T) {
107111
timeV, _ := time.Parse("2006-01-02 15:04:05", "2021-07-11 00:00:00")
108112
dateV, _ := time.Parse("2006-01-02", "2021-07-11")
109113

114+
enum8VA := "a8"
115+
enum8VB := "b8"
116+
117+
enum8V1 := int8(1)
118+
enum8V2 := int8(2)
119+
120+
enum16VA := "a16"
121+
enum16VB := "b16"
122+
123+
enum16V1 := int16(1)
124+
enum16V2 := int16(2)
125+
110126
var timeNil *time.Time
111127

112-
if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?debug=true"); assert.NoError(t, err) {
128+
if connect, err := sql.Open("clickhouse", "tcp://0.0.0.0:9000?debug=true"); assert.NoError(t, err) {
113129
if tx, err := connect.Begin(); assert.NoError(t, err) {
114130
if _, err := connect.Exec("DROP TABLE IF EXISTS clickhouse_test_nullable_array"); assert.NoError(t, err) {
115131
if _, err := tx.Exec(ddl); assert.NoError(t, err) {
116132
if tx, err := connect.Begin(); assert.NoError(t, err) {
117133
stmt, err := tx.Prepare(dml)
118134
if assert.NoError(t, err) {
119-
for i := 0; i < 100; i++ {
135+
for i := 0; i < 10; i++ {
120136
if _, err := stmt.Exec(
121137
[]*float64{&decV, nil, &decV},
122138
[]*int8{&int8V, nil, &int8V},
@@ -140,6 +156,11 @@ func Test_NullableArray(t *testing.T) {
140156

141157
[]*time.Time{&dateV, nil, &dateV},
142158
[]*time.Time{&timeV, nil, &timeV},
159+
160+
[]*string{&enum8VA, nil, &enum8VB},
161+
[]*int8{&enum8V1, nil, &enum8V2},
162+
[]*string{&enum16VA, nil, &enum16VB},
163+
[]*int16{&enum16V1, nil, &enum16V2},
143164
); !assert.NoError(t, err) {
144165
t.Fatal(err)
145166
}
@@ -152,23 +173,27 @@ func Test_NullableArray(t *testing.T) {
152173
if rows, err := connect.Query(query); assert.NoError(t, err) {
153174
for rows.Next() {
154175
var (
155-
ArrDecimal = make([]*int64, 0)
156-
ArrInt8 = make([]*int8, 0)
157-
ArrInt16 = make([]*int16, 0)
158-
ArrInt32 = make([]*int32, 0)
159-
ArrInt64 = make([]*int64, 0)
160-
ArrUInt8 = make([]*uint8, 0)
161-
ArrUInt16 = make([]*uint16, 0)
162-
ArrUInt32 = make([]*uint32, 0)
163-
ArrUInt64 = make([]*uint64, 0)
164-
ArrFloat32 = make([]*float32, 0)
165-
ArrFloat64 = make([]*float64, 0)
166-
ArrIpv6 = make([]*net.IP, 0)
167-
ArrIpv4 = make([]*net.IP, 0)
168-
ArrString = make([]*string, 0)
169-
ArrArrString = make([][]*string, 0)
170-
ArrDate = make([]*time.Time, 0)
171-
ArrDateTime = make([]*time.Time, 0)
176+
ArrDecimal = make([]*int64, 0)
177+
ArrInt8 = make([]*int8, 0)
178+
ArrInt16 = make([]*int16, 0)
179+
ArrInt32 = make([]*int32, 0)
180+
ArrInt64 = make([]*int64, 0)
181+
ArrUInt8 = make([]*uint8, 0)
182+
ArrUInt16 = make([]*uint16, 0)
183+
ArrUInt32 = make([]*uint32, 0)
184+
ArrUInt64 = make([]*uint64, 0)
185+
ArrFloat32 = make([]*float32, 0)
186+
ArrFloat64 = make([]*float64, 0)
187+
ArrIpv6 = make([]*net.IP, 0)
188+
ArrIpv4 = make([]*net.IP, 0)
189+
ArrString = make([]*string, 0)
190+
ArrArrString = make([][]*string, 0)
191+
ArrDate = make([]*time.Time, 0)
192+
ArrDateTime = make([]*time.Time, 0)
193+
ArrEnum8Str = make([]*string, 0)
194+
ArrEnum8Int8 = make([]*string, 0)
195+
ArrEnum16Str = make([]*string, 0)
196+
ArrEnum16Int16 = make([]*string, 0)
172197
)
173198
if err := rows.Scan(
174199
&ArrDecimal,
@@ -188,6 +213,10 @@ func Test_NullableArray(t *testing.T) {
188213
&ArrArrString,
189214
&ArrDate,
190215
&ArrDateTime,
216+
&ArrEnum8Str,
217+
&ArrEnum8Int8,
218+
&ArrEnum16Str,
219+
&ArrEnum16Int16,
191220
); assert.NoError(t, err) {
192221
assert.Equal(t, ArrDecimal, []*int64{&int64Dec, nil, &int64Dec})
193222
assert.Equal(t, ArrInt8, []*int8{&int8V, nil, &int8V})
@@ -213,6 +242,11 @@ func Test_NullableArray(t *testing.T) {
213242
assert.True(t, len(ArrDateTime) == 3)
214243
assert.Equal(t, ArrDate[1], timeNil)
215244
assert.Equal(t, ArrDateTime[1], timeNil)
245+
246+
assert.Equal(t, ArrEnum8Int8, []*string{&enum8VA, nil, &enum8VB})
247+
assert.Equal(t, ArrEnum16Int16, []*string{&enum16VA, nil, &enum16VB})
248+
assert.Equal(t, ArrEnum8Str, []*string{&enum8VA, nil, &enum8VB})
249+
assert.Equal(t, ArrEnum16Str, []*string{&enum16VA, nil, &enum16VB})
216250
}
217251
}
218252
}

‎lib/column/enum.go‎

Lines changed: 45 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,16 +39,7 @@ func (enum *Enum) Read(decoder *binary.Decoder, isNull bool) (interface{}, error
3939
func (enum *Enum) Write(encoder *binary.Encoder, v interface{}) error {
4040
switch v := v.(type) {
4141
case string:
42-
ident, found := enum.iv[v]
43-
if !found {
44-
return fmt.Errorf("invalid Enum ident: %s", v)
45-
}
46-
switch ident := ident.(type) {
47-
case int8:
48-
return encoder.Int8(ident)
49-
case int16:
50-
return encoder.Int16(ident)
51-
}
42+
return enum.encodeFromString(v, encoder)
5243
case uint8:
5344
if _, ok := enum.baseType.(int8); ok {
5445
return encoder.Int8(int8(v))
@@ -72,13 +63,57 @@ func (enum *Enum) Write(encoder *binary.Encoder, v interface{}) error {
7263
case int16:
7364
return encoder.Int16(int16(v))
7465
}
66+
// nullable enums
67+
case *string:
68+
return enum.encodeFromString(*v, encoder)
69+
case *uint8:
70+
if _, ok := enum.baseType.(int8); ok {
71+
return encoder.Int8(int8(*v))
72+
}
73+
case *int8:
74+
if _, ok := enum.baseType.(int8); ok {
75+
return encoder.Int8(*v)
76+
}
77+
case *uint16:
78+
if _, ok := enum.baseType.(int16); ok {
79+
return encoder.Int16(int16(*v))
80+
}
81+
case *int16:
82+
if _, ok := enum.baseType.(int16); ok {
83+
return encoder.Int16(*v)
84+
}
85+
case *int64:
86+
switch enum.baseType.(type) {
87+
case int8:
88+
return encoder.Int8(int8(*v))
89+
case int16:
90+
return encoder.Int16(int16(*v))
91+
}
7592
}
7693
return &ErrUnexpectedType{
7794
T: v,
7895
Column: enum,
7996
}
8097
}
8198

99+
func (enum *Enum) encodeFromString(v string, encoder *binary.Encoder) error {
100+
ident, found := enum.iv[v]
101+
if !found {
102+
return fmt.Errorf("invalid Enum ident: %s", v)
103+
}
104+
switch ident := ident.(type) {
105+
case int8:
106+
return encoder.Int8(ident)
107+
case int16:
108+
return encoder.Int16(ident)
109+
default:
110+
return &ErrUnexpectedType{
111+
T: ident,
112+
Column: enum,
113+
}
114+
}
115+
}
116+
82117
func (enum *Enum) defaultValue() interface{} {
83118
return enum.baseType
84119
}

0 commit comments

Comments
 (0)