gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/util/wrapper/wrapper.go (about)

     1  package wrapper
     2  
     3  import (
     4  	"context"
     5  
     6  	"gitee.com/liuxuezhan/go-micro-v1.18.0/client"
     7  	"gitee.com/liuxuezhan/go-micro-v1.18.0/metadata"
     8  )
     9  
    10  type clientWrapper struct {
    11  	client.Client
    12  	headers metadata.Metadata
    13  }
    14  
    15  var (
    16  	HeaderPrefix = "Micro-"
    17  )
    18  
    19  func (c *clientWrapper) setHeaders(ctx context.Context) context.Context {
    20  	// copy metadata
    21  	mda, _ := metadata.FromContext(ctx)
    22  	md := metadata.Copy(mda)
    23  
    24  	// set headers
    25  	for k, v := range c.headers {
    26  		if _, ok := md[k]; !ok {
    27  			md[k] = v
    28  		}
    29  	}
    30  
    31  	return metadata.NewContext(ctx, md)
    32  }
    33  
    34  func (c *clientWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error {
    35  	ctx = c.setHeaders(ctx)
    36  	return c.Client.Call(ctx, req, rsp, opts...)
    37  }
    38  
    39  func (c *clientWrapper) Stream(ctx context.Context, req client.Request, opts ...client.CallOption) (client.Stream, error) {
    40  	ctx = c.setHeaders(ctx)
    41  	return c.Client.Stream(ctx, req, opts...)
    42  }
    43  
    44  func (c *clientWrapper) Publish(ctx context.Context, p client.Message, opts ...client.PublishOption) error {
    45  	ctx = c.setHeaders(ctx)
    46  	return c.Client.Publish(ctx, p, opts...)
    47  }
    48  
    49  // FromService wraps a client to inject From-Service header into metadata
    50  func FromService(name string, c client.Client) client.Client {
    51  	return &clientWrapper{
    52  		c,
    53  		metadata.Metadata{
    54  			HeaderPrefix + "From-Service": name,
    55  		},
    56  	}
    57  }