Fix JSON decode (#25)

* Call String() directly from array and document type

* fix test

* Add experimental json parser

* Add workaround for bson.M unmarshal

Author: haohanyang

pkg/models/column.go

@@ -130,13 +130,22 @@ func (c *Column) AppendValue(rv bson.RawValue) error {
 			return fmt.Errorf("field %s should have type %s, but got %s", c.Name, c.Type().ItemTypeString(), rv.Type.String())
 		}
 
-		c.Field.Append(pointer(rv.Document().String()))
+		v, err := rawDocToJson(rv)
+		if err != nil {
+			return err
+		}
+
+		c.Field.Append(&v)
 	case bson.TypeArray:
 		if c.Type() != data.FieldTypeNullableString {
 			return fmt.Errorf("field %s should have type %s, but got %s", c.Name, c.Type().ItemTypeString(), rv.Type.String())
 		}
 
-		c.Field.Append(pointer(rv.Array().String()))
+		v, err := rawArrayToJson(rv)
+		if err != nil {
+			return err
+		}
+		c.Field.Append(&v)
 
 	default:
 		if c.Type() != data.FieldTypeNullableString {
@@ -158,7 +167,7 @@ func (c *Column) Type() data.FieldType {
 	return c.Field.Type()
 }
 
-func NewColumn(rowIndex int, element bson.RawElement) *Column {
+func NewColumn(rowIndex int, element bson.RawElement) (*Column, error) {
 	key := element.Key()
 	value := element.Value()
 	var field *data.Field
@@ -194,11 +203,21 @@ func NewColumn(rowIndex int, element bson.RawElement) *Column {
 
 	case bson.TypeEmbeddedDocument:
 		field = data.NewField(key, nil, make([]*string, rowIndex+1))
-		field.Set(rowIndex, pointer(value.Document().String()))
+
+		v, err := rawDocToJson(value)
+		if err != nil {
+			return nil, err
+		}
+		field.Set(rowIndex, &v)
 
 	case bson.TypeArray:
 		field = data.NewField(key, nil, make([]*string, rowIndex+1))
-		field.Set(rowIndex, pointer(value.Array().String()))
+
+		v, err := rawArrayToJson(value)
+		if err != nil {
+			return nil, err
+		}
+		field.Set(rowIndex, &v)
 
 	default:
 		field = data.NewField(key, nil, make([]*string, rowIndex+1))
@@ -209,7 +228,7 @@ func NewColumn(rowIndex int, element bson.RawElement) *Column {
 		Name:      key,
 		Field:     field,
 		BsonTypes: []bsontype.Type{value.Type},
-	}
+	}, nil
 }
 
 // Convert array and embedded document type values to json.RawMessage if allowed

pkg/models/utils.go

@@ -0,0 +1,37 @@
+package models
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"go.mongodb.org/mongo-driver/bson"
+)
+
+func rawArrayToJson(value bson.RawValue) (string, error) {
+	var bsonDoc bson.M
+	err := bson.UnmarshalExtJSON([]byte(fmt.Sprintf(`{"data":%s}`, value.String())), true, &bsonDoc)
+	if err != nil {
+		return "", err
+	}
+
+	rawBytes, err := json.Marshal(bsonDoc["data"])
+	if err != nil {
+		return "", err
+	}
+	return string(rawBytes), nil
+}
+
+func rawDocToJson(value bson.RawValue) (string, error) {
+	var bsonMap bson.M
+
+	err := bson.UnmarshalExtJSON([]byte(value.String()), true, &bsonMap)
+	if err != nil {
+		return "", err
+	}
+
+	rawBytes, err := json.Marshal(bsonMap)
+	if err != nil {
+		return "", err
+	}
+	return string(rawBytes), nil
+}

pkg/plugin/query.go

@@ -81,7 +81,11 @@ func CreateTableFramesFromQuery(ctx context.Context, tableName string, cursor *m
 				if element.Value().Type == bson.TypeNull {
 					continue
 				}
-				columns[name] = models.NewColumn(rowIndex, element)
+				nc, err := models.NewColumn(rowIndex, element)
+				if err != nil {
+					return nil, err
+				}
+				columns[name] = nc
 			}
 		}
 

pkg/plugin/query_test.go

@@ -740,14 +740,22 @@ func TestCreateTableFramesFromQuery(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		v, _ := frame.Fields[0].ConcreteAt(0)
-		var doc bson.M
-		bson.UnmarshalExtJSON(v.(json.RawMessage), true, &doc)
-		assertEq(t, doc["0"], int32(1))
-		assertEq(t, doc["1"], "bar")
+		v, notNull := frame.Fields[0].ConcreteAt(0)
+		if !notNull {
+			t.Fatal("null value")
+		}
+
+		var doc bson.A
+		err = bson.UnmarshalExtJSON(v.(json.RawMessage), true, &doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		assertEq(t, doc[0], int32(1))
+		assertEq(t, doc[1], "bar")
 
-		docIn := doc["2"].(bson.M)
-		assertEq(t, docIn["baz"], int32(1))
+		docIn := doc[2].(bson.D)
+		assertEq(t, docIn.Map()["baz"], int32(1))
 	})
 
 	t.Run("array and embedded docs can exist in the same field", func(t *testing.T) {