@@ -116,7 +116,7 @@ func (r requestID) augmentErrorWithRequestID(err error) error {
116
116
}
117
117
}
118
118
119
- func gRPCCallOptionsToRequestID (opts []grpc.CallOption ) (reqID requestID , found bool ) {
119
+ func gRPCCallOptionsToRequestID (opts []grpc.CallOption ) (md metadata. MD , reqID requestID , found bool ) {
120
120
for _ , opt := range opts {
121
121
hdrOpt , ok := opt .(grpc.HeaderCallOption )
122
122
if ! ok {
@@ -126,6 +126,7 @@ func gRPCCallOptionsToRequestID(opts []grpc.CallOption) (reqID requestID, found
126
126
metadata := hdrOpt .HeaderAddr
127
127
reqIDs := metadata .Get (xSpannerRequestIDHeader )
128
128
if len (reqIDs ) != 0 && len (reqIDs [0 ]) != 0 {
129
+ md = * metadata
129
130
reqID = requestID (reqIDs [0 ])
130
131
found = true
131
132
break
@@ -137,7 +138,11 @@ func gRPCCallOptionsToRequestID(opts []grpc.CallOption) (reqID requestID, found
137
138
func (wr * requestIDHeaderInjector ) interceptUnary (ctx context.Context , method string , req , reply any , cc * grpc.ClientConn , invoker grpc.UnaryInvoker , opts ... grpc.CallOption ) error {
138
139
// It is imperative to search for the requestID before the call
139
140
// because gRPC's internals will consume the headers.
140
- reqID , foundRequestID := gRPCCallOptionsToRequestID (opts )
141
+ metadataWithRequestID , reqID , foundRequestID := gRPCCallOptionsToRequestID (opts )
142
+ if foundRequestID {
143
+ ctx = metadata .NewOutgoingContext (ctx , metadataWithRequestID )
144
+ }
145
+
141
146
err := invoker (ctx , method , req , reply , cc , opts ... )
142
147
if ! foundRequestID {
143
148
return err
@@ -174,7 +179,11 @@ type requestIDHeaderInjector int
174
179
func (wr * requestIDHeaderInjector ) interceptStream (ctx context.Context , desc * grpc.StreamDesc , cc * grpc.ClientConn , method string , streamer grpc.Streamer , opts ... grpc.CallOption ) (grpc.ClientStream , error ) {
175
180
// It is imperative to search for the requestID before the call
176
181
// because gRPC's internals will consume the headers.
177
- reqID , foundRequestID := gRPCCallOptionsToRequestID (opts )
182
+ metadataWithRequestID , reqID , foundRequestID := gRPCCallOptionsToRequestID (opts )
183
+ if foundRequestID {
184
+ ctx = metadata .NewOutgoingContext (ctx , metadataWithRequestID )
185
+ }
186
+
178
187
cs , err := streamer (ctx , desc , cc , method , opts ... )
179
188
if ! foundRequestID {
180
189
return cs , err
0 commit comments