Repository: zeromicro/go-zero
Pull request #3780 (2023-12-11 19:22:24)

feature: add a mongo registry option to convert type easier.
Issue #3734.

If users want to convert their custom type to a MongoDB primitive type, they have to define a WithRegistry() function and set up the Encoder/Decoder.

`mongo/mongo.go`
```go
package mongo

// ...

func WithRegistry(registry *bsoncodec.Registry) mon.Option {
	return func(opts *mopt.ClientOptions) {
		opts.SetRegistry(registry)
	}
}

// Mongo type Encoder/Decoder
var MongoDecimalEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {...}
var MongoDecimalDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {...}
var MongoDateTimeEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {...}
var MongoDateTimeDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {...}
```

And `RegisterTypeEncoder` and `RegisterTypeDecoder` like below:

`main.go`
```go
package main

// ...

type User struct {
	ID        primitive.ObjectID `bson:"_id,omitempty" json:"id,omitempty"`
	Name      string             `bson:"name,omitempty" json:"name,omitempty"`
	Money     decimal.Decimal    `bson:"money,omitempty" json:"money,omitempty"`
	CreatedAt time.Time          `bson:"created_at,omitempty" json:"created_at,omitempty"`
}

func main() {
	// ...
	registry := bson.NewRegistry()
	registry.RegisterTypeEncoder(reflect.TypeOf(decimal.Decimal{}), mongo.MongoDecimalEncoder)
	registry.RegisterTypeDecoder(reflect.TypeOf(decimal.Decimal{}), mongo.MongoDecimalDecoder)
	registry.RegisterTypeEncoder(reflect.TypeOf(time.Time{}), mongo.MongoDateTimeEncoder)
	registry.RegisterTypeDecoder(reflect.TypeOf(time.Time{}), mongo.MongoDateTimeDecoder)
	opts := mongo.WithRegistry(registry)
	conn := mon.MustNewModel("mongodb://localhost:27017", "test", "collection", opts)

	// find
	var users []User
	err := conn.Find(ctx, &users, bson.D{{}})
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(users)
}
```

In order to make this easier, I added a `WithRegister()` function so that users don't need to implement it themselves. I also used a RegisterType struct to store custom types and their respective Encoder/Decoder functions.
In this feature, users only need to define the `Encoder/Decoder` and use `WithRegistry()`.

`mongo/mongo.go`
```go
package mongo

// ...

// No more use
// func WithRegistry(registry *bsoncodec.Registry) mon.Option {
// 	return func(opts *mopt.ClientOptions) {
// 		opts.SetRegistry(registry)
// 	}
// }

// Mongo type Encoder/Decoder
var MongoDecimalEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {...}
var MongoDecimalDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {...}
var MongoDateTimeEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {...}
var MongoDateTimeDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {...}
```

`main.go`
```go
package main

// ...

type User struct {
	ID        primitive.ObjectID `bson:"_id,omitempty" json:"id,omitempty"`
	Name      string             `bson:"name,omitempty" json:"name,omitempty"`
	Money     decimal.Decimal    `bson:"money,omitempty" json:"money,omitempty"`
	CreatedAt time.Time          `bson:"created_at,omitempty" json:"created_at,omitempty"`
}

func main() {
	// ...
	registerType := []mon.RegisterType{
		{
			ValueType: reflect.TypeOf(decimal.Decimal{}),
			Encoder:   mongo.MongoDecimalEncoder,
			Decoder:   mongo.MongoDecimalDecoder,
		},
		{
			ValueType: reflect.TypeOf(time.Time{}),
			Encoder:   mongo.MongoDateTimeEncoder,
			Decoder:   mongo.MongoDateTimeDecoder,
		},
	}
	opts := mon.WithRegistry(registerType...)
	conn := mon.MustNewModel("mongodb://localhost:27017", "test", "collection", opts)
	
	// find
	var users []User
	err := conn.Find(ctx, &users, bson.D{{}})
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(users)
}
```

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zeromicro/go-zero/pull/3780?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) Report
> Merging [#3780](https://app.codecov.io/gh/zeromicro/go-zero/pull/3780?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) (62ff333) into [master](https://app.codecov.io/gh/zeromicro/go-zero/commit/a1bbac3c6c2009dc8fc3b70a0ac2fd3406d8e322?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) (a1bbac3) will **increase** coverage by `0.00%`.
> Report is 8 commits behind head on master.
> The diff coverage is `100.00%`.

<details><summary>Additional details and impacted files</summary>



| [Files](https://app.codecov.io/gh/zeromicro/go-zero/pull/3780?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) | Coverage Δ | |
|---|---|---|
| [core/stores/mon/options.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/3780?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-Y29yZS9zdG9yZXMvbW9uL29wdGlvbnMuZ28=) | `100.00% <100.00%> (ø)` | |

</details>
kevwan: Thanks for your PR!

I recommend not to import decimal package only because of unit test.
POABOB: Do I need to update the go-zero document that let people know this function?
POABOB: @kevwan Do I need to do anything extra?

Patch:
diff --git a/core/stores/mon/options.go b/core/stores/mon/options.go
index 4097328f..14a6666a 100644
--- a/core/stores/mon/options.go
+++ b/core/stores/mon/options.go
@@ -1,9 +1,12 @@
 package mon
 
 import (
+	"reflect"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/syncx"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
@@ -20,6 +23,13 @@ type (
 
 	// Option defines the method to customize a mongo model.
 	Option func(opts *options)
+
+	// RegisterType A struct store With custom type and Encoder/Decoder
+	RegisterType struct {
+		ValueType reflect.Type
+		Encoder   bsoncodec.ValueEncoder
+		Decoder   bsoncodec.ValueDecoder
+	}
 )
 
 // DisableLog disables logging of mongo commands, includes info and slow logs.
@@ -50,3 +60,15 @@ func WithTimeout(timeout time.Duration) Option {
 		opts.SetTimeout(timeout)
 	}
 }
+
+// WithRegistry set the Registry to convert custom type to mongo primitive type more easily.
+func WithRegistry(registerType ...RegisterType) Option {
+	return func(opts *options) {
+		registry := bson.NewRegistry()
+		for _, v := range registerType {
+			registry.RegisterTypeEncoder(v.ValueType, v.Encoder)
+			registry.RegisterTypeDecoder(v.ValueType, v.Decoder)
+		}
+		opts.SetRegistry(registry)
+	}
+}


Test patch:
diff --git a/core/stores/mon/options_test.go b/core/stores/mon/options_test.go
index ee944f94..f5a858ad 100644
--- a/core/stores/mon/options_test.go
+++ b/core/stores/mon/options_test.go
@@ -1,10 +1,14 @@
 package mon
 
 import (
+	"fmt"
+	"reflect"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/bson/bsonrw"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
@@ -51,3 +55,56 @@ func TestDisableInfoLog(t *testing.T) {
 	assert.False(t, logMon.True())
 	assert.True(t, logSlowMon.True())
 }
+
+func TestWithRegistryForTimestampRegisterType(t *testing.T) {
+	opts := mopt.Client()
+
+	// mongoDateTimeEncoder allow user convert time.Time to primitive.DateTime.
+	var mongoDateTimeEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error {
+		// Use reflect, determine if it can be converted to time.Time.
+		dec, ok := value.Interface().(time.Time)
+		if !ok {
+			return fmt.Errorf("value %v to encode is not of type time.Time", value)
+		}
+		return w.WriteDateTime(dec.Unix())
+	}
+
+	// mongoDateTimeEncoder allow user convert primitive.DateTime to time.Time.
+	var mongoDateTimeDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error {
+		primTime, err := r.ReadDateTime()
+		if err != nil {
+			return fmt.Errorf("error reading primitive.DateTime from ValueReader: %v", err)
+		}
+		value.Set(reflect.ValueOf(time.Unix(primTime, 0)))
+		return nil
+	}
+
+	registerType := []RegisterType{
+		{
+			ValueType: reflect.TypeOf(time.Time{}),
+			Encoder:   mongoDateTimeEncoder,
+			Decoder:   mongoDateTimeDecoder,
+		},
+	}
+	WithRegistry(registerType...)(opts)
+
+	for _, v := range registerType {
+		// Validate Encoder
+		enc, err := opts.Registry.LookupEncoder(v.ValueType)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if assert.ObjectsAreEqual(v.Encoder, enc) {
+			t.Errorf("Encoder got from Registry: %v, but want: %v", enc, v.Encoder)
+		}
+
+		// Validate Decoder
+		dec, err := opts.Registry.LookupDecoder(v.ValueType)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if assert.ObjectsAreEqual(v.Decoder, dec) {
+			t.Errorf("Decoder got from Registry: %v, but want: %v", dec, v.Decoder)
+		}
+	}
+}