使用 mongo-driver 自定义 BSON 编组和解组

Custom BSON marshal and unmarshal using mongo-driver

我有一个如下所示的结构字段。我还在数据库中存储了相同结构的原始 protobuf。现在每次获取或保存数据到 mongo。当我想保存到数据库时,我必须从 proto 更新 ReallyBigRaw,当获取时,我必须将 ReallyBigRaw 解组为 ReallyBigObj 以给出响应。有没有一种方法可以实现一些接口或提供一些回调函数,以便 mongo 驱动程序在从数据库保存或获取数据之前自动执行此操作。

此外,我使用的是官方 golang mongo 驱动程序而不是 mgo,我已经阅读了一些可以在 mgo golang 库中完成的答案。

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
}

我实现了 MarshalerUnmarshaler 接口。由于 mongo 驱动程序调用 MarshalBSONUnmarshalBSON 如果类型实现 MarshalerUnmarshaler 我们也会陷入无限循环。为了避免这种情况,我们创建了 type 的别名。 Golang 中的别名仅继承字段而不是方法,因此我们最终调用正常的 bson.Marshalbson.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
}