From e7bcf80c95a4e3bcc82f189b041e45786f88a637 Mon Sep 17 00:00:00 2001
From: "junzhou.guo" <junzhou.guo@yunzhanghu.com>
Date: Fri, 11 Nov 2022 18:55:17 +0800
Subject: [PATCH] TransBase setup func

---
 client/dtmgrpc/msg.go     | 10 ++++++++--
 client/dtmgrpc/options.go | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)
 create mode 100644 client/dtmgrpc/options.go

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
+	}
+}