GODRIVER-1199 Add UnmarshalJSON and MarshalJSON for primitive.Decimal128 (#556)

Author: iwysiu

bson/primitive/decimal.go

@@ -10,6 +10,7 @@
 package primitive
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"math/big"
@@ -211,6 +212,49 @@ func (d Decimal128) IsZero() bool {
 	return d.h == 0 && d.l == 0
 }
 
+// MarshalJSON returns Decimal128 as a string.
+func (d Decimal128) MarshalJSON() ([]byte, error) {
+	return json.Marshal(d.String())
+}
+
+// UnmarshalJSON creates a primitive.Decimal128 from a JSON string, an extended JSON $numberDecimal value, or the string
+// "null". If b is a JSON string or extended JSON value, d will have the value of that string, and if b is "null", d will
+// be unchanged.
+func (d *Decimal128) UnmarshalJSON(b []byte) error {
+	// Ignore "null" to keep parity with the standard library. Decoding a JSON null into a non-pointer Decimal128 field
+	// will leave the field unchanged. For pointer values, encoding/json will set the pointer to nil and will not
+	// enter the UnmarshalJSON hook.
+	if string(b) == "null" {
+		return nil
+	}
+
+	var res interface{}
+	err := json.Unmarshal(b, &res)
+	if err != nil {
+		return err
+	}
+	str, ok := res.(string)
+
+	// Extended JSON
+	if !ok {
+		m, ok := res.(map[string]interface{})
+		if !ok {
+			return errors.New("not an extended JSON Decimal128: expected document")
+		}
+		d128, ok := m["$numberDecimal"]
+		if !ok {
+			return errors.New("not an extended JSON Decimal128: expected key $numberDecimal")
+		}
+		str, ok = d128.(string)
+		if !ok {
+			return errors.New("not an extended JSON Decimal128: expected decimal to be string")
+		}
+	}
+
+	*d, err = ParseDecimal128(str)
+	return err
+}
+
 func divmod(h, l uint64, div uint32) (qh, ql uint64, rem uint32) {
 	div64 := uint64(div)
 	a := h >> 32

bson/primitive/decimal_test.go

@@ -1,9 +1,13 @@
 package primitive
 
 import (
-	"github.com/stretchr/testify/require"
+	"encoding/json"
+	"fmt"
 	"math/big"
 	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 )
 
 type bigIntTestCase struct {
@@ -144,3 +148,61 @@ func TestParseDecimal128(t *testing.T) {
 		}
 	}
 }
+
+func TestDecimal128_JSON(t *testing.T) {
+	t.Run("roundTrip", func(t *testing.T) {
+		decimal := NewDecimal128(0x3040000000000000, 12345)
+		bytes, err := json.Marshal(decimal)
+		assert.Nil(t, err, "json.Marshal error: %v", err)
+		got := NewDecimal128(0, 0)
+		err = json.Unmarshal([]byte(bytes), &got)
+		assert.Nil(t, err, "json.Unmarshal error: %v", err)
+		assert.Equal(t, decimal.h, got.h, "expected h: %v got: %v", decimal.h, got.h)
+		assert.Equal(t, decimal.l, got.l, "expected l: %v got: %v", decimal.l, got.l)
+	})
+	t.Run("unmarshal extendedJSON", func(t *testing.T) {
+		want := NewDecimal128(0x3040000000000000, 12345)
+		extJSON := fmt.Sprintf(`{"$numberDecimal": %q}`, want.String())
+
+		got := NewDecimal128(0, 0)
+		err := json.Unmarshal([]byte(extJSON), &got)
+		assert.Nil(t, err, "json.Unmarshal error: %v", err)
+		assert.Equal(t, want.h, got.h, "expected h: %v got: %v", want.h, got.h)
+		assert.Equal(t, want.l, got.l, "expected l: %v got: %v", want.l, got.l)
+	})
+	t.Run("unmarshal null", func(t *testing.T) {
+		want := NewDecimal128(0, 0)
+		extJSON := `null`
+
+		got := NewDecimal128(0, 0)
+		err := json.Unmarshal([]byte(extJSON), &got)
+		assert.Nil(t, err, "json.Unmarshal error: %v", err)
+		assert.Equal(t, want.h, got.h, "expected h: %v got: %v", want.h, got.h)
+		assert.Equal(t, want.l, got.l, "expected l: %v got: %v", want.l, got.l)
+	})
+	t.Run("unmarshal", func(t *testing.T) {
+		cases := append(bigIntTestCases,
+			[]bigIntTestCase{
+				{s: "-0001231.453454000000565600000000E-21", h: 0xafe6000003faa269, l: 0x81cfeceaabdb1800},
+				{s: "12345E+21", h: 0x306a000000000000, l: 12345},
+				{s: "0.10000000000000000000000000000000000000000001", remark: "parse fail"},
+				{s: ".125e1", h: 0x303c000000000000, l: 125},
+				{s: ".125", h: 0x303a000000000000, l: 125},
+			}...)
+		for _, c := range cases {
+			input := fmt.Sprintf(`{"foo": %q}`, c.s)
+			var got map[string]Decimal128
+			err := json.Unmarshal([]byte(input), &got)
+
+			switch c.remark {
+			case "overflow", "parse fail":
+				assert.NotNil(t, err, "expected Unmarshal error, got nil")
+			default:
+				assert.Nil(t, err, "Unmarshal error: %v", err)
+				gotDecimal := got["foo"]
+				assert.Equal(t, c.h, gotDecimal.h, "expected h: %v got: %v", c.h, gotDecimal.l)
+				assert.Equal(t, c.l, gotDecimal.l, "expected l: %v got: %v", c.l, gotDecimal.h)
+			}
+		}
+	})
+}