Getters and setters for the JSON type (#325)

* initial commit for JSON type support

* comment on missing setters

* get and set more inputs

* formatting

* add JSON example

Makefile

@@ -7,9 +7,10 @@ install:
 
 .PHONY: examples
 examples:
-	go run examples/simple/main.go
 	go run examples/appender/main.go
+	go run examples/json/main.go
 	go run examples/scalar_udf/main.go
+	go run examples/simple/main.go
 	go run examples/table_udf/main.go
 	go run examples/table_udf_parallel/main.go
 

duckdb_test.go

@@ -279,31 +279,32 @@ func TestQuery(t *testing.T) {
 func TestJSON(t *testing.T) {
 	t.Parallel()
 	db := openDB(t)
-	var data string
 
-	t.Run("select empty JSON", func(t *testing.T) {
-		require.NoError(t, db.QueryRow("SELECT '{}'::JSON").Scan(&data))
-		require.Equal(t, "{}", string(data))
+	t.Run("SELECT an empty JSON", func(t *testing.T) {
+		var res Composite[map[string]any]
+		require.NoError(t, db.QueryRow(`SELECT '{}'::JSON`).Scan(&res))
+		require.Equal(t, 0, len(res.Get()))
 	})
 
-	t.Run("select from marshalled JSON", func(t *testing.T) {
-		val, _ := json.Marshal(struct {
+	t.Run("SELECT a marshalled JSON", func(t *testing.T) {
+		val, err := json.Marshal(struct {
 			Foo string `json:"foo"`
 		}{
 			Foo: "bar",
 		})
-		require.NoError(t, db.QueryRow(`SELECT ?::JSON->>'foo'`, string(val)).Scan(&data))
-		require.Equal(t, "bar", data)
-	})
+		require.NoError(t, err)
 
-	t.Run("select JSON array", func(t *testing.T) {
-		require.NoError(t, db.QueryRow("SELECT json_array('foo', 'bar')").Scan(&data))
-		require.Equal(t, `["foo","bar"]`, data)
+		var res string
+		require.NoError(t, db.QueryRow(`SELECT ?::JSON->>'foo'`, string(val)).Scan(&res))
+		require.Equal(t, "bar", res)
+	})
 
-		var items []string
-		require.NoError(t, json.Unmarshal([]byte(data), &items))
-		require.Equal(t, len(items), 2)
-		require.Equal(t, items, []string{"foo", "bar"})
+	t.Run("SELECT a JSON array", func(t *testing.T) {
+		var res Composite[[]any]
+		require.NoError(t, db.QueryRow(`SELECT json_array('foo', 'bar')`).Scan(&res))
+		require.Equal(t, 2, len(res.Get()))
+		require.Equal(t, "foo", res.Get()[0])
+		require.Equal(t, "bar", res.Get()[1])
 	})
 
 	require.NoError(t, db.Close())
@@ -620,7 +621,7 @@ func TestMultipleStatements(t *testing.T) {
 	var family string
 	err = rows.Scan(&family)
 	require.NoError(t, err)
-	require.Equal(t, "\"anatidae\"", family)
+	require.Equal(t, "anatidae", family)
 	require.False(t, rows.Next())
 	err = rows.Close()
 	require.NoError(t, err)

examples/json/main.go

@@ -0,0 +1,44 @@
+package main
+
+import (
+	"database/sql"
+	"log"
+
+	"github.com/marcboeker/go-duckdb"
+)
+
+var db *sql.DB
+
+func main() {
+	var err error
+	db, err = sql.Open("duckdb", "?access_mode=READ_WRITE")
+
+	check(err)
+	defer db.Close()
+
+	check(db.Ping())
+
+	var jsonArray duckdb.Composite[[]any]
+	row := db.QueryRow(`SELECT json_array('foo', 'bar');`)
+	check(row.Err())
+	check(row.Scan(&jsonArray))
+
+	log.Printf("first element: %s \n", jsonArray.Get()[0])
+	log.Printf("second element: %s \n", jsonArray.Get()[1])
+
+	var jsonMap duckdb.Composite[map[string]any]
+	row = db.QueryRow(`SELECT '{"family": "anatidae", "species": ["duck", "goose"], "coolness": 42.42}'::JSON;`)
+	check(row.Err())
+	check(row.Scan(&jsonMap))
+
+	log.Printf("family: %s", jsonMap.Get()["family"])
+	log.Printf("species: %s", jsonMap.Get()["species"])
+	log.Printf("coolness: %f", jsonMap.Get()["coolness"])
+}
+
+func check(args ...interface{}) {
+	err := args[len(args)-1]
+	if err != nil {
+		panic(err)
+	}
+}

examples/simple/main.go

@@ -22,17 +22,15 @@ type user struct {
 func main() {
 	var err error
 	db, err = sql.Open("duckdb", "?access_mode=READ_WRITE")
-	if err != nil {
-		log.Fatal(err)
-	}
+	check(err)
 	defer db.Close()
 
 	check(db.Ping())
 
 	setting := db.QueryRowContext(context.Background(), "SELECT current_setting('access_mode')")
-	var am string
-	check(setting.Scan(&am))
-	log.Printf("DB opened with access mode %s", am)
+	var accessMode string
+	check(setting.Scan(&accessMode))
+	log.Printf("DB opened with access mode %s", accessMode)
 
 	check(db.ExecContext(context.Background(), "CREATE TABLE users(name VARCHAR, age INTEGER, height FLOAT, awesome BOOLEAN, bday DATE)"))
 	check(db.ExecContext(context.Background(), "INSERT INTO users VALUES('marc', 99, 1.91, true, '1970-01-01')"))

type.go

@@ -97,3 +97,5 @@ var typeToStringMap = map[Type]string{
 	TYPE_VARINT:       "VARINT",
 	TYPE_SQLNULL:      "SQLNULL",
 }
+
+const aliasJSON = "JSON"

types_test.go

@@ -55,6 +55,11 @@ type testTypesRow struct {
 	Array_col        Composite[[3]int32]
 	Time_tz_col      time.Time
 	Timestamp_tz_col time.Time
+	Json_col_map     Composite[map[string]any]
+	Json_col_array   Composite[[]any]
+	Json_col_string  string
+	Json_col_bool    bool
+	Json_col_float64 float64
 }
 
 const testTypesTableSQL = `CREATE TABLE test (
@@ -85,7 +90,12 @@ const testTypesTableSQL = `CREATE TABLE test (
 	Map_col MAP(INTEGER, VARCHAR),
 	Array_col INTEGER[3],
 	Time_tz_col TIMETZ,
-	Timestamp_tz_col TIMESTAMPTZ
+	Timestamp_tz_col TIMESTAMPTZ,
+	Json_col_map JSON,
+	Json_col_array JSON,
+	Json_col_string JSON,
+	Json_col_bool JSON,
+	Json_col_float64 JSON
 )`
 
 func (r *testTypesRow) toUTC() {
@@ -129,6 +139,15 @@ func testTypesGenerateRow[T require.TestingT](t T, i int) testTypesRow {
 	arrayCol := Composite[[3]int32]{
 		[3]int32{int32(i), int32(i), int32(i)},
 	}
+	jsonMapCol := Composite[map[string]any]{
+		map[string]any{
+			"hello": float64(42),
+			"world": float64(84),
+		},
+	}
+	jsonArrayCol := Composite[[]any]{
+		[]any{"hello", "world"},
+	}
 
 	return testTypesRow{
 		i%2 == 1,
@@ -159,6 +178,11 @@ func testTypesGenerateRow[T require.TestingT](t T, i int) testTypesRow {
 		arrayCol,
 		timeTZ,
 		ts,
+		jsonMapCol,
+		jsonArrayCol,
+		varcharCol,
+		i%2 == 1,
+		float64(i),
 	}
 }
 
@@ -208,7 +232,12 @@ func testTypes[T require.TestingT](t T, c *Connector, a *Appender, expectedRows
 			r.Map_col,
 			r.Array_col.Get(),
 			r.Time_tz_col,
-			r.Timestamp_tz_col)
+			r.Timestamp_tz_col,
+			r.Json_col_map.Get(),
+			r.Json_col_array.Get(),
+			r.Json_col_string,
+			r.Json_col_bool,
+			r.Json_col_float64)
 		require.NoError(t, err)
 	}
 	require.NoError(t, a.Flush())
@@ -248,7 +277,12 @@ func testTypes[T require.TestingT](t T, c *Connector, a *Appender, expectedRows
 			&r.Map_col,
 			&r.Array_col,
 			&r.Time_tz_col,
-			&r.Timestamp_tz_col)
+			&r.Timestamp_tz_col,
+			&r.Json_col_map,
+			&r.Json_col_array,
+			&r.Json_col_string,
+			&r.Json_col_bool,
+			&r.Json_col_float64)
 		require.NoError(t, err)
 		actualRows = append(actualRows, r)
 	}
@@ -828,3 +862,30 @@ func TestInterval(t *testing.T) {
 
 	require.NoError(t, db.Close())
 }
+
+func TestJSONType(t *testing.T) {
+	t.Parallel()
+	db := openDB(t)
+
+	_, err := db.Exec(`CREATE TABLE test (c1 STRUCT(index INTEGER))`)
+	require.NoError(t, err)
+
+	_, err = db.Exec(`INSERT INTO test VALUES ({index: 1}), ({index: 2}), ({index: 2}), ({index: 3}), ({index: 3}), ({index: 3})`)
+	require.NoError(t, err)
+
+	// Verify results.
+	row := db.QueryRowContext(context.Background(), `
+	SELECT json_group_object(t2.status, t2.count) AS result
+	FROM (
+		SELECT json_extract(c1, '$.index') AS status, COUNT(*) AS count
+		FROM test
+		GROUP BY status
+	) AS t2`)
+
+	var res Composite[map[string]any]
+	require.NoError(t, row.Scan(&res))
+	require.Equal(t, float64(1), res.Get()["1"])
+	require.Equal(t, float64(2), res.Get()["2"])
+	require.Equal(t, float64(3), res.Get()["3"])
+	require.NoError(t, db.Close())
+}

vector.go

@@ -46,6 +46,15 @@ func (vec *vector) init(logicalType C.duckdb_logical_type, colIdx int) error {
 		return addIndexToError(unsupportedTypeError(name), colIdx)
 	}
 
+	cStr := C.duckdb_logical_type_get_alias(logicalType)
+	alias := C.GoString(cStr)
+	C.duckdb_free(unsafe.Pointer(cStr))
+	switch alias {
+	case aliasJSON:
+		vec.initJSON()
+		return nil
+	}
+
 	switch t {
 	case TYPE_BOOLEAN:
 		initBool(vec)
@@ -266,7 +275,7 @@ func (vec *vector) initBytes(t Type) {
 		if vec.getNull(rowIdx) {
 			return nil
 		}
-		return vec.getCString(rowIdx)
+		return vec.getBytes(rowIdx)
 	}
 	vec.setFn = func(vec *vector, rowIdx C.idx_t, val any) error {
 		if val == nil {
@@ -278,6 +287,23 @@ func (vec *vector) initBytes(t Type) {
 	vec.Type = t
 }
 
+func (vec *vector) initJSON() {
+	vec.getFn = func(vec *vector, rowIdx C.idx_t) any {
+		if vec.getNull(rowIdx) {
+			return nil
+		}
+		return vec.getJSON(rowIdx)
+	}
+	vec.setFn = func(vec *vector, rowIdx C.idx_t, val any) error {
+		if val == nil {
+			vec.setNull(rowIdx)
+			return nil
+		}
+		return setJSON(vec, rowIdx, val)
+	}
+	vec.Type = TYPE_VARCHAR
+}
+
 func (vec *vector) initDecimal(logicalType C.duckdb_logical_type, colIdx int) error {
 	vec.decimalWidth = uint8(C.duckdb_decimal_width(logicalType))
 	vec.decimalScale = uint8(C.duckdb_decimal_scale(logicalType))

vector_getters.go

@@ -6,6 +6,7 @@ package duckdb
 import "C"
 
 import (
+	"encoding/json"
 	"math/big"
 	"time"
 	"unsafe"
@@ -109,7 +110,7 @@ func (vec *vector) getHugeint(rowIdx C.idx_t) *big.Int {
 	return hugeIntToNative(hugeInt)
 }
 
-func (vec *vector) getCString(rowIdx C.idx_t) any {
+func (vec *vector) getBytes(rowIdx C.idx_t) any {
 	cStr := getPrimitive[duckdb_string_t](vec, rowIdx)
 
 	var blob []byte
@@ -127,6 +128,13 @@ func (vec *vector) getCString(rowIdx C.idx_t) any {
 	return blob
 }
 
+func (vec *vector) getJSON(rowIdx C.idx_t) any {
+	bytes := vec.getBytes(rowIdx).(string)
+	var value any
+	_ = json.Unmarshal([]byte(bytes), &value)
+	return value
+}
+
 func (vec *vector) getDecimal(rowIdx C.idx_t) Decimal {
 	var val *big.Int
 	switch vec.internalType {

vector_setters.go

@@ -6,6 +6,7 @@ package duckdb
 import "C"
 
 import (
+	"encoding/json"
 	"math/big"
 	"reflect"
 	"strconv"
@@ -258,6 +259,14 @@ func setBytes[S any](vec *vector, rowIdx C.idx_t, val S) error {
 	return nil
 }
 
+func setJSON[S any](vec *vector, rowIdx C.idx_t, val S) error {
+	bytes, err := json.Marshal(val)
+	if err != nil {
+		return err
+	}
+	return setBytes(vec, rowIdx, bytes)
+}
+
 func setDecimal[S any](vec *vector, rowIdx C.idx_t, val S) error {
 	switch vec.internalType {
 	case TYPE_SMALLINT: