github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/query/client.go

package query

import (
	"context"
	"time"

	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Operations"
	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
	"google.golang.org/grpc"

	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/session"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/types"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
	"github.com/ydb-platform/ydb-go-sdk/v3/query"
	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
)

//go:generate mockgen -destination grpc_client_mock_test.go --typed -package query -write_package_comment=false github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1 QueryServiceClient,QueryService_AttachSessionClient,QueryService_ExecuteQueryClient

var (
	_ query.Client = (*Client)(nil)
	_ sessionPool  = (*pool.Pool[*Session, Session])(nil)
)

type (
	sessionPool interface {
		closer.Closer

		Stats() pool.Stats
		With(ctx context.Context, f func(ctx context.Context, s *Session) error, opts ...retry.Option) error
	}
	Client struct {
		config *config.Config
		client Ydb_Query_V1.QueryServiceClient
		pool   sessionPool

		done chan struct{}
	}
)

func fetchScriptResults(ctx context.Context,
	client Ydb_Query_V1.QueryServiceClient,
	opID string, opts ...options.FetchScriptOption,
) (*options.FetchScriptResult, error) {
	r, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
		request := &options.FetchScriptResultsRequest{
			FetchScriptResultsRequest: Ydb_Query.FetchScriptResultsRequest{
				OperationId: opID,
			},
		}
		for _, opt := range opts {
			if opt != nil {
				opt(request)
			}
		}

		response, err := client.FetchScriptResults(ctx, &request.FetchScriptResultsRequest)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		rs := response.GetResultSet()
		columns := rs.GetColumns()
		columnNames := make([]string, len(columns))
		columnTypes := make([]types.Type, len(columns))
		for i := range columns {
			columnNames[i] = columns[i].GetName()
			columnTypes[i] = types.TypeFromYDB(columns[i].GetType())
		}
		rows := make([]query.Row, len(rs.GetRows()))
		for i, r := range rs.GetRows() {
			rows[i] = NewRow(columns, r)
		}

		return &options.FetchScriptResult{
			ResultSetIndex: response.GetResultSetIndex(),
			ResultSet:      MaterializedResultSet(int(response.GetResultSetIndex()), columnNames, columnTypes, rows),
			NextToken:      response.GetNextFetchToken(),
		}, nil
	}, retry.WithIdempotent(true))
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

func (c *Client) FetchScriptResults(ctx context.Context,
	opID string, opts ...options.FetchScriptOption,
) (*options.FetchScriptResult, error) {
	r, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
		r, err := fetchScriptResults(ctx, c.client, opID,
			append(opts, func(request *options.FetchScriptResultsRequest) {
				request.Trace = c.config.Trace()
			})...,
		)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		return r, nil
	}, retry.WithIdempotent(true))
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

type executeScriptSettings struct {
	executeSettings
	ttl             time.Duration
	operationParams *Ydb_Operations.OperationParams
}

func (s *executeScriptSettings) OperationParams() *Ydb_Operations.OperationParams {
	return s.operationParams
}

func (s *executeScriptSettings) ResultsTTL() time.Duration {
	return s.ttl
}

func executeScript(ctx context.Context,
	client Ydb_Query_V1.QueryServiceClient, request *Ydb_Query.ExecuteScriptRequest, grpcOpts ...grpc.CallOption,
) (*options.ExecuteScriptOperation, error) {
	op, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.ExecuteScriptOperation, error) {
		response, err := client.ExecuteScript(ctx, request, grpcOpts...)
		if err != nil {
			return nil, xerrors.WithStackTrace(err)
		}

		return &options.ExecuteScriptOperation{
			ID:            response.GetId(),
			ConsumedUnits: response.GetCostInfo().GetConsumedUnits(),
			Metadata:      options.ToMetadataExecuteQuery(response.GetMetadata()),
		}, nil
	}, retry.WithIdempotent(true))
	if err != nil {
		return op, xerrors.WithStackTrace(err)
	}

	return op, nil
}

func (c *Client) ExecuteScript(
	ctx context.Context, q string, ttl time.Duration, opts ...options.Execute,
) (
	op *options.ExecuteScriptOperation, err error,
) {
	a := allocator.New()
	defer a.Free()

	settings := &executeScriptSettings{
		executeSettings: options.ExecuteSettings(opts...),
		ttl:             ttl,
		operationParams: operation.Params(
			ctx,
			c.config.OperationTimeout(),
			c.config.OperationCancelAfter(),
			operation.ModeSync,
		),
	}

	request, grpcOpts := executeQueryScriptRequest(a, q, settings)

	op, err = executeScript(ctx, c.client, request, grpcOpts...)
	if err != nil {
		return op, xerrors.WithStackTrace(err)
	}

	return op, nil
}

func (c *Client) Close(ctx context.Context) error {
	close(c.done)

	if err := c.pool.Close(ctx); err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func do(
	ctx context.Context,
	pool sessionPool,
	op func(ctx context.Context, s *Session) error,
	opts ...retry.Option,
) (finalErr error) {
	err := pool.With(ctx, func(ctx context.Context, s *Session) error {
		s.SetStatus(session.StatusInUse)

		err := op(ctx, s)
		if err != nil {
			s.SetStatus(session.StatusError)

			return xerrors.WithStackTrace(err)
		}

		s.SetStatus(session.StatusIdle)

		return nil
	}, opts...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func (c *Client) Do(ctx context.Context, op query.Operation, opts ...options.DoOption) (finalErr error) {
	ctx, cancel := xcontext.WithDone(ctx, c.done)
	defer cancel()

	var (
		settings = options.ParseDoOpts(c.config.Trace(), opts...)
		onDone   = trace.QueryOnDo(settings.Trace(), &ctx,
			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Do"),
		)
		attempts = 0
	)
	defer func() {
		onDone(attempts, finalErr)
	}()

	err := do(ctx, c.pool,
		func(ctx context.Context, s *Session) error {
			return op(ctx, s)
		},
		append([]retry.Option{
			retry.WithTrace(&trace.Retry{
				OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
					return func(info trace.RetryLoopDoneInfo) {
						attempts = info.Attempts
					}
				},
			}),
		}, settings.RetryOpts()...)...,
	)

	return err
}

func doTx(
	ctx context.Context,
	pool sessionPool,
	op query.TxOperation,
	txSettings tx.Settings,
	opts ...retry.Option,
) (finalErr error) {
	err := do(ctx, pool, func(ctx context.Context, s *Session) (opErr error) {
		tx, err := s.Begin(ctx, txSettings)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		defer func() {
			_ = tx.Rollback(ctx)

			if opErr != nil {
				s.SetStatus(session.StatusError)
			}
		}()

		err = op(ctx, tx)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		err = tx.CommitTx(ctx)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, opts...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func clientQueryRow(
	ctx context.Context, pool sessionPool, q string, settings executeSettings, resultOpts ...resultOption,
) (row query.Row, finalErr error) {
	err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
		row, err = s.queryRow(ctx, q, settings, resultOpts...)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return row, nil
}

// QueryRow is a helper which read only one row from first result set in result
func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute) (_ query.Row, finalErr error) {
	ctx, cancel := xcontext.WithDone(ctx, c.done)
	defer cancel()

	onDone := trace.QueryOnQueryRow(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).QueryRow"),
		q,
	)
	defer func() {
		onDone(finalErr)
	}()

	row, err := clientQueryRow(ctx, c.pool, q, options.ExecuteSettings(opts...), withTrace(c.config.Trace()))
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return row, nil
}

func clientExec(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (finalErr error) {
	settings := options.ExecuteSettings(opts...)
	err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
		streamResult, err := execute(ctx, s.ID(), s.client, q, settings, withTrace(s.trace))
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		err = readAll(ctx, streamResult)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func (c *Client) Exec(ctx context.Context, q string, opts ...options.Execute) (finalErr error) {
	ctx, cancel := xcontext.WithDone(ctx, c.done)
	defer cancel()

	onDone := trace.QueryOnExec(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Exec"),
		q,
	)
	defer func() {
		onDone(finalErr)
	}()

	err := clientExec(ctx, c.pool, q, opts...)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (
	r query.Result, err error,
) {
	settings := options.ExecuteSettings(opts...)
	err = do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
		streamResult, err := execute(ctx, s.ID(), s.client, q,
			options.ExecuteSettings(opts...), withTrace(s.trace),
		)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		if err != nil {
			return xerrors.WithStackTrace(err)
		}
		defer func() {
			_ = streamResult.Close(ctx)
		}()

		r, err = resultToMaterializedResult(ctx, streamResult)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

func (c *Client) Query(ctx context.Context, q string, opts ...options.Execute) (r query.Result, err error) {
	ctx, cancel := xcontext.WithDone(ctx, c.done)
	defer cancel()

	onDone := trace.QueryOnQuery(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Query"),
		q,
	)
	defer func() {
		onDone(err)
	}()

	r, err = clientQuery(ctx, c.pool, q, opts...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return r, nil
}

func clientQueryResultSet(
	ctx context.Context, pool sessionPool, q string, settings executeSettings, resultOpts ...resultOption,
) (rs result.ClosableResultSet, finalErr error) {
	err := do(ctx, pool, func(ctx context.Context, s *Session) error {
		streamResult, err := execute(ctx, s.ID(), s.client, q, settings, resultOpts...)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		rs, err = readMaterializedResultSet(ctx, streamResult)
		if err != nil {
			return xerrors.WithStackTrace(err)
		}

		return nil
	}, settings.RetryOpts()...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return rs, nil
}

// QueryResultSet is a helper which read all rows from first result set in result
func (c *Client) QueryResultSet(
	ctx context.Context, q string, opts ...options.Execute,
) (rs result.ClosableResultSet, finalErr error) {
	ctx, cancel := xcontext.WithDone(ctx, c.done)
	defer cancel()

	onDone := trace.QueryOnQueryResultSet(c.config.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).QueryResultSet"),
		q,
	)
	defer func() {
		onDone(finalErr)
	}()

	rs, err := clientQueryResultSet(ctx, c.pool, q, options.ExecuteSettings(opts...), withTrace(c.config.Trace()))
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return rs, nil
}

func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options.DoTxOption) (finalErr error) {
	ctx, cancel := xcontext.WithDone(ctx, c.done)
	defer cancel()

	var (
		settings = options.ParseDoTxOpts(c.config.Trace(), opts...)
		onDone   = trace.QueryOnDoTx(settings.Trace(), &ctx,
			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).DoTx"),
		)
		attempts = 0
	)
	defer func() {
		onDone(attempts, finalErr)
	}()

	err := doTx(ctx, c.pool, op,
		settings.TxSettings(),
		append(
			[]retry.Option{
				retry.WithTrace(&trace.Retry{
					OnRetry: func(info trace.RetryLoopStartInfo) func(trace.RetryLoopDoneInfo) {
						return func(info trace.RetryLoopDoneInfo) {
							attempts = info.Attempts
						}
					},
				}),
			},
			settings.RetryOpts()...,
		)...,
	)
	if err != nil {
		return xerrors.WithStackTrace(err)
	}

	return nil
}

func New(ctx context.Context, cc grpc.ClientConnInterface, cfg *config.Config) *Client {
	onDone := trace.QueryOnNew(cfg.Trace(), &ctx,
		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.New"),
	)
	defer onDone()

	client := Ydb_Query_V1.NewQueryServiceClient(cc)

	return &Client{
		config: cfg,
		client: client,
		done:   make(chan struct{}),
		pool: pool.New(ctx,
			pool.WithLimit[*Session, Session](cfg.PoolLimit()),
			pool.WithItemUsageLimit[*Session, Session](cfg.PoolSessionUsageLimit()),
			pool.WithTrace[*Session, Session](poolTrace(cfg.Trace())),
			pool.WithCreateItemTimeout[*Session, Session](cfg.SessionCreateTimeout()),
			pool.WithCloseItemTimeout[*Session, Session](cfg.SessionDeleteTimeout()),
			pool.WithIdleTimeToLive[*Session, Session](cfg.SessionIdleTimeToLive()),
			pool.WithCreateItemFunc(func(ctx context.Context) (_ *Session, err error) {
				var (
					createCtx    context.Context
					cancelCreate context.CancelFunc
				)
				if d := cfg.SessionCreateTimeout(); d > 0 {
					createCtx, cancelCreate = xcontext.WithTimeout(ctx, d)
				} else {
					createCtx, cancelCreate = xcontext.WithCancel(ctx)
				}
				defer cancelCreate()

				s, err := createSession(createCtx, client,
					session.WithConn(cc),
					session.WithDeleteTimeout(cfg.SessionDeleteTimeout()),
					session.WithTrace(cfg.Trace()),
				)
				if err != nil {
					return nil, xerrors.WithStackTrace(err)
				}

				s.laztTx = cfg.LazyTx()

				return s, nil
			}),
		),
	}
}

func poolTrace(t *trace.Query) *pool.Trace {
	return &pool.Trace{
		OnNew: func(ctx *context.Context, call stack.Caller) func(limit int) {
			onDone := trace.QueryOnPoolNew(t, ctx, call)

			return func(limit int) {
				onDone(limit)
			}
		},
		OnClose: func(ctx *context.Context, call stack.Caller) func(err error) {
			onDone := trace.QueryOnClose(t, ctx, call)

			return func(err error) {
				onDone(err)
			}
		},
		OnTry: func(ctx *context.Context, call stack.Caller) func(err error) {
			onDone := trace.QueryOnPoolTry(t, ctx, call)

			return func(err error) {
				onDone(err)
			}
		},
		OnWith: func(ctx *context.Context, call stack.Caller) func(attempts int, err error) {
			onDone := trace.QueryOnPoolWith(t, ctx, call)

			return func(attempts int, err error) {
				onDone(attempts, err)
			}
		},
		OnPut: func(ctx *context.Context, call stack.Caller, item any) func(err error) {
			onDone := trace.QueryOnPoolPut(t, ctx, call, item.(*Session)) //nolint:forcetypeassert

			return func(err error) {
				onDone(err)
			}
		},
		OnGet: func(ctx *context.Context, call stack.Caller) func(item any, attempts int, err error) {
			onDone := trace.QueryOnPoolGet(t, ctx, call)

			return func(item any, attempts int, err error) {
				onDone(item.(*Session), attempts, err) //nolint:forcetypeassert
			}
		},
		OnChange: func(stats pool.Stats) {
			trace.QueryOnPoolChange(t, stats.Limit, stats.Index, stats.Idle, stats.Wait, stats.CreateInProgress)
		},
	}
}