go.uber.org/yarpc@v1.72.1/internal/grpcctx/grpcctx.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  // Package grpcctx contains helper functionality for testing with grpc-go.
    22  package grpcctx
    23  
    24  import (
    25  	"context"
    26  
    27  	"google.golang.org/grpc/metadata"
    28  )
    29  
    30  // ContextWrapper wraps a context for grpc-go with the required headers for yarpc.
    31  //
    32  // This is a convenience object for use when using grpc-go clients. You must set
    33  // certain yarpc-specific headers when using native grpc-go clients calling into yarpc
    34  // servers, and this object makes that simpler.
    35  type ContextWrapper struct {
    36  	md metadata.MD
    37  }
    38  
    39  // NewContextWrapper returns a new ContextWrapper.
    40  //
    41  // The only fields that a grpc-go client needs to set are caller and service.
    42  // Encoding is also required if content-type is not set properly.
    43  // See the documention on EncodingHeader.
    44  func NewContextWrapper() *ContextWrapper {
    45  	return &ContextWrapper{metadata.New(nil)}
    46  }
    47  
    48  // Wrap wraps the given context with the headers.
    49  func (c *ContextWrapper) Wrap(ctx context.Context) context.Context {
    50  	return metadata.NewOutgoingContext(ctx, c.md)
    51  }
    52  
    53  // WithCaller returns a new ContextWrapper with the given caller.
    54  func (c *ContextWrapper) WithCaller(caller string) *ContextWrapper {
    55  	return c.copyAndAdd("rpc-caller", caller)
    56  }
    57  
    58  // WithService returns a new ContextWrapper with the given service.
    59  func (c *ContextWrapper) WithService(service string) *ContextWrapper {
    60  	return c.copyAndAdd("rpc-service", service)
    61  }
    62  
    63  // WithShardKey returns a new ContextWrapper with the given shard key.
    64  func (c *ContextWrapper) WithShardKey(shardKey string) *ContextWrapper {
    65  	return c.copyAndAdd("rpc-shard-key", shardKey)
    66  }
    67  
    68  // WithRoutingKey returns a new ContextWrapper with the given routing key.
    69  func (c *ContextWrapper) WithRoutingKey(routingKey string) *ContextWrapper {
    70  	return c.copyAndAdd("rpc-routing-key", routingKey)
    71  }
    72  
    73  // WithRoutingDelegate returns a new ContextWrapper with the given routing delegate.
    74  func (c *ContextWrapper) WithRoutingDelegate(routingDelegate string) *ContextWrapper {
    75  	return c.copyAndAdd("rpc-routing-delegate", routingDelegate)
    76  }
    77  
    78  // WithEncoding returns a new ContextWrapper with the given encoding.
    79  func (c *ContextWrapper) WithEncoding(encoding string) *ContextWrapper {
    80  	return c.copyAndAdd("rpc-encoding", encoding)
    81  }
    82  
    83  // WithHeader returns a new ContextWrapper with the given header.
    84  func (c *ContextWrapper) WithHeader(key, value string) *ContextWrapper {
    85  	return c.copyAndAdd(key, value)
    86  }
    87  
    88  func (c *ContextWrapper) copyAndAdd(key, value string) *ContextWrapper {
    89  	md := c.md
    90  	if md == nil {
    91  		md = metadata.New(nil)
    92  	} else {
    93  		md = md.Copy()
    94  	}
    95  	md[key] = []string{value}
    96  	return &ContextWrapper{md}
    97  }