parse enum type

Author: kshvakov

README.md

@@ -24,11 +24,9 @@ Golang SQL database driver for [Yandex ClickHouse](https://clickhouse.yandex/)
 * FixedString(N)
 * Date 
 * DateTime
+* Enum
 * [Array(T) (one-dimensional)](https://clickhouse.yandex/reference_en.html#Array(T)) [godoc](https://godoc.org/github.com/kshvakov/clickhouse#Array)
 
-## TODO 
-
-* Compression
 
 example:
 ```

clickhouse_negative_test.go

@@ -0,0 +1,39 @@
+package clickhouse_test
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/kshvakov/clickhouse"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_Negative_OpenConnectAndPing(t *testing.T) {
+	if connect, err := sql.Open("clickhouse", ""); assert.NoError(t, err) {
+		assert.Error(t, connect.Ping())
+	}
+	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:10000"); assert.NoError(t, err) {
+		assert.Error(t, connect.Ping())
+	}
+	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?username=invalid"); assert.NoError(t, err) {
+		if err := connect.Ping(); assert.Error(t, err) {
+			if exception, ok := err.(*clickhouse.Exception); assert.True(t, ok) {
+				assert.Equal(t, int32(192), exception.Code)
+			}
+		}
+	}
+	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?password=invalid"); assert.NoError(t, err) {
+		if err := connect.Ping(); assert.Error(t, err) {
+			if exception, ok := err.(*clickhouse.Exception); assert.True(t, ok) {
+				assert.Equal(t, int32(193), exception.Code)
+			}
+		}
+	}
+	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?database=invalid"); assert.NoError(t, err) {
+		if err := connect.Ping(); assert.Error(t, err) {
+			if exception, ok := err.(*clickhouse.Exception); assert.True(t, ok) {
+				assert.Equal(t, int32(81), exception.Code)
+			}
+		}
+	}
+}

clickhouse_test.go

@@ -15,33 +15,6 @@ func Test_OpenConnectAndPing(t *testing.T) {
 	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?debug=true"); assert.NoError(t, err) {
 		assert.NoError(t, connect.Ping())
 	}
-	if connect, err := sql.Open("clickhouse", ""); assert.NoError(t, err) {
-		assert.Error(t, connect.Ping())
-	}
-	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:10000"); assert.NoError(t, err) {
-		assert.Error(t, connect.Ping())
-	}
-	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?username=invalid"); assert.NoError(t, err) {
-		if err := connect.Ping(); assert.Error(t, err) {
-			if exception, ok := err.(*clickhouse.Exception); assert.True(t, ok) {
-				assert.Equal(t, int32(192), exception.Code)
-			}
-		}
-	}
-	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?password=invalid"); assert.NoError(t, err) {
-		if err := connect.Ping(); assert.Error(t, err) {
-			if exception, ok := err.(*clickhouse.Exception); assert.True(t, ok) {
-				assert.Equal(t, int32(193), exception.Code)
-			}
-		}
-	}
-	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?database=invalid"); assert.NoError(t, err) {
-		if err := connect.Ping(); assert.Error(t, err) {
-			if exception, ok := err.(*clickhouse.Exception); assert.True(t, ok) {
-				assert.Equal(t, int32(81), exception.Code)
-			}
-		}
-	}
 }
 
 func Test_CreateTable(t *testing.T) {

enum.go

@@ -0,0 +1,70 @@
+package clickhouse
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type enum struct {
+	iv map[string]interface{}
+	vi map[interface{}]string
+}
+
+func (e *enum) toIdent(v interface{}) (string, error) {
+	if ident, found := e.vi[v]; found {
+		return ident, nil
+	}
+	return "", fmt.Errorf("invalid Enum value: %v", v)
+}
+func (e enum) toValue(ident string) (interface{}, error) {
+	if value, found := e.iv[ident]; found {
+		return value, nil
+	}
+	return "", fmt.Errorf("invalid Enum ident: %s", ident)
+}
+
+func parseEnum(str string) (enum, error) {
+	var (
+		data     string
+		isEnum16 bool
+	)
+	switch {
+	case strings.HasPrefix(str, "Enum8"):
+		data = str[6:]
+	case strings.HasPrefix(str, "Enum16"):
+		data = str[7:]
+		isEnum16 = true
+	default:
+		return enum{}, fmt.Errorf("'%s' is not Enum type", str)
+	}
+	enum := enum{
+		iv: make(map[string]interface{}),
+		vi: make(map[interface{}]string),
+	}
+	for _, block := range strings.Split(data[:len(data)-1], ",") {
+		parts := strings.Split(block, "=")
+		if len(parts) != 2 {
+			return enum, fmt.Errorf("invalid Enum format: %s", str)
+		}
+		var (
+			ident      = strings.TrimSpace(parts[0])
+			value, err = strconv.ParseInt(strings.TrimSpace(parts[1]), 10, 16)
+		)
+		if err != nil {
+			return enum, fmt.Errorf("invalid Enum value: %v", err)
+		}
+		{
+			var (
+				ident             = ident[1 : len(ident)-1]
+				value interface{} = int16(value)
+			)
+			if !isEnum16 {
+				value = int8(value.(int16))
+			}
+			enum.iv[ident] = value
+			enum.vi[value] = ident
+		}
+	}
+	return enum, nil
+}

enum_test.go

@@ -0,0 +1,54 @@
+package clickhouse
+
+import "testing"
+import "github.com/stretchr/testify/assert"
+
+func Test_Parse_Enum(t *testing.T) {
+	{
+		if enum, err := parseEnum("Enum8('a' = 2, 'b' = 1)"); assert.NoError(t, err) {
+			if value, err := enum.toValue("b"); assert.NoError(t, err) {
+				if v, ok := value.(int8); assert.True(t, ok) {
+					assert.Equal(t, int8(1), v)
+				}
+			}
+			if value, err := enum.toValue("a"); assert.NoError(t, err) {
+				if v, ok := value.(int8); assert.True(t, ok) {
+					assert.Equal(t, int8(2), v)
+				}
+			}
+		}
+
+		if enum, err := parseEnum("Enum8('a' = 2, 'b' = 1)"); assert.NoError(t, err) {
+			if ident, err := enum.toIdent(int8(1)); assert.NoError(t, err) {
+				assert.Equal(t, "b", ident)
+			}
+			if ident, err := enum.toIdent(int8(2)); assert.NoError(t, err) {
+				assert.Equal(t, "a", ident)
+			}
+		}
+	}
+
+	{
+		if enum, err := parseEnum("Enum16('a' = 2, 'b' = 1)"); assert.NoError(t, err) {
+			if value, err := enum.toValue("b"); assert.NoError(t, err) {
+				if v, ok := value.(int16); assert.True(t, ok) {
+					assert.Equal(t, int16(1), v)
+				}
+			}
+			if value, err := enum.toValue("a"); assert.NoError(t, err) {
+				if v, ok := value.(int16); assert.True(t, ok) {
+					assert.Equal(t, int16(2), v)
+				}
+			}
+		}
+
+		if enum, err := parseEnum("Enum16('a' = 2, 'b' = 1)"); assert.NoError(t, err) {
+			if ident, err := enum.toIdent(int16(1)); assert.NoError(t, err) {
+				assert.Equal(t, "b", ident)
+			}
+			if ident, err := enum.toIdent(int16(2)); assert.NoError(t, err) {
+				assert.Equal(t, "a", ident)
+			}
+		}
+	}
+}