Search code examples
mongodbgoprotocol-buffersmarshallingbson

Custom BSON marshal and unmarshal using mongo-driver


I have a struct field like this below. I also store raw protobuf of the same struct in db. Now every time fetch or save data to mongo. I have to update ReallyBigRaw, from the proto when I want to save to DB and when fetch I have to unmarshal ReallyBigRaw to ReallyBigObj to give out responses. Is there a way I can implement some interface or provide some callback functions so that the mongo driver does this automatically before saving or fetching data from DB.

Also, I am using the offical golang mongo driver not mgo, I have read some answers where can be done in mgo golang library.

import (
    "github.com/golang/protobuf/jsonpb"
    "go.mongodb.org/mongo-driver/bson"
    "go.mongodb.org/mongo-driver/bson/primitive"

    proto "github.com/dinesh/api/go"
)

type ReallyBig struct {
    ID           string                 `bson:"_id,omitempty"`
    DraftID      string                 `bson:"draft_id,omitempty"`
// Marshaled ReallyBigObj proto to map[string]interface{} stored in DB
    ReallyBigRaw map[string]interface{} `bson:"raw,omitempty"`
    ReallyBigObj *proto.ReallyBig       `bson:"-"`
    CreatedAt    primitive.DateTime     `bson:"created_at,omitempty"`
    UpdatedAt    primitive.DateTime     `bson:"updated_at,omitempty"`
}

func (r *ReallyBig) GetProto() (*proto.ReallyBig, error) {
    if r.ReallyBigObj != nil {
        return r.ReallyBigObj, nil
    }
    Obj, err := getProto(r.ReallyBigRaw)
    if err != nil {
        return nil, err
    }
    r.ReallyBigObj = Obj
    return r.ReallyBigObj, nil
}

func getRaw(r *proto.ReallyBig) (map[string]interface{}, error) {
    m := jsonpb.Marshaler{}
    b := bytes.NewBuffer([]byte{})

    // marshals proto to json format
    err := m.Marshal(b, r)
    if err != nil {
        return nil, err
    }
    var raw map[string]interface{}
    // unmarshal the raw data to an interface
    err = json.Unmarshal(b.Bytes(), &raw)
    if err != nil {
        return nil, err
    }
    return raw, nil
}

func getProto(raw map[string]interface{}) (*proto.ReallyBig, error) {
    b, err := json.Marshal(raw)
    if err != nil {
        return nil, err
    }
    u := jsonpb.Unmarshaler{}
    var reallyBigProto proto.ReallyBig
    err = u.Unmarshal(bytes.NewReader(b), &recipeProto)
    if err != nil {
        return nil, err
    }
    return &reallyBigProto, nil
}

Solution

  • I implemented the Marshaler and Unmarshaler interface. Since mongo driver calls MarshalBSON and UnmarshalBSON if the type implements Marshaler and Unmarshaler we also end up in infinite loop. To avoid that we create a Alias of the type. Alias in Golang inherit only the fields not the methods so we end up calling normal bson.Marshal and bson.Unmarshal

    func (r *ReallyBig) MarshalBSON() ([]byte, error) {
        type ReallyBigAlias ReallyBig
        reallyBigRaw, err := getRaw(r.ReallyBigObj)
        if err != nil {
            return nil, err
        }
        r.ReallyBigRaw = reallyBigRaw
        return bson.Marshal((*ReallyBigAlias)(r))
    }
    
    func (r *ReallyBig) UnmarshalBSON(data []byte) error {
        type ReallyBigAlias ReallyBig
        err := bson.Unmarshal(data, (*ReallyBigAlias)(r))
        if err != nil {
            return err
        }
        reallyBigProto, err := getProto(r.ReallyBigRaw)
        if err != nil {
            return err
        }
        r.ReallyBigObj = reallyBigProto
        return nil
    }