diff --git a/v4/registry/nats/nats.go b/v4/registry/nats/nats.go index 2b7ee173..6e0d6ebf 100644 --- a/v4/registry/nats/nats.go +++ b/v4/registry/nats/nats.go @@ -14,11 +14,12 @@ import ( ) type natsRegistry struct { - addrs []string - opts registry.Options - nopts nats.Options - queryTopic string - watchTopic string + addrs []string + opts registry.Options + nopts nats.Options + queryTopic string + watchTopic string + registerAction string sync.RWMutex conn *nats.Conn @@ -27,8 +28,9 @@ type natsRegistry struct { } var ( - defaultQueryTopic = "micro.registry.nats.query" - defaultWatchTopic = "micro.registry.nats.watch" + defaultQueryTopic = "micro.registry.nats.query" + defaultWatchTopic = "micro.registry.nats.watch" + defaultRegisterAction = "create" ) func init() { @@ -55,6 +57,11 @@ func configure(n *natsRegistry, opts ...registry.Option) error { watchTopic = wt } + registerAction := defaultRegisterAction + if ra, ok := n.opts.Context.Value(registerActionKey{}).(string); ok { + registerAction = ra + } + // registry.Options have higher priority than nats.Options // only if Addrs, Secure or TLSConfig were not set through a registry.Option // we read them from nats.Option @@ -78,6 +85,7 @@ func configure(n *natsRegistry, opts ...registry.Option) error { n.nopts = natsOptions n.queryTopic = queryTopic n.watchTopic = watchTopic + n.registerAction = registerAction return nil } @@ -322,7 +330,7 @@ func (n *natsRegistry) Register(s *registry.Service, opts ...registry.RegisterOp return err } - b, err := json.Marshal(®istry.Result{Action: "create", Service: s}) + b, err := json.Marshal(®istry.Result{Action: n.registerAction, Service: s}) if err != nil { return err } diff --git a/v4/registry/nats/options.go b/v4/registry/nats/options.go index a3efb678..22db3b94 100644 --- a/v4/registry/nats/options.go +++ b/v4/registry/nats/options.go @@ -11,6 +11,7 @@ type contextQuorumKey struct{} type optionsKey struct{} type watchTopicKey struct{} type queryTopicKey struct{} +type registerActionKey struct{} var ( DefaultQuorum = 0 @@ -70,3 +71,17 @@ func WatchTopic(s string) registry.Option { o.Context = context.WithValue(o.Context, watchTopicKey{}, s) } } + +// RegisterAction allows to set the action to use when registering to nats. +// As of now there are three different options: +// - "create" (default) only registers if there is noone already registered under the same key. +// - "update" only updates the registration if it already exists. +// - "put" creates or updates a registration +func RegisterAction(s string) registry.Option { + return func(o *registry.Options) { + if o.Context == nil { + o.Context = context.Background() + } + o.Context = context.WithValue(o.Context, registerActionKey{}, s) + } +}