GODRIVER-2119 Implement Text(Un)Marshaler for primitive.ObjectID (#715)

Author: Luberry

bson/primitive/objectid.go

@@ -12,6 +12,7 @@ package primitive
 import (
 	"bytes"
 	"crypto/rand"
+	"encoding"
 	"encoding/binary"
 	"encoding/hex"
 	"encoding/json"
@@ -34,6 +35,9 @@ var NilObjectID ObjectID
 var objectIDCounter = readRandomUint32()
 var processUnique = processUniqueBytes()
 
+var _ encoding.TextMarshaler = ObjectID{}
+var _ encoding.TextUnmarshaler = ObjectID{}
+
 // NewObjectID generates a new ObjectID.
 func NewObjectID() ObjectID {
 	return NewObjectIDFromTimestamp(time.Now())
@@ -94,6 +98,19 @@ func IsValidObjectID(s string) bool {
 	return err == nil
 }
 
+// MarshalText returns the ObjectID as UTF-8-encoded text. Implementing this allows us to use ObjectID
+// as a map key when marshalling JSON. See https://pkg.go.dev/encoding#TextMarshaler
+func (id ObjectID) MarshalText() ([]byte, error) {
+	return []byte(id.Hex()), nil
+}
+
+// UnmarshalText populates the byte slice with the ObjectID. Implementing this allows us to use ObjectID
+// as a map key when unmarshalling JSON. See https://pkg.go.dev/encoding#TextUnmarshaler
+func (id ObjectID) UnmarshalText(b []byte) error {
+	id, err := ObjectIDFromHex(string(b))
+	return err
+}
+
 // MarshalJSON returns the ObjectID as a string
 func (id ObjectID) MarshalJSON() ([]byte, error) {
 	return json.Marshal(id.Hex())

bson/primitive/objectid_test.go

@@ -172,6 +172,38 @@ func TestCounterOverflow(t *testing.T) {
 	require.Equal(t, uint32(0), objectIDCounter)
 }
 
+func TestObjectID_MarshalJSONMap(t *testing.T) {
+	type mapOID struct {
+		Map map[ObjectID]string
+	}
+
+	oid := NewObjectID()
+	expectedJSON := []byte(fmt.Sprintf(`{"Map":{%q:"foo"}}`, oid.Hex()))
+	data := mapOID{
+		Map: map[ObjectID]string{oid: "foo"},
+	}
+
+	out, err := json.Marshal(&data)
+	require.NoError(t, err)
+	require.Equal(t, expectedJSON, out)
+}
+
+func TestObjectID_UnmarshalJSONMap(t *testing.T) {
+	type mapOID struct {
+		Map map[ObjectID]string
+	}
+	oid := NewObjectID()
+	mapOIDJSON := []byte(fmt.Sprintf(`{"Map":{%q:"foo"}}`, oid.Hex()))
+	expectedData := mapOID{
+		Map: map[ObjectID]string{oid: "foo"},
+	}
+
+	data := mapOID{}
+	err := json.Unmarshal(mapOIDJSON, &data)
+	require.NoError(t, err)
+	require.Equal(t, expectedData, data)
+}
+
 func TestObjectID_UnmarshalJSON(t *testing.T) {
 	oid := NewObjectID()