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

     1  package query
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
     7  
     8  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/session"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
    12  	baseTx "github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/query"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    16  )
    17  
    18  var _ query.Session = (*Session)(nil)
    19  
    20  type (
    21  	Session struct {
    22  		session.Core
    23  
    24  		client Ydb_Query_V1.QueryServiceClient
    25  		trace  *trace.Query
    26  		laztTx bool
    27  	}
    28  )
    29  
    30  func (s *Session) QueryResultSet(
    31  	ctx context.Context, q string, opts ...options.Execute,
    32  ) (rs result.ClosableResultSet, finalErr error) {
    33  	onDone := trace.QueryOnSessionQueryResultSet(s.trace, &ctx,
    34  		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).QueryResultSet"), s, q)
    35  	defer func() {
    36  		onDone(finalErr)
    37  	}()
    38  
    39  	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
    40  	if err != nil {
    41  		return nil, xerrors.WithStackTrace(err)
    42  	}
    43  
    44  	rs, err = readResultSet(ctx, r)
    45  	if err != nil {
    46  		return nil, xerrors.WithStackTrace(err)
    47  	}
    48  
    49  	return rs, nil
    50  }
    51  
    52  func (s *Session) queryRow(
    53  	ctx context.Context, q string, settings executeSettings, resultOpts ...resultOption,
    54  ) (row query.Row, finalErr error) {
    55  	r, err := execute(ctx, s.ID(), s.client, q, settings, resultOpts...)
    56  	if err != nil {
    57  		return nil, xerrors.WithStackTrace(err)
    58  	}
    59  
    60  	row, err = readRow(ctx, r)
    61  	if err != nil {
    62  		return nil, xerrors.WithStackTrace(err)
    63  	}
    64  
    65  	return row, nil
    66  }
    67  
    68  func (s *Session) QueryRow(ctx context.Context, q string, opts ...options.Execute) (_ query.Row, finalErr error) {
    69  	onDone := trace.QueryOnSessionQueryRow(s.trace, &ctx,
    70  		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).QueryRow"), s, q)
    71  	defer func() {
    72  		onDone(finalErr)
    73  	}()
    74  
    75  	row, err := s.queryRow(ctx, q, options.ExecuteSettings(opts...), withTrace(s.trace))
    76  	if err != nil {
    77  		return nil, xerrors.WithStackTrace(err)
    78  	}
    79  
    80  	return row, nil
    81  }
    82  
    83  func createSession(
    84  	ctx context.Context, client Ydb_Query_V1.QueryServiceClient, opts ...session.Option,
    85  ) (*Session, error) {
    86  	core, err := session.Open(ctx, client, opts...)
    87  	if err != nil {
    88  		return nil, xerrors.WithStackTrace(err)
    89  	}
    90  
    91  	return &Session{
    92  		Core:   core,
    93  		trace:  core.Trace,
    94  		client: core.Client,
    95  	}, nil
    96  }
    97  
    98  func (s *Session) Begin(
    99  	ctx context.Context,
   100  	txSettings query.TransactionSettings,
   101  ) (
   102  	tx query.Transaction, finalErr error,
   103  ) {
   104  	onDone := trace.QueryOnSessionBegin(s.trace, &ctx,
   105  		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Begin"), s)
   106  	defer func() {
   107  		if finalErr != nil {
   108  			onDone(finalErr, nil)
   109  		} else {
   110  			onDone(nil, tx)
   111  		}
   112  	}()
   113  
   114  	if s.laztTx {
   115  		return &Transaction{
   116  			s:          s,
   117  			txSettings: txSettings,
   118  		}, nil
   119  	}
   120  
   121  	txID, err := begin(ctx, s.client, s.ID(), txSettings)
   122  	if err != nil {
   123  		return nil, xerrors.WithStackTrace(err)
   124  	}
   125  
   126  	return &Transaction{
   127  		LazyID: baseTx.ID(txID),
   128  		s:      s,
   129  	}, nil
   130  }
   131  
   132  func (s *Session) Exec(
   133  	ctx context.Context, q string, opts ...options.Execute,
   134  ) (finalErr error) {
   135  	onDone := trace.QueryOnSessionExec(s.trace, &ctx,
   136  		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Exec"), s, q)
   137  	defer func() {
   138  		onDone(finalErr)
   139  	}()
   140  
   141  	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
   142  	if err != nil {
   143  		return xerrors.WithStackTrace(err)
   144  	}
   145  
   146  	err = readAll(ctx, r)
   147  	if err != nil {
   148  		return xerrors.WithStackTrace(err)
   149  	}
   150  
   151  	return nil
   152  }
   153  
   154  func (s *Session) Query(
   155  	ctx context.Context, q string, opts ...options.Execute,
   156  ) (_ query.Result, finalErr error) {
   157  	onDone := trace.QueryOnSessionQuery(s.trace, &ctx,
   158  		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Query"), s, q)
   159  	defer func() {
   160  		onDone(finalErr)
   161  	}()
   162  
   163  	r, err := execute(ctx, s.ID(), s.client, q, options.ExecuteSettings(opts...), withTrace(s.trace))
   164  	if err != nil {
   165  		return nil, xerrors.WithStackTrace(err)
   166  	}
   167  
   168  	return r, nil
   169  }