github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/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  // WithUserAgent returns a copy of parent context with custom user-agent info
    27  func WithUserAgent(ctx context.Context, userAgent string) context.Context {
    28  	return metadata.AppendToOutgoingContext(ctx, HeaderUserAgent, userAgent)
    29  }
    30  
    31  // WithRequestType returns a copy of parent context with custom request type
    32  func WithRequestType(ctx context.Context, requestType string) context.Context {
    33  	return metadata.AppendToOutgoingContext(ctx, HeaderRequestType, requestType)
    34  }
    35  
    36  // WithAllowFeatures returns a copy of parent context with allowed client feature
    37  func WithAllowFeatures(ctx context.Context, features []string) context.Context {
    38  	kv := make([]string, 0, len(features)*2)
    39  	for _, feature := range features {
    40  		kv = append(kv, HeaderClientCapabilities, feature)
    41  	}
    42  
    43  	return metadata.AppendToOutgoingContext(ctx, kv...)
    44  }