Fix parsing int2 values, add a test for typecasting VALUES(...)

BemiHQ · Dec 23, 2024 · 819143c · 819143c
1 parent 77fa0e8
commit 819143c
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/src/main.go b/src/main.go
@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-const VERSION = "0.28.3"
+const VERSION = "0.28.4"
 
 func main() {
 	config := LoadConfig()

diff --git a/src/query_handler.go b/src/query_handler.go
@@ -444,6 +444,9 @@ func (queryHandler *QueryHandler) generateDataRow(rows *sql.Rows, cols []*sql.Co
 	valuePtrs := make([]interface{}, len(cols))
 	for i, col := range cols {
 		switch col.ScanType().String() {
+		case "int16":
+			var value sql.NullInt16
+			valuePtrs[i] = &value
 		case "int32":
 			var value sql.NullInt32
 			valuePtrs[i] = &value
@@ -487,6 +490,12 @@ func (queryHandler *QueryHandler) generateDataRow(rows *sql.Rows, cols []*sql.Co
 	var values [][]byte
 	for i, valuePtr := range valuePtrs {
 		switch value := valuePtr.(type) {
+		case *sql.NullInt16:
+			if value.Valid {
+				values = append(values, []byte(strconv.Itoa(int(value.Int16))))
+			} else {
+				values = append(values, nil)
+			}
 		case *sql.NullInt32:
 			if value.Valid {
 				values = append(values, []byte(strconv.Itoa(int(value.Int32))))

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -478,6 +478,10 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"word"},
 			"values":      {"abort"},
 		},
+		"SELECT t.x FROM (VALUES (1::int2, 'pg_type'::regclass)) t(x, y)": {
+			"description": {"x"},
+			"values":      {"1"},
+		},
 		// SHOW
 		"SHOW search_path": {
 			"description": {"search_path"},