Merge pull request #1117 from ydb-platform/query-logs

logging for `query.{Do,DoTx}` + refactoring of `query.{Do,DoTx}Options`

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,7 @@
+* Added logs over query service internals
+* Changed `trace.Query` events
+* Changed visibility of `query.{Do,DoTx}Options` from public to private
+
 ## v3.57.0
 * Added experimental implementation of query service client
 * Fixed sometime panic on topic writer closing

internal/conn/grpc_client_stream.go

@@ -135,9 +135,10 @@ func (s *grpcClientStream) wrapError(err error) error {
 		return nil
 	}
 
-	nodeErr := newConnError(s.c.endpoint.NodeID(), s.c.endpoint.Address(), err)
-
-	return xerrors.WithStackTrace(nodeErr, xerrors.WithSkipDepth(1))
+	return xerrors.WithStackTrace(
+		newConnError(s.c.endpoint.NodeID(), s.c.endpoint.Address(), err),
+		xerrors.WithSkipDepth(1),
+	)
 }
 
 func createPinger(c *conn) context.CancelFunc {

internal/pool/pool.go

@@ -107,15 +107,15 @@ func (p *Pool[T]) try(ctx context.Context, f func(ctx context.Context, item *T)
 	return nil
 }
 
-func (p *Pool[T]) With(ctx context.Context, f func(ctx context.Context, item *T) error) error {
+func (p *Pool[T]) With(ctx context.Context, f func(ctx context.Context, item *T) error, opts ...retry.Option) error {
 	err := retry.Retry(ctx, func(ctx context.Context) error {
 		err := p.try(ctx, f)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
 
 		return nil
-	})
+	}, opts...)
 	if err != nil {
 		return xerrors.WithStackTrace(err)
 	}

internal/query/client.go

@@ -12,10 +12,13 @@ import (
 
 	"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/stack"
 	"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 -package query -write_package_comment=false github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1 QueryServiceClient,QueryService_AttachSessionClient,QueryService_ExecuteQueryClient
@@ -27,6 +30,7 @@ type balancer interface {
 var _ query.Client = (*Client)(nil)
 
 type Client struct {
+	config     *config.Config
 	grpcClient Ydb_Query_V1.QueryServiceClient
 	pool       *pool.Pool[Session]
 }
@@ -40,39 +44,63 @@ func (c Client) Close(ctx context.Context) error {
 	return nil
 }
 
-func do(ctx context.Context, pool *pool.Pool[Session], op query.Operation, opts *query.DoOptions) error {
-	return retry.Retry(ctx, func(ctx context.Context) error {
-		err := pool.With(ctx, func(ctx context.Context, s *Session) error {
-			err := op(ctx, s)
-			if err != nil {
-				return xerrors.WithStackTrace(err)
-			}
-
-			return nil
-		})
+func do(
+	ctx context.Context,
+	pool *pool.Pool[Session],
+	op query.Operation,
+	t *trace.Query,
+	opts ...options.DoOption,
+) (finalErr error) {
+	doOpts := options.ParseDoOpts(t, opts...)
+
+	err := pool.With(ctx, func(ctx context.Context, s *Session) error {
+		err := op(ctx, s)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
 
 		return nil
-	}, opts.RetryOptions...)
-}
-
-func (c Client) Do(ctx context.Context, op query.Operation, opts ...query.DoOption) error {
-	doOptions := query.NewDoOptions(opts...)
-	if doOptions.Label != "" {
-		doOptions.RetryOptions = append(doOptions.RetryOptions, retry.WithLabel(doOptions.Label))
-	}
-	if doOptions.Idempotent {
-		doOptions.RetryOptions = append(doOptions.RetryOptions, retry.WithIdempotent(doOptions.Idempotent))
+	}, append(doOpts.RetryOpts(), retry.WithTrace(&trace.Retry{
+		OnRetry: func(
+			info trace.RetryLoopStartInfo,
+		) func(
+			trace.RetryLoopIntermediateInfo,
+		) func(
+			trace.RetryLoopDoneInfo,
+		) {
+			onIntermediate := trace.QueryOnDo(doOpts.Trace(), &ctx, stack.FunctionID(""))
+
+			return func(info trace.RetryLoopIntermediateInfo) func(trace.RetryLoopDoneInfo) {
+				onDone := onIntermediate(info.Error)
+
+				return func(info trace.RetryLoopDoneInfo) {
+					onDone(info.Attempts, info.Error)
+				}
+			}
+		},
+	}))...)
+	if err != nil {
+		return xerrors.WithStackTrace(err)
 	}
 
-	return do(ctx, c.pool, op, &doOptions)
+	return nil
+}
+
+func (c Client) Do(ctx context.Context, op query.Operation, opts ...options.DoOption) error {
+	return do(ctx, c.pool, op, c.config.Trace(), opts...)
 }
 
-func doTx(ctx context.Context, pool *pool.Pool[Session], op query.TxOperation, opts *query.DoTxOptions) error {
-	return do(ctx, pool, func(ctx context.Context, s query.Session) error {
-		tx, err := s.Begin(ctx, opts.TxSettings)
+func doTx(
+	ctx context.Context,
+	pool *pool.Pool[Session],
+	op query.TxOperation,
+	t *trace.Query,
+	opts ...options.DoTxOption,
+) error {
+	doTxOpts := options.ParseDoTxOpts(t, opts...)
+
+	err := do(ctx, pool, func(ctx context.Context, s query.Session) error {
+		tx, err := s.Begin(ctx, doTxOpts.TxSettings())
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
@@ -96,19 +124,16 @@ func doTx(ctx context.Context, pool *pool.Pool[Session], op query.TxOperation, o
 		}
 
 		return nil
-	}, &opts.DoOptions)
-}
-
-func (c Client) DoTx(ctx context.Context, op query.TxOperation, opts ...query.DoTxOption) error {
-	doTxOptions := query.NewDoTxOptions(opts...)
-	if doTxOptions.Label != "" {
-		doTxOptions.RetryOptions = append(doTxOptions.RetryOptions, retry.WithLabel(doTxOptions.Label))
-	}
-	if doTxOptions.Idempotent {
-		doTxOptions.RetryOptions = append(doTxOptions.RetryOptions, retry.WithIdempotent(doTxOptions.Idempotent))
+	}, t, doTxOpts.DoOpts()...)
+	if err != nil {
+		return xerrors.WithStackTrace(err)
 	}
 
-	return doTx(ctx, c.pool, op, &doTxOptions)
+	return nil
+}
+
+func (c Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options.DoTxOption) error {
+	return doTx(ctx, c.pool, op, c.config.Trace(), opts...)
 }
 
 func deleteSession(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, sessionID string) error {
@@ -231,6 +256,7 @@ func createSession(
 
 func New(ctx context.Context, balancer balancer, config *config.Config) (*Client, error) {
 	client := &Client{
+		config:     config,
 		grpcClient: Ydb_Query_V1.NewQueryServiceClient(balancer),
 	}
 

internal/query/client_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 func TestCreateSession(t *testing.T) {
@@ -182,7 +183,7 @@ func TestDo(t *testing.T) {
 			return newTestSession()
 		}), func(ctx context.Context, s query.Session) error {
 			return nil
-		}, &query.DoOptions{})
+		}, &trace.Query{})
 		require.NoError(t, err)
 	})
 	t.Run("RetryableError", func(t *testing.T) {
@@ -196,7 +197,7 @@ func TestDo(t *testing.T) {
 			}
 
 			return nil
-		}, &query.DoOptions{})
+		}, &trace.Query{})
 		require.NoError(t, err)
 		require.Equal(t, 10, counter)
 	})
@@ -217,7 +218,7 @@ func TestDoTx(t *testing.T) {
 			return newTestSessionWithClient(client)
 		}), func(ctx context.Context, tx query.TxActor) error {
 			return nil
-		}, &query.DoTxOptions{})
+		}, &trace.Query{})
 		require.NoError(t, err)
 	})
 	t.Run("RetryableError", func(t *testing.T) {
@@ -242,7 +243,7 @@ func TestDoTx(t *testing.T) {
 			}
 
 			return nil
-		}, &query.DoTxOptions{})
+		}, &trace.Query{})
 		require.NoError(t, err)
 		require.Equal(t, 10, counter)
 	})

internal/query/config/options.go

@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/config"
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 type Option func(*Config)
@@ -15,6 +16,13 @@ func With(config config.Common) Option {
 	}
 }
 
+// WithTrace appends table trace to early defined traces
+func WithTrace(trace *trace.Query, opts ...trace.QueryComposeOption) Option {
+	return func(c *Config) {
+		c.trace = c.trace.Compose(trace, opts...)
+	}
+}
+
 // WithSizeLimit defines upper bound of pooled sessions.
 // If sizeLimit is less than or equal to zero then the
 // DefaultPoolMaxSize variable is used as a limit.

internal/query/errors.go

@@ -7,7 +7,6 @@ import (
 var (
 	ErrNotImplemented          = errors.New("not implemented yet")
 	errWrongNextResultSetIndex = errors.New("wrong result set index")
-	errInterruptedStream       = errors.New("interrupted stream")
 	errClosedResult            = errors.New("result closed early")
 	errWrongResultSetIndex     = errors.New("critical violation of the logic - wrong result set index")
 )

internal/query/execute_query.go

@@ -9,16 +9,17 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
 	"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"
 )
 
 type executeConfig interface {
-	ExecMode() query.ExecMode
-	StatsMode() query.StatsMode
+	ExecMode() options.ExecMode
+	StatsMode() options.StatsMode
 	TxControl() *query.TransactionControl
-	Syntax() query.Syntax
+	Syntax() options.Syntax
 	Params() *params.Parameters
 	CallOptions() []grpc.CallOption
 }