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

     1  package query
     2  
     3  import (
     4  	"context"
     5  	"sync/atomic"
     6  
     7  	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
     8  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
     9  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
    10  
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/query"
    14  )
    15  
    16  var _ query.Session = (*Session)(nil)
    17  
    18  type Session struct {
    19  	id          string
    20  	nodeID      int64
    21  	queryClient Ydb_Query_V1.QueryServiceClient
    22  	status      query.SessionStatus
    23  	close       func()
    24  }
    25  
    26  func (s *Session) Close(ctx context.Context) error {
    27  	s.close()
    28  
    29  	return nil
    30  }
    31  
    32  func begin(
    33  	ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessionID string, txSettings query.TransactionSettings,
    34  ) (*transaction, error) {
    35  	a := allocator.New()
    36  	defer a.Free()
    37  	response, err := client.BeginTransaction(ctx,
    38  		&Ydb_Query.BeginTransactionRequest{
    39  			SessionId:  sessionID,
    40  			TxSettings: txSettings.ToYDB(a),
    41  		},
    42  	)
    43  	if err != nil {
    44  		return nil, xerrors.WithStackTrace(xerrors.Transport(err))
    45  	}
    46  	if response.GetStatus() != Ydb.StatusIds_SUCCESS {
    47  		return nil, xerrors.WithStackTrace(xerrors.FromOperation(response))
    48  	}
    49  
    50  	return &transaction{
    51  		id: response.GetTxMeta().GetId(),
    52  	}, nil
    53  }
    54  
    55  func (s *Session) Begin(ctx context.Context, txSettings query.TransactionSettings) (query.Transaction, error) {
    56  	tx, err := begin(ctx, s.queryClient, s.id, txSettings)
    57  	if err != nil {
    58  		return nil, xerrors.WithStackTrace(err)
    59  	}
    60  	tx.s = s
    61  
    62  	return tx, nil
    63  }
    64  
    65  func (s *Session) ID() string {
    66  	return s.id
    67  }
    68  
    69  func (s *Session) NodeID() int64 {
    70  	return s.nodeID
    71  }
    72  
    73  func (s *Session) Status() query.SessionStatus {
    74  	return query.SessionStatus(atomic.LoadUint32((*uint32)(&s.status)))
    75  }
    76  
    77  func (s *Session) Execute(
    78  	ctx context.Context, q string, opts ...query.ExecuteOption,
    79  ) (query.Transaction, query.Result, error) {
    80  	tx, r, err := execute(ctx, s, s.queryClient, q, query.ExecuteSettings(opts...))
    81  	if err != nil {
    82  		return nil, nil, xerrors.WithStackTrace(err)
    83  	}
    84  
    85  	return tx, r, nil
    86  }