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

     1  package meta
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"google.golang.org/grpc/metadata"
     8  
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/credentials"
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/version"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    14  )
    15  
    16  func New(
    17  	database string,
    18  	credentials credentials.Credentials,
    19  	trace *trace.Driver,
    20  	opts ...Option,
    21  ) *Meta {
    22  	m := &Meta{
    23  		trace:       trace,
    24  		credentials: credentials,
    25  		database:    database,
    26  	}
    27  	for _, o := range opts {
    28  		if o != nil {
    29  			o(m)
    30  		}
    31  	}
    32  
    33  	return m
    34  }
    35  
    36  type Option func(m *Meta)
    37  
    38  func WithUserAgentOption(userAgent string) Option {
    39  	return func(m *Meta) {
    40  		m.userAgents = append(m.userAgents, userAgent)
    41  	}
    42  }
    43  
    44  func WithRequestTypeOption(requestType string) Option {
    45  	return func(m *Meta) {
    46  		m.requestsType = requestType
    47  	}
    48  }
    49  
    50  func AllowOption(feature string) Option {
    51  	return func(m *Meta) {
    52  		m.capabilities = append(m.capabilities, feature)
    53  	}
    54  }
    55  
    56  func ForbidOption(feature string) Option {
    57  	return func(m *Meta) {
    58  		n := 0
    59  		for _, capability := range m.capabilities {
    60  			if capability != feature {
    61  				m.capabilities[n] = capability
    62  				n++
    63  			}
    64  		}
    65  		m.capabilities = m.capabilities[:n]
    66  	}
    67  }
    68  
    69  type Meta struct {
    70  	trace        *trace.Driver
    71  	credentials  credentials.Credentials
    72  	database     string
    73  	requestsType string
    74  	userAgents   []string
    75  	capabilities []string
    76  }
    77  
    78  func (m *Meta) meta(ctx context.Context) (_ metadata.MD, err error) {
    79  	md, has := metadata.FromOutgoingContext(ctx)
    80  	if !has {
    81  		md = metadata.MD{}
    82  	}
    83  
    84  	if len(md.Get(HeaderDatabase)) == 0 {
    85  		md.Set(HeaderDatabase, m.database)
    86  	}
    87  
    88  	if len(md.Get(HeaderVersion)) == 0 {
    89  		md.Set(HeaderVersion, version.FullVersion)
    90  	}
    91  
    92  	if m.requestsType != "" {
    93  		if len(md.Get(HeaderRequestType)) == 0 {
    94  			md.Set(HeaderRequestType, m.requestsType)
    95  		}
    96  	}
    97  
    98  	if len(m.userAgents) != 0 {
    99  		md.Append(HeaderUserAgent, m.userAgents...)
   100  	}
   101  
   102  	if len(m.capabilities) > 0 {
   103  		md.Append(HeaderClientCapabilities, m.capabilities...)
   104  	}
   105  
   106  	if m.credentials == nil {
   107  		return md, nil
   108  	}
   109  
   110  	var token string
   111  
   112  	done := trace.DriverOnGetCredentials(m.trace, &ctx, stack.FunctionID(""))
   113  	defer func() {
   114  		done(token, err)
   115  	}()
   116  
   117  	token, err = m.credentials.Token(ctx)
   118  	if err != nil {
   119  		if stringer, ok := m.credentials.(fmt.Stringer); ok {
   120  			return nil, xerrors.WithStackTrace(fmt.Errorf("%w: %s", err, stringer.String()))
   121  		}
   122  
   123  		return nil, xerrors.WithStackTrace(err)
   124  	}
   125  
   126  	md.Set(HeaderTicket, token)
   127  
   128  	return md, nil
   129  }
   130  
   131  func (m *Meta) Context(ctx context.Context) (_ context.Context, err error) {
   132  	md, err := m.meta(ctx)
   133  	if err != nil {
   134  		return ctx, xerrors.WithStackTrace(err)
   135  	}
   136  
   137  	return metadata.NewOutgoingContext(ctx, md), nil
   138  }