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

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