Move wrapper internally since its not top level relevant
This commit is contained in:
		
							
								
								
									
										57
									
								
								util/wrapper/wrapper.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										57
									
								
								util/wrapper/wrapper.go
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,57 @@ | ||||
| package wrapper | ||||
|  | ||||
| import ( | ||||
| 	"context" | ||||
|  | ||||
| 	"github.com/micro/go-micro/client" | ||||
| 	"github.com/micro/go-micro/metadata" | ||||
| ) | ||||
|  | ||||
| type clientWrapper struct { | ||||
| 	client.Client | ||||
| 	headers metadata.Metadata | ||||
| } | ||||
|  | ||||
| var ( | ||||
| 	HeaderPrefix = "Micro-" | ||||
| ) | ||||
|  | ||||
| func (c *clientWrapper) setHeaders(ctx context.Context) context.Context { | ||||
| 	// copy metadata | ||||
| 	mda, _ := metadata.FromContext(ctx) | ||||
| 	md := metadata.Copy(mda) | ||||
|  | ||||
| 	// set headers | ||||
| 	for k, v := range c.headers { | ||||
| 		if _, ok := md[k]; !ok { | ||||
| 			md[k] = v | ||||
| 		} | ||||
| 	} | ||||
|  | ||||
| 	return metadata.NewContext(ctx, md) | ||||
| } | ||||
|  | ||||
| func (c *clientWrapper) Call(ctx context.Context, req client.Request, rsp interface{}, opts ...client.CallOption) error { | ||||
| 	ctx = c.setHeaders(ctx) | ||||
| 	return c.Client.Call(ctx, req, rsp, opts...) | ||||
| } | ||||
|  | ||||
| func (c *clientWrapper) Stream(ctx context.Context, req client.Request, opts ...client.CallOption) (client.Stream, error) { | ||||
| 	ctx = c.setHeaders(ctx) | ||||
| 	return c.Client.Stream(ctx, req, opts...) | ||||
| } | ||||
|  | ||||
| func (c *clientWrapper) Publish(ctx context.Context, p client.Message, opts ...client.PublishOption) error { | ||||
| 	ctx = c.setHeaders(ctx) | ||||
| 	return c.Client.Publish(ctx, p, opts...) | ||||
| } | ||||
|  | ||||
| // FromService wraps a client to inject From-Service header into metadata | ||||
| func FromService(name string, c client.Client) client.Client { | ||||
| 	return &clientWrapper{ | ||||
|                 c, | ||||
|                 metadata.Metadata{ | ||||
|                         HeaderPrefix + "From-Service": name, | ||||
|                 }, | ||||
|         } | ||||
| } | ||||
		Reference in New Issue
	
	Block a user