fix value

gcpug · Aug 22, 2023 · b996b14 · b996b14
1 parent 35b7b47
commit b996b14
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 12 deletions.
diff --git a/server/database_function.go b/server/database_function.go
@@ -484,8 +484,11 @@ var customFunctions map[string]CustomFunction = map[string]CustomFunction{
 	},
 	"GENERATE_ARRAY": {
 		Func:  sqlite3FnGenerateArray,
-		NArgs: 2,
+		NArgs: -2,
 		ArgTypes: func(vts []ValueType) bool {
+			if len(vts) < 2 || len(vts) > 3 {
+				return false
+			}
 			if vts[0].Code == TCInt64 && vts[1].Code == TCInt64 {
 				if len(vts) == 3 {
 					return vts[2].Code == TCInt64
@@ -495,14 +498,22 @@ var customFunctions map[string]CustomFunction = map[string]CustomFunction{
 			return false
 		},
 		ReturnType: func(vts []ValueType) ValueType {
-			return ValueType{Code: TCArray}
+			return ValueType{
+				Code:      TCArray,
+				ArrayType: &ValueType{Code: TCInt64},
+			}
 		},
 	},
 }
 
-func sqlite3FnGenerateArray(a, b, c int64) []byte {
-	res := make([]byte, 0, c)
-	for i := a; i < b; i += c {
+func sqlite3FnGenerateArray(xs ...int64) []byte {
+	step := int64(1)
+	if len(xs) == 3 {
+		step = xs[2]
+	}
+	a, b := xs[0], xs[1]
+	res := make([]byte, 0, b-a)
+	for i := a; i <= b; i += step {
 		res = append(res, byte(i))
 	}
 	return res

diff --git a/server/database_query_test.go b/server/database_query_test.go
@@ -3431,14 +3431,14 @@ func TestQuery(t *testing.T) {
 				name: "Function_GENERATE_ARRAY",
 				sql:  `SELECT GENERATE_ARRAY(1, 3)`,
 				expected: [][]interface{}{
-					[]interface{}{int64(1), int64(2), int64(3)},
+					{[]int64{1, 2, 3}},
 				},
 			},
 			{
 				name: "Function_GENERATE_ARRAY2",
-				sql:  `SELECT GENERATE_ARRAY(1, 3, 2)`,
+				sql:  `SELECT GENERATE_ARRAY(1, 10, 2)`,
 				expected: [][]interface{}{
-					[]interface{}{int64(1), int64(3)},
+					{[]int64{1, 3, 5, 7, 9}},
 				},
 			},
 		},

diff --git a/server/value.go b/server/value.go
@@ -298,12 +298,19 @@ func (a *ArrayValueDecoder) Scan(src interface{}) error {
 		return nil
 	}
 
-	v, ok := src.(string)
-	if !ok {
+	switch v := src.(type) {
+	case string:
+		return a.UnmarshalJSON([]byte(v))
+	case []uint8:
+		res := make([]int64, len(v))
+		for i, _ := range v {
+			res[i] = int64(v[i])
+		}
+		a.Values = res
+	default:
 		return fmt.Errorf("unexpected type %T for %T", src, a)
 	}
-
-	return a.UnmarshalJSON([]byte(v))
+	return nil
 }
 
 func (a *ArrayValueDecoder) UnmarshalJSON(b []byte) error {