Repository: zeromicro/go-zero
Pull request #2753 (2023-01-05 04:23:34)

feat: add mongo options

codecov-commenter: # [Codecov](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#2753](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (bcad98e) into [master](https://codecov.io/gh/zeromicro/go-zero/commit/69d355eb4ba722144fdbae21677d8a7ca2c600be?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (69d355e) will **decrease** coverage by `0.01%`.
> The diff coverage is `50.00%`.


| [Impacted Files](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [core/stores/mon/clientmanager.go](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvbW9uL2NsaWVudG1hbmFnZXIuZ28=) | `17.39% <16.66%> (-3.67%)` | :arrow_down: |
| [core/stores/mon/model.go](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvbW9uL21vZGVsLmdv) | `93.61% <100.00%> (ø)` | |
| [core/stores/mon/options.go](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvbW9uL29wdGlvbnMuZ28=) | `100.00% <100.00%> (ø)` | |
| [core/stores/sqlx/sqlconn.go](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvc3FseC9zcWxjb25uLmdv) | `56.37% <0.00%> (-1.48%)` | :arrow_down: |
| [core/stores/sqlx/tx.go](https://codecov.io/gh/zeromicro/go-zero/pull/2753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9zdG9yZXMvc3FseC90eC5nbw==) | `23.30% <0.00%> (+3.88%)` | :arrow_up: |


Patch:
diff --git a/core/stores/mon/clientmanager.go b/core/stores/mon/clientmanager.go
index 5a6c0278..24b8f4cd 100644
--- a/core/stores/mon/clientmanager.go
+++ b/core/stores/mon/clientmanager.go
@@ -3,15 +3,12 @@ package mon
 import (
 	"context"
 	"io"
-	"time"
 
 	"github.com/zeromicro/go-zero/core/syncx"
 	"go.mongodb.org/mongo-driver/mongo"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
-const defaultTimeout = time.Second
-
 var clientManager = syncx.NewResourceManager()
 
 // ClosableClient wraps *mongo.Client and provides a Close method.
@@ -30,9 +27,20 @@ func Inject(key string, client *mongo.Client) {
 	clientManager.Inject(key, &ClosableClient{client})
 }
 
-func getClient(url string) (*mongo.Client, error) {
+func getClient(url string, opts ...Option) (*mongo.Client, error) {
 	val, err := clientManager.GetResource(url, func() (io.Closer, error) {
-		cli, err := mongo.Connect(context.Background(), mopt.Client().ApplyURI(url))
+		o := mopt.Client().ApplyURI(url)
+		opts = append([]Option{defaultTimeoutOption()}, opts...)
+		for _, opt := range opts {
+			opt(o)
+		}
+
+		cli, err := mongo.Connect(context.Background(), o)
+		if err != nil {
+			return nil, err
+		}
+
+		err = cli.Ping(context.Background(), nil)
 		if err != nil {
 			return nil, err
 		}
diff --git a/core/stores/mon/model.go b/core/stores/mon/model.go
index 1f27ff72..d7ec04f2 100644
--- a/core/stores/mon/model.go
+++ b/core/stores/mon/model.go
@@ -48,7 +48,7 @@ func MustNewModel(uri, db, collection string, opts ...Option) *Model {
 
 // NewModel returns a Model.
 func NewModel(uri, db, collection string, opts ...Option) (*Model, error) {
-	cli, err := getClient(uri)
+	cli, err := getClient(uri, opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/core/stores/mon/options.go b/core/stores/mon/options.go
index 31f9ecdb..ca13bff8 100644
--- a/core/stores/mon/options.go
+++ b/core/stores/mon/options.go
@@ -4,14 +4,15 @@ import (
 	"time"
 
 	"github.com/zeromicro/go-zero/core/syncx"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
+const defaultTimeout = time.Second * 3
+
 var slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
 
 type (
-	options struct {
-		timeout time.Duration
-	}
+	options = mopt.ClientOptions
 
 	// Option defines the method to customize a mongo model.
 	Option func(opts *options)
@@ -22,8 +23,15 @@ func SetSlowThreshold(threshold time.Duration) {
 	slowThreshold.Set(threshold)
 }
 
-func defaultOptions() *options {
-	return &options{
-		timeout: defaultTimeout,
+func defaultTimeoutOption() Option {
+	return func(opts *options) {
+		opts.SetTimeout(defaultTimeout)
+	}
+}
+
+// WithTimeout set the mon client operation timeout.
+func WithTimeout(timeout time.Duration) Option {
+	return func(opts *options) {
+		opts.SetTimeout(timeout)
 	}
 }
diff --git a/rest/httpc/requests.go b/rest/httpc/requests.go
index a4ae197d..4ddb15a8 100644
--- a/rest/httpc/requests.go
+++ b/rest/httpc/requests.go
@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"net/http/httptrace"
 	nurl "net/url"
 	"strings"
 
@@ -176,11 +175,6 @@ func request(r *http.Request, cli client) (*http.Response, error) {
 		respHandlers[i] = h
 	}
 
-	clientTrace := httptrace.ContextClientTrace(ctx)
-	if clientTrace != nil {
-		ctx = httptrace.WithClientTrace(ctx, clientTrace)
-	}
-
 	r = r.WithContext(ctx)
 	propagator.Inject(ctx, propagation.HeaderCarrier(r.Header))
 


Test patch:
diff --git a/core/stores/mon/options_test.go b/core/stores/mon/options_test.go
index 4cefad2a..5582f492 100644
--- a/core/stores/mon/options_test.go
+++ b/core/stores/mon/options_test.go
@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
 )
 
 func TestSetSlowThreshold(t *testing.T) {
@@ -13,6 +14,14 @@ func TestSetSlowThreshold(t *testing.T) {
 	assert.Equal(t, time.Second, slowThreshold.Load())
 }
 
-func TestDefaultOptions(t *testing.T) {
-	assert.Equal(t, defaultTimeout, defaultOptions().timeout)
+func Test_defaultTimeoutOption(t *testing.T) {
+	opts := mopt.Client()
+	defaultTimeoutOption()(opts)
+	assert.Equal(t, defaultTimeout, *opts.Timeout)
+}
+
+func TestWithTimeout(t *testing.T) {
+	opts := mopt.Client()
+	WithTimeout(time.Second)(opts)
+	assert.Equal(t, time.Second, *opts.Timeout)
 }
diff --git a/rest/httpc/requests_test.go b/rest/httpc/requests_test.go
index b5bf9612..59fdfed5 100644
--- a/rest/httpc/requests_test.go
+++ b/rest/httpc/requests_test.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/http/httptrace"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -205,12 +206,14 @@ func TestDo_WithClientHttpTrace(t *testing.T) {
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
 	defer svr.Close()
 
+	enter := false
 	_, err := Do(httptrace.WithClientTrace(context.Background(),
 		&httptrace.ClientTrace{
-			DNSStart: func(info httptrace.DNSStartInfo) {
-				assert.Equal(t, "localhost", info.Host)
+			GetConn: func(hostPort string) {
+				assert.Equal(t, "127.0.0.1", strings.Split(hostPort, ":")[0])
+				enter = true
 			},
 		}), http.MethodGet, svr.URL, nil)
 	assert.Nil(t, err)
-
+	assert.True(t, enter)
 }