diff --git a/client/dtmgrpc/msg.go b/client/dtmgrpc/msg.go index e08d2f3..e9b59da 100644 --- a/client/dtmgrpc/msg.go +++ b/client/dtmgrpc/msg.go @@ -24,8 +24,14 @@ type MsgGrpc struct { } // NewMsgGrpc create new msg -func NewMsgGrpc(server string, gid string) *MsgGrpc { - return &MsgGrpc{Msg: *dtmcli.NewMsg(server, gid)} +func NewMsgGrpc(server string, gid string, opts ...TransBaseOption) *MsgGrpc { + mg := &MsgGrpc{Msg: *dtmcli.NewMsg(server, gid)} + + for _, opt := range opts { + opt(&mg.TransBase) + } + + return mg } // Add add a new step diff --git a/client/dtmgrpc/options.go b/client/dtmgrpc/options.go new file mode 100644 index 0000000..1ea61c7 --- /dev/null +++ b/client/dtmgrpc/options.go @@ -0,0 +1,15 @@ +package dtmgrpc + +import ( + "github.com/dtm-labs/dtm/client/dtmcli/dtmimp" +) + +// TransBaseOption setup func for TransBase +type TransBaseOption func(tb *dtmimp.TransBase) + +// WithBranchHeaders setup TransBase.BranchHeaders +func WithBranchHeaders(headers map[string]string) TransBaseOption { + return func(tb *dtmimp.TransBase) { + tb.BranchHeaders = headers + } +}