Skip to content

Commit b96c5f2

Browse files
asmyasnikovasmyasnikov
authored
Merge pull request #1249 from ydb-platform/query-helper
* Added experimental query client helpers
2 parents a5b8c0a + eecac61 commit b96c5f2

27 files changed

+1251
-300
lines changed

‎CHANGELOG.md‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,12 @@
1-
* Added experimental `query.ReadRow` method for execute query and read only one row from result
1+
* Added experimental method for execute query and read only one row from result:
2+
* `query/Client.ReadRow`
3+
* `query/Session.ReadRow`
4+
* `query/Transaction.ReadRow`
5+
* Added experimental method for execute query and read only one result set from result:
6+
* `query/Client.ReadResultSet`
7+
* `query/Session.ReadResultSet`
8+
* `query/Transaction.ReadResultSet`
9+
* Added experimental `sugar.UnmarshallRow[T]` and `sugar.UnmarshallResultSet[T]` helpers for converts YDB rows to typed objects
210

311
## v3.68.1
412
* Downgraded minimal version of Go to 1.20

‎internal/query/client.go‎

Lines changed: 32 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package query
22

33
import (
44
"context"
5-
"io"
65

76
"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
87
"google.golang.org/grpc"
@@ -149,65 +148,62 @@ func doTx(
149148
return attempts, nil
150149
}
151150

152-
func readRow(ctx context.Context,
153-
pool *pool.Pool[*Session, Session],
154-
q string,
155-
t *trace.Query,
156-
opts ...options.ExecuteOption,
157-
) (row query.Row, err error) {
158-
_, err = do(ctx, pool, func(ctx context.Context, s query.Session) (err error) {
159-
_, r, err := s.Execute(ctx, q, opts...)
160-
if err != nil {
161-
return xerrors.WithStackTrace(err)
162-
}
163-
defer func() {
164-
_ = r.Close(ctx)
165-
}()
166-
rs, err := r.NextResultSet(ctx)
167-
if err != nil {
168-
return xerrors.WithStackTrace(err)
169-
}
170-
row, err = rs.NextRow(ctx)
171-
if err != nil {
172-
return xerrors.WithStackTrace(err)
173-
}
151+
// ReadRow is a helper which read only one row from first result set in result
152+
func (c *Client) ReadRow(ctx context.Context, q string, opts ...options.ExecuteOption) (row query.Row, err error) {
153+
ctx, cancel := xcontext.WithDone(ctx, c.done)
154+
defer cancel()
174155

175-
if _, err = rs.NextRow(ctx); err == nil || !xerrors.Is(err, io.EOF) {
176-
return xerrors.WithStackTrace(errMoreThanOneRow)
177-
}
156+
onDone := trace.QueryOnReadRow(c.config.Trace(), &ctx,
157+
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/query.(*Client).ReadRow"),
158+
q,
159+
)
160+
defer func() {
161+
onDone(err)
162+
}()
178163

179-
if _, err = r.NextResultSet(ctx); err == nil || !xerrors.Is(err, io.EOF) {
180-
return xerrors.WithStackTrace(errMoreThanOneResultSet)
164+
_, err = do(ctx, c.pool, func(ctx context.Context, s query.Session) (err error) {
165+
row, err = s.ReadRow(ctx, q, opts...)
166+
if err != nil {
167+
return xerrors.WithStackTrace(err)
181168
}
182169

183-
return r.Err()
184-
}, t)
170+
return nil
171+
}, c.config.Trace())
185172
if err != nil {
186173
return nil, xerrors.WithStackTrace(err)
187174
}
188175

189176
return row, nil
190177
}
191178

192-
// ReadRow is a helper which read only one row from first result set in result
193-
func (c *Client) ReadRow(ctx context.Context, q string, opts ...options.ExecuteOption) (row query.Row, err error) {
179+
// ReadResultSet is a helper which read all rows from first result set in result
180+
func (c *Client) ReadResultSet(
181+
ctx context.Context, q string, opts ...options.ExecuteOption,
182+
) (rs query.ResultSet, err error) {
194183
ctx, cancel := xcontext.WithDone(ctx, c.done)
195184
defer cancel()
196185

197-
onDone := trace.QueryOnReadRow(c.config.Trace(), &ctx,
198-
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/query.(*Client).ReadRow"),
186+
onDone := trace.QueryOnReadResultSet(c.config.Trace(), &ctx,
187+
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/query.(*Client).ReadResultSet"),
199188
q,
200189
)
201190
defer func() {
202191
onDone(err)
203192
}()
204193

205-
row, err = readRow(ctx, c.pool, q, c.config.Trace(), opts...)
194+
_, err = do(ctx, c.pool, func(ctx context.Context, s query.Session) (err error) {
195+
rs, err = s.ReadResultSet(ctx, q, opts...)
196+
if err != nil {
197+
return xerrors.WithStackTrace(err)
198+
}
199+
200+
return nil
201+
}, c.config.Trace())
206202
if err != nil {
207203
return nil, xerrors.WithStackTrace(err)
208204
}
209205

210-
return row, nil
206+
return rs, nil
211207
}
212208

213209
func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options.DoTxOption) (err error) {

‎internal/query/client_test.go‎

Lines changed: 0 additions & 235 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package query
33
import (
44
"context"
55
"errors"
6-
"io"
76
"testing"
87
"time"
98

@@ -257,237 +256,3 @@ func TestDoTx(t *testing.T) {
257256
require.Equal(t, 10, counter)
258257
})
259258
}
260-
261-
func TestReadRow(t *testing.T) {
262-
ctx := xtest.Context(t)
263-
t.Run("HappyWay", func(t *testing.T) {
264-
ctrl := gomock.NewController(t)
265-
row, err := readRow(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
266-
stream := NewMockQueryService_ExecuteQueryClient(ctrl)
267-
stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
268-
Status: Ydb.StatusIds_SUCCESS,
269-
TxMeta: &Ydb_Query.TransactionMeta{
270-
Id: "456",
271-
},
272-
ResultSetIndex: 0,
273-
ResultSet: &Ydb.ResultSet{
274-
Columns: []*Ydb.Column{
275-
{
276-
Name: "a",
277-
Type: &Ydb.Type{
278-
Type: &Ydb.Type_TypeId{
279-
TypeId: Ydb.Type_UINT64,
280-
},
281-
},
282-
},
283-
{
284-
Name: "b",
285-
Type: &Ydb.Type{
286-
Type: &Ydb.Type_TypeId{
287-
TypeId: Ydb.Type_UTF8,
288-
},
289-
},
290-
},
291-
},
292-
Rows: []*Ydb.Value{
293-
{
294-
Items: []*Ydb.Value{{
295-
Value: &Ydb.Value_Uint64Value{
296-
Uint64Value: 1,
297-
},
298-
}, {
299-
Value: &Ydb.Value_TextValue{
300-
TextValue: "1",
301-
},
302-
}},
303-
},
304-
},
305-
},
306-
}, nil)
307-
stream.EXPECT().Recv().Return(nil, io.EOF)
308-
client := NewMockQueryServiceClient(ctrl)
309-
client.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
310-
311-
return &Session{
312-
id: "123",
313-
statusCode: statusIdle,
314-
cfg: config.New(),
315-
grpcClient: client,
316-
}, nil
317-
}), "", nil)
318-
require.NoError(t, err)
319-
var (
320-
a uint64
321-
b string
322-
)
323-
err = row.Scan(&a, &b)
324-
require.NoError(t, err)
325-
require.EqualValues(t, 1, a)
326-
require.EqualValues(t, "1", b)
327-
})
328-
t.Run("MoreThanOneRow", func(t *testing.T) {
329-
ctrl := gomock.NewController(t)
330-
row, err := readRow(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
331-
stream := NewMockQueryService_ExecuteQueryClient(ctrl)
332-
stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
333-
Status: Ydb.StatusIds_SUCCESS,
334-
TxMeta: &Ydb_Query.TransactionMeta{
335-
Id: "456",
336-
},
337-
ResultSetIndex: 0,
338-
ResultSet: &Ydb.ResultSet{
339-
Columns: []*Ydb.Column{
340-
{
341-
Name: "a",
342-
Type: &Ydb.Type{
343-
Type: &Ydb.Type_TypeId{
344-
TypeId: Ydb.Type_UINT64,
345-
},
346-
},
347-
},
348-
{
349-
Name: "b",
350-
Type: &Ydb.Type{
351-
Type: &Ydb.Type_TypeId{
352-
TypeId: Ydb.Type_UTF8,
353-
},
354-
},
355-
},
356-
},
357-
Rows: []*Ydb.Value{
358-
{
359-
Items: []*Ydb.Value{{
360-
Value: &Ydb.Value_Uint64Value{
361-
Uint64Value: 1,
362-
},
363-
}, {
364-
Value: &Ydb.Value_TextValue{
365-
TextValue: "1",
366-
},
367-
}},
368-
},
369-
{
370-
Items: []*Ydb.Value{{
371-
Value: &Ydb.Value_Uint64Value{
372-
Uint64Value: 2,
373-
},
374-
}, {
375-
Value: &Ydb.Value_TextValue{
376-
TextValue: "2",
377-
},
378-
}},
379-
},
380-
},
381-
},
382-
}, nil)
383-
client := NewMockQueryServiceClient(ctrl)
384-
client.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
385-
386-
return &Session{
387-
id: "123",
388-
statusCode: statusIdle,
389-
cfg: config.New(),
390-
grpcClient: client,
391-
}, nil
392-
}), "", nil)
393-
require.ErrorIs(t, err, errMoreThanOneRow)
394-
require.Nil(t, row)
395-
})
396-
t.Run("MoreThanOneResultSet", func(t *testing.T) {
397-
ctrl := gomock.NewController(t)
398-
row, err := readRow(ctx, testPool(ctx, func(ctx context.Context) (*Session, error) {
399-
stream := NewMockQueryService_ExecuteQueryClient(ctrl)
400-
stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
401-
Status: Ydb.StatusIds_SUCCESS,
402-
TxMeta: &Ydb_Query.TransactionMeta{
403-
Id: "456",
404-
},
405-
ResultSetIndex: 0,
406-
ResultSet: &Ydb.ResultSet{
407-
Columns: []*Ydb.Column{
408-
{
409-
Name: "a",
410-
Type: &Ydb.Type{
411-
Type: &Ydb.Type_TypeId{
412-
TypeId: Ydb.Type_UINT64,
413-
},
414-
},
415-
},
416-
{
417-
Name: "b",
418-
Type: &Ydb.Type{
419-
Type: &Ydb.Type_TypeId{
420-
TypeId: Ydb.Type_UTF8,
421-
},
422-
},
423-
},
424-
},
425-
Rows: []*Ydb.Value{
426-
{
427-
Items: []*Ydb.Value{{
428-
Value: &Ydb.Value_Uint64Value{
429-
Uint64Value: 1,
430-
},
431-
}, {
432-
Value: &Ydb.Value_TextValue{
433-
TextValue: "1",
434-
},
435-
}},
436-
},
437-
},
438-
},
439-
}, nil)
440-
stream.EXPECT().Recv().Return(&Ydb_Query.ExecuteQueryResponsePart{
441-
Status: Ydb.StatusIds_SUCCESS,
442-
TxMeta: &Ydb_Query.TransactionMeta{
443-
Id: "456",
444-
},
445-
ResultSetIndex: 0,
446-
ResultSet: &Ydb.ResultSet{
447-
Columns: []*Ydb.Column{
448-
{
449-
Name: "a",
450-
Type: &Ydb.Type{
451-
Type: &Ydb.Type_TypeId{
452-
TypeId: Ydb.Type_UINT64,
453-
},
454-
},
455-
},
456-
{
457-
Name: "b",
458-
Type: &Ydb.Type{
459-
Type: &Ydb.Type_TypeId{
460-
TypeId: Ydb.Type_UTF8,
461-
},
462-
},
463-
},
464-
},
465-
Rows: []*Ydb.Value{
466-
{
467-
Items: []*Ydb.Value{{
468-
Value: &Ydb.Value_Uint64Value{
469-
Uint64Value: 1,
470-
},
471-
}, {
472-
Value: &Ydb.Value_TextValue{
473-
TextValue: "1",
474-
},
475-
}},
476-
},
477-
},
478-
},
479-
}, nil)
480-
client := NewMockQueryServiceClient(ctrl)
481-
client.EXPECT().ExecuteQuery(gomock.Any(), gomock.Any()).Return(stream, nil)
482-
483-
return &Session{
484-
id: "123",
485-
statusCode: statusIdle,
486-
cfg: config.New(),
487-
grpcClient: client,
488-
}, nil
489-
}), "", nil)
490-
require.ErrorIs(t, err, errMoreThanOneRow)
491-
require.Nil(t, row)
492-
})
493-
}

‎internal/query/errors.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,5 @@ var (
1111
errWrongResultSetIndex = errors.New("critical violation of the logic - wrong result set index")
1212
errMoreThanOneRow = errors.New("unexpected more than one row in result set")
1313
errMoreThanOneResultSet = errors.New("unexpected more than one result set")
14+
errNoResultSets = errors.New("no result sets")
1415
)

0 commit comments

Comments
 (0)