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

     1  package xsql
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/ydb-platform/ydb-go-sdk/v3/table"
     7  	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
     8  )
     9  
    10  type (
    11  	ctxTransactionControlKey struct{}
    12  	ctxDataQueryOptionsKey   struct{}
    13  	ctxScanQueryOptionsKey   struct{}
    14  	ctxModeTypeKey           struct{}
    15  	ctxTxControlHookKey      struct{}
    16  
    17  	txControlHook func(txControl *table.TransactionControl)
    18  )
    19  
    20  func WithTxControlHook(ctx context.Context, hook txControlHook) context.Context {
    21  	return context.WithValue(ctx, ctxTxControlHookKey{}, hook)
    22  }
    23  
    24  // WithQueryMode returns a copy of context with given QueryMode
    25  func WithQueryMode(ctx context.Context, m QueryMode) context.Context {
    26  	return context.WithValue(ctx, ctxModeTypeKey{}, m)
    27  }
    28  
    29  // queryModeFromContext returns defined QueryMode or DefaultQueryMode
    30  func queryModeFromContext(ctx context.Context, defaultQueryMode QueryMode) QueryMode {
    31  	if m, ok := ctx.Value(ctxModeTypeKey{}).(QueryMode); ok {
    32  		return m
    33  	}
    34  
    35  	return defaultQueryMode
    36  }
    37  
    38  func WithTxControl(ctx context.Context, txc *table.TransactionControl) context.Context {
    39  	return context.WithValue(ctx, ctxTransactionControlKey{}, txc)
    40  }
    41  
    42  func txControl(ctx context.Context, defaultTxControl *table.TransactionControl) (txControl *table.TransactionControl) {
    43  	defer func() {
    44  		if hook, has := ctx.Value(ctxTxControlHookKey{}).(txControlHook); has && hook != nil {
    45  			hook(txControl)
    46  		}
    47  	}()
    48  	if txc, ok := ctx.Value(ctxTransactionControlKey{}).(*table.TransactionControl); ok {
    49  		return txc
    50  	}
    51  
    52  	return defaultTxControl
    53  }
    54  
    55  func (c *conn) WithScanQueryOptions(ctx context.Context, opts ...options.ExecuteScanQueryOption) context.Context {
    56  	return context.WithValue(ctx,
    57  		ctxScanQueryOptionsKey{},
    58  		append(
    59  			append([]options.ExecuteScanQueryOption{}, c.scanQueryOptions(ctx)...),
    60  			opts...,
    61  		),
    62  	)
    63  }
    64  
    65  func (c *conn) scanQueryOptions(ctx context.Context) []options.ExecuteScanQueryOption {
    66  	if opts, ok := ctx.Value(ctxScanQueryOptionsKey{}).([]options.ExecuteScanQueryOption); ok {
    67  		return append(c.scanOpts, opts...)
    68  	}
    69  
    70  	return c.scanOpts
    71  }
    72  
    73  func (c *conn) WithDataQueryOptions(ctx context.Context, opts ...options.ExecuteDataQueryOption) context.Context {
    74  	return context.WithValue(ctx,
    75  		ctxDataQueryOptionsKey{},
    76  		append(
    77  			append([]options.ExecuteDataQueryOption{}, c.dataQueryOptions(ctx)...),
    78  			opts...,
    79  		),
    80  	)
    81  }
    82  
    83  func (c *conn) dataQueryOptions(ctx context.Context) []options.ExecuteDataQueryOption {
    84  	if opts, ok := ctx.Value(ctxDataQueryOptionsKey{}).([]options.ExecuteDataQueryOption); ok {
    85  		return append(c.dataOpts, opts...)
    86  	}
    87  
    88  	return c.dataOpts
    89  }
    90  
    91  func (c *conn) withKeepInCache(ctx context.Context) context.Context {
    92  	return c.WithDataQueryOptions(ctx, options.WithKeepInCache(true))
    93  }