Skip to content

Commit

Permalink
fix(spanner): inject "x-goog-spanner-request-id" into outgoing client…
Browse files Browse the repository at this point in the history
… context

Injects metadata derived from the associated gax.CallOptions that were
correctly being generated and tested, but accidentally not injected
into the outgoing context.

Fixes #11543
  • Loading branch information
odeke-em committed Jan 31, 2025
1 parent 0dd7d3d commit bfa4afc
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions spanner/request_id_header.go
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ func (r requestID) augmentErrorWithRequestID(err error) error {
}
}

func gRPCCallOptionsToRequestID(opts []grpc.CallOption) (reqID requestID, found bool) {
func gRPCCallOptionsToRequestID(opts []grpc.CallOption) (md metadata.MD, reqID requestID, found bool) {
for _, opt := range opts {
hdrOpt, ok := opt.(grpc.HeaderCallOption)
if !ok {
Expand All @@ -126,6 +126,7 @@ func gRPCCallOptionsToRequestID(opts []grpc.CallOption) (reqID requestID, found
metadata := hdrOpt.HeaderAddr
reqIDs := metadata.Get(xSpannerRequestIDHeader)
if len(reqIDs) != 0 && len(reqIDs[0]) != 0 {
md = *metadata
reqID = requestID(reqIDs[0])
found = true
break
Expand All @@ -137,7 +138,11 @@ func gRPCCallOptionsToRequestID(opts []grpc.CallOption) (reqID requestID, found
func (wr *requestIDHeaderInjector) interceptUnary(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
// It is imperative to search for the requestID before the call
// because gRPC's internals will consume the headers.
reqID, foundRequestID := gRPCCallOptionsToRequestID(opts)
metadataWithRequestID, reqID, foundRequestID := gRPCCallOptionsToRequestID(opts)
if foundRequestID {
ctx = metadata.NewOutgoingContext(ctx, metadataWithRequestID)
}

err := invoker(ctx, method, req, reply, cc, opts...)
if !foundRequestID {
return err
Expand Down Expand Up @@ -174,7 +179,11 @@ type requestIDHeaderInjector int
func (wr *requestIDHeaderInjector) interceptStream(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
// It is imperative to search for the requestID before the call
// because gRPC's internals will consume the headers.
reqID, foundRequestID := gRPCCallOptionsToRequestID(opts)
metadataWithRequestID, reqID, foundRequestID := gRPCCallOptionsToRequestID(opts)
if foundRequestID {
ctx = metadata.NewOutgoingContext(ctx, metadataWithRequestID)
}

cs, err := streamer(ctx, desc, cc, method, opts...)
if !foundRequestID {
return cs, err
Expand Down

0 comments on commit bfa4afc

Please sign in to comment.