Skip to content

Commit b3d2e7e

Browse files
asmyasnikovasmyasnikov
authored
Merge pull request #1696 from hurd54/metric-export-extended-info-with-labels
metric-export-extended-info-with-labels
2 parents 4082973 + ae792aa commit b3d2e7e

File tree

12 files changed

+349
-133
lines changed

12 files changed

+349
-133
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
* Added export of advanced metric information for QueryService calls
2+
13
## v3.104.0
24
* Added binding `ydb.WithWideTimeTypes()` which interprets `time.Time` and `time.Duration` as `Timestamp64` and `Interval64` YDB types
35

‎internal/query/client.go‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,7 @@ func (c *Client) Do(ctx context.Context, op query.Operation, opts ...options.DoO
237237
settings = options.ParseDoOpts(c.config.Trace(), opts...)
238238
onDone = trace.QueryOnDo(settings.Trace(), &ctx,
239239
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Do"),
240+
settings.Label(),
240241
)
241242
attempts = 0
242243
)
@@ -321,15 +322,17 @@ func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute
321322
ctx, cancel := xcontext.WithDone(ctx, c.done)
322323
defer cancel()
323324

325+
settings := options.ExecuteSettings(opts...)
326+
324327
onDone := trace.QueryOnQueryRow(c.config.Trace(), &ctx,
325328
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).QueryRow"),
326-
q,
329+
q, settings.Label(),
327330
)
328331
defer func() {
329332
onDone(finalErr)
330333
}()
331334

332-
row, err := clientQueryRow(ctx, c.pool, q, options.ExecuteSettings(opts...), withTrace(c.config.Trace()))
335+
row, err := clientQueryRow(ctx, c.pool, q, settings, withTrace(c.config.Trace()))
333336
if err != nil {
334337
return nil, xerrors.WithStackTrace(err)
335338
}
@@ -366,9 +369,11 @@ func (c *Client) Exec(ctx context.Context, q string, opts ...options.Execute) (f
366369
ctx, cancel := xcontext.WithDone(ctx, c.done)
367370
defer cancel()
368371

372+
settings := options.ExecuteSettings(opts...)
369373
onDone := trace.QueryOnExec(c.config.Trace(), &ctx,
370374
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Exec"),
371375
q,
376+
settings.Label(),
372377
)
373378
defer func() {
374379
onDone(finalErr)
@@ -413,9 +418,10 @@ func (c *Client) Query(ctx context.Context, q string, opts ...options.Execute) (
413418
ctx, cancel := xcontext.WithDone(ctx, c.done)
414419
defer cancel()
415420

421+
settings := options.ExecuteSettings(opts...)
416422
onDone := trace.QueryOnQuery(c.config.Trace(), &ctx,
417423
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).Query"),
418-
q,
424+
q, settings.Label(),
419425
)
420426
defer func() {
421427
onDone(err)
@@ -431,7 +437,7 @@ func (c *Client) Query(ctx context.Context, q string, opts ...options.Execute) (
431437

432438
func clientQueryResultSet(
433439
ctx context.Context, pool sessionPool, q string, settings executeSettings, resultOpts ...resultOption,
434-
) (rs result.ClosableResultSet, finalErr error) {
440+
) (rs result.ClosableResultSet, rowsCount int, finalErr error) {
435441
err := do(ctx, pool, func(ctx context.Context, s *Session) error {
436442
streamResult, err := s.execute(ctx, q, settings, resultOpts...)
437443
if err != nil {
@@ -441,18 +447,18 @@ func clientQueryResultSet(
441447
_ = streamResult.Close(ctx)
442448
}()
443449

444-
rs, err = readMaterializedResultSet(ctx, streamResult)
450+
rs, rowsCount, err = readMaterializedResultSet(ctx, streamResult)
445451
if err != nil {
446452
return xerrors.WithStackTrace(err)
447453
}
448454

449455
return nil
450456
}, settings.RetryOpts()...)
451457
if err != nil {
452-
return nil, xerrors.WithStackTrace(err)
458+
return nil, 0, xerrors.WithStackTrace(err)
453459
}
454460

455-
return rs, nil
461+
return rs, rowsCount, nil
456462
}
457463

458464
// QueryResultSet is a helper which read all rows from first result set in result
@@ -462,15 +468,21 @@ func (c *Client) QueryResultSet(
462468
ctx, cancel := xcontext.WithDone(ctx, c.done)
463469
defer cancel()
464470

471+
var (
472+
settings = options.ExecuteSettings(opts...)
473+
rowsCount int
474+
err error
475+
)
476+
465477
onDone := trace.QueryOnQueryResultSet(c.config.Trace(), &ctx,
466478
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).QueryResultSet"),
467-
q,
479+
q, settings.Label(),
468480
)
469481
defer func() {
470-
onDone(finalErr)
482+
onDone(finalErr, rowsCount)
471483
}()
472484

473-
rs, err := clientQueryResultSet(ctx, c.pool, q, options.ExecuteSettings(opts...), withTrace(c.config.Trace()))
485+
rs, rowsCount, err = clientQueryResultSet(ctx, c.pool, q, settings, withTrace(c.config.Trace()))
474486
if err != nil {
475487
return nil, xerrors.WithStackTrace(err)
476488
}
@@ -486,6 +498,7 @@ func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options
486498
settings = options.ParseDoTxOpts(c.config.Trace(), opts...)
487499
onDone = trace.QueryOnDoTx(settings.Trace(), &ctx,
488500
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Client).DoTx"),
501+
settings.Label(),
489502
)
490503
attempts = 0
491504
)

‎internal/query/client_test.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1083,7 +1083,7 @@ func TestClient(t *testing.T) {
10831083
t.Run("QueryResultSet", func(t *testing.T) {
10841084
t.Run("HappyWay", func(t *testing.T) {
10851085
ctrl := gomock.NewController(t)
1086-
rs, err := clientQueryResultSet(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
1086+
rs, rowsCount, err := clientQueryResultSet(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
10871087
stream := NewMockQueryService_ExecuteQueryClient(ctrl)
10881088
stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
10891089
Status: Ydb.StatusIds_SUCCESS,
@@ -1185,6 +1185,7 @@ func TestClient(t *testing.T) {
11851185
}), "", options.ExecuteSettings())
11861186
require.NoError(t, err)
11871187
require.NotNil(t, rs)
1188+
require.Equal(t, 5, rowsCount)
11881189
{
11891190
require.NoError(t, err)
11901191
r1, err := rs.NextRow(ctx)
@@ -1228,7 +1229,7 @@ func TestClient(t *testing.T) {
12281229
})
12291230
t.Run("MoreThanOneResultSet", func(t *testing.T) {
12301231
ctrl := gomock.NewController(t)
1231-
rs, err := clientQueryResultSet(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
1232+
rs, rowsCount, err := clientQueryResultSet(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
12321233
stream := NewMockQueryService_ExecuteQueryClient(ctrl)
12331234
stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
12341235
Status: Ydb.StatusIds_SUCCESS,
@@ -1394,6 +1395,7 @@ func TestClient(t *testing.T) {
13941395
}), "", options.ExecuteSettings())
13951396
require.ErrorIs(t, err, errMoreThanOneResultSet)
13961397
require.Nil(t, rs)
1398+
require.Equal(t, 0, rowsCount)
13971399
})
13981400
})
13991401
t.Run("QueryRow", func(t *testing.T) {

‎internal/query/execute_query.go‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ type executeSettings interface {
3232
RetryOpts() []retry.Option
3333
ResourcePool() string
3434
ResponsePartLimitSizeBytes() int64
35+
Label() string
3536
}
3637

3738
type executeScriptConfig interface {
@@ -179,14 +180,16 @@ func readResultSet(ctx context.Context, r *streamResult) (_ *resultSetWithClose,
179180
}, nil
180181
}
181182

182-
func readMaterializedResultSet(ctx context.Context, r *streamResult) (_ *materializedResultSet, finalErr error) {
183+
func readMaterializedResultSet(ctx context.Context, r *streamResult) (
184+
_ *materializedResultSet, rowsCount int, finalErr error,
185+
) {
183186
defer func() {
184187
_ = r.Close(ctx)
185188
}()
186189

187190
rs, err := r.nextResultSet(ctx)
188191
if err != nil {
189-
return nil, xerrors.WithStackTrace(err)
192+
return nil, 0, xerrors.WithStackTrace(err)
190193
}
191194

192195
var rows []query.Row
@@ -197,21 +200,21 @@ func readMaterializedResultSet(ctx context.Context, r *streamResult) (_ *materia
197200
break
198201
}
199202

200-
return nil, xerrors.WithStackTrace(err)
203+
return nil, 0, xerrors.WithStackTrace(err)
201204
}
202205

203206
rows = append(rows, row)
204207
}
205208

206209
_, err = r.nextResultSet(ctx)
207210
if err == nil {
208-
return nil, xerrors.WithStackTrace(errMoreThanOneResultSet)
211+
return nil, 0, xerrors.WithStackTrace(errMoreThanOneResultSet)
209212
}
210213
if !xerrors.Is(err, io.EOF) {
211-
return nil, xerrors.WithStackTrace(err)
214+
return nil, 0, xerrors.WithStackTrace(err)
212215
}
213216

214-
return MaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows), nil
217+
return MaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows), len(rows), nil
215218
}
216219

217220
func readRow(ctx context.Context, r *streamResult) (_ *Row, finalErr error) {

‎internal/query/options/execute.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ type (
4646
txControl TxControl
4747
retryOptions []retry.Option
4848
responsePartLimitBytes int64
49+
label string
4950
}
5051

5152
// Execute is an interface for execute method options
@@ -164,6 +165,7 @@ func defaultExecuteSettings() executeSettings {
164165
statsMode: StatsModeNone,
165166
txControl: tx.DefaultTxControl(),
166167
params: &params.Params{},
168+
label: "undefined",
167169
}
168170
}
169171

@@ -211,6 +213,10 @@ func (s *executeSettings) ResponsePartLimitSizeBytes() int64 {
211213
return s.responsePartLimitBytes
212214
}
213215

216+
func (s *executeSettings) Label() string {
217+
return s.label
218+
}
219+
214220
func WithParameters(params params.Parameters) parametersOption {
215221
return parametersOption{
216222
params: params,

‎internal/query/options/retry.go‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ type (
2424
doSettings struct {
2525
retryOpts []retry.Option
2626
trace *trace.Query
27+
label string
2728
}
2829

2930
DoTxOption interface {
@@ -44,6 +45,22 @@ type (
4445
}
4546
)
4647

48+
type LabelOption string
49+
50+
func (label LabelOption) applyDoOption(s *doSettings) {
51+
s.label = string(label)
52+
RetryOptionsOption{retry.WithLabel(s.label)}.applyDoOption(s)
53+
}
54+
55+
func (label LabelOption) applyDoTxOption(s *doTxSettings) {
56+
s.label = string(label)
57+
RetryOptionsOption{retry.WithLabel(s.label)}.applyDoTxOption(s)
58+
}
59+
60+
func (label LabelOption) applyExecuteOption(s *executeSettings) {
61+
s.label = string(label)
62+
}
63+
4764
func (opts RetryOptionsOption) applyExecuteOption(s *executeSettings) {
4865
s.retryOptions = append(s.retryOptions, opts...)
4966
}
@@ -56,6 +73,10 @@ func (s *doSettings) RetryOpts() []retry.Option {
5673
return s.retryOpts
5774
}
5875

76+
func (s *doSettings) Label() string {
77+
return s.label
78+
}
79+
5980
func (s *doTxSettings) TxSettings() tx.Settings {
6081
return s.txSettings
6182
}
@@ -88,8 +109,8 @@ func WithIdempotent() RetryOptionsOption {
88109
return []retry.Option{retry.WithIdempotent(true)}
89110
}
90111

91-
func WithLabel(lbl string) RetryOptionsOption {
92-
return []retry.Option{retry.WithLabel(lbl)}
112+
func WithLabel(lbl string) LabelOption {
113+
return LabelOption(lbl)
93114
}
94115

95116
func WithTrace(t *trace.Query) TraceOption {
@@ -103,6 +124,7 @@ func WithRetryBudget(b budget.Budget) RetryOptionsOption {
103124
func ParseDoOpts(t *trace.Query, opts ...DoOption) (s *doSettings) {
104125
s = &doSettings{
105126
trace: t,
127+
label: "undefined",
106128
}
107129

108130
for _, opt := range opts {

‎internal/query/session.go‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -152,11 +152,14 @@ func (s *Session) execute(
152152
return r, nil
153153
}
154154

155-
func (s *Session) Exec(
156-
ctx context.Context, q string, opts ...options.Execute,
157-
) (finalErr error) {
155+
func (s *Session) Exec(ctx context.Context, q string, opts ...options.Execute) (finalErr error) {
156+
settings := options.ExecuteSettings(opts...)
158157
onDone := trace.QueryOnSessionExec(s.trace, &ctx,
159-
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Exec"), s, q)
158+
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Exec"),
159+
s,
160+
q,
161+
settings.Label(),
162+
)
160163
defer func() {
161164
onDone(finalErr)
162165
}()
@@ -177,11 +180,14 @@ func (s *Session) Exec(
177180
return nil
178181
}
179182

180-
func (s *Session) Query(
181-
ctx context.Context, q string, opts ...options.Execute,
182-
) (_ query.Result, finalErr error) {
183+
func (s *Session) Query(ctx context.Context, q string, opts ...options.Execute) (_ query.Result, finalErr error) {
184+
settings := options.ExecuteSettings(opts...)
183185
onDone := trace.QueryOnSessionQuery(s.trace, &ctx,
184-
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Query"), s, q)
186+
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*Session).Query"),
187+
s,
188+
q,
189+
settings.Label(),
190+
)
185191
defer func() {
186192
onDone(finalErr)
187193
}()

0 commit comments

Comments
 (0)