diff --git a/core/stores/mongo/collection.go b/core/stores/mongo/collection.go index 6fb199b2..1592dfa2 100644 --- a/core/stores/mongo/collection.go +++ b/core/stores/mongo/collection.go @@ -43,11 +43,11 @@ type ( } ) -func newCollection(collection *mgo.Collection) Collection { +func newCollection(collection *mgo.Collection, brk breaker.Breaker) Collection { return &decoratedCollection{ name: collection.FullName, collection: collection, - brk: breaker.NewBreaker(), + brk: brk, } } diff --git a/core/stores/mongo/collection_test.go b/core/stores/mongo/collection_test.go index 2b6c7825..4723bf58 100644 --- a/core/stores/mongo/collection_test.go +++ b/core/stores/mongo/collection_test.go @@ -71,7 +71,7 @@ func TestNewCollection(t *testing.T) { Database: nil, Name: "foo", FullName: "bar", - }) + }, breaker.GetBreaker("localhost")) assert.Equal(t, "bar", col.(*decoratedCollection).name) } diff --git a/core/stores/mongo/model.go b/core/stores/mongo/model.go index f8c298ac..f42186d0 100644 --- a/core/stores/mongo/model.go +++ b/core/stores/mongo/model.go @@ -5,6 +5,7 @@ import ( "time" "github.com/globalsign/mgo" + "github.com/tal-tech/go-zero/core/breaker" ) type ( @@ -20,6 +21,7 @@ type ( session *concurrentSession db *mgo.Database collection string + brk breaker.Breaker opts []Option } ) @@ -46,6 +48,7 @@ func NewModel(url, collection string, opts ...Option) (*Model, error) { // If name is empty, the database name provided in the dialed URL is used instead db: session.DB(""), collection: collection, + brk: breaker.GetBreaker(url), opts: opts, }, nil } @@ -66,7 +69,7 @@ func (mm *Model) FindId(id interface{}) (Query, error) { // GetCollection returns a Collection with given session. func (mm *Model) GetCollection(session *mgo.Session) Collection { - return newCollection(mm.db.C(mm.collection).With(session)) + return newCollection(mm.db.C(mm.collection).With(session), mm.brk) } // Insert inserts docs into mm.