github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/meta/context.go (about)

     1  package meta
     2  
     3  import (
     4  	"context"
     5  
     6  	"google.golang.org/grpc/metadata"
     7  )
     8  
     9  // WithTraceID returns a copy of parent context with traceID
    10  func WithTraceID(ctx context.Context, traceID string) context.Context {
    11  	if md, has := metadata.FromOutgoingContext(ctx); !has || len(md[HeaderTraceID]) == 0 {
    12  		return metadata.AppendToOutgoingContext(ctx, HeaderTraceID, traceID)
    13  	}
    14  
    15  	return ctx
    16  }
    17  
    18  func traceID(ctx context.Context) (string, bool) {
    19  	if md, has := metadata.FromOutgoingContext(ctx); has && len(md[HeaderTraceID]) > 0 {
    20  		return md[HeaderTraceID][0], true
    21  	}
    22  
    23  	return "", false
    24  }
    25  
    26  // WithApplicationName returns a copy of parent context with custom user-agent info
    27  func WithApplicationName(ctx context.Context, applicationName string) context.Context {
    28  	md, has := metadata.FromOutgoingContext(ctx)
    29  	if !has {
    30  		md = metadata.MD{}
    31  	}
    32  	md.Set(HeaderApplicationName, applicationName)
    33  
    34  	return metadata.NewOutgoingContext(ctx, md)
    35  }
    36  
    37  // WithRequestType returns a copy of parent context with custom request type
    38  func WithRequestType(ctx context.Context, requestType string) context.Context {
    39  	return metadata.AppendToOutgoingContext(ctx, HeaderRequestType, requestType)
    40  }
    41  
    42  // WithAllowFeatures returns a copy of parent context with allowed client feature
    43  func WithAllowFeatures(ctx context.Context, features ...string) context.Context {
    44  	kv := make([]string, 0, len(features)*2) //nolint:gomnd
    45  	for _, feature := range features {
    46  		kv = append(kv, HeaderClientCapabilities, feature)
    47  	}
    48  
    49  	return metadata.AppendToOutgoingContext(ctx, kv...)
    50  }
    51  
    52  // WithTraceParent returns a copy of parent context with traceparent header
    53  func WithTraceParent(ctx context.Context, traceparent string) context.Context {
    54  	return metadata.AppendToOutgoingContext(ctx, HeaderTraceParent, traceparent)
    55  }