Skip to content

Commit c2b8607

Browse files
asmyasnikovasmyasnikov
authored
Merge pull request #1313 from ydb-platform/add-ResultSetIndex-method
Add `query.ResultSet.Index()` method
2 parents 44ddaa6 + f2f56df commit c2b8607

File tree

8 files changed

+54
-30
lines changed

8 files changed

+54
-30
lines changed

‎CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
* Added `query.ResultSet.Index()` method
2+
13
## v3.74.5
24
* Fixed bug with reading empty result set parts.
35
* Fixed nil pointer dereference when closing result set

‎internal/query/result.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ func (r *result) nextResultSet(ctx context.Context) (_ *resultSet, err error) {
191191
}
192192
if part.GetResultSetIndex() < r.resultSetIndex {
193193
return nil, xerrors.WithStackTrace(fmt.Errorf(
194-
"next result set index %d less than last result set index %d: %w",
194+
"next result set rowIndex %d less than last result set index %d: %w",
195195
part.GetResultSetIndex(), r.resultSetIndex, errWrongNextResultSetIndex,
196196
))
197197
}
@@ -324,7 +324,7 @@ func exactlyOneResultSetFromResult(ctx context.Context, r query.Result) (rs quer
324324
return nil, xerrors.WithStackTrace(err)
325325
}
326326

327-
return NewMaterializedResultSet(rs.Columns(), rs.ColumnTypes(), rows), nil
327+
return NewMaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows), nil
328328
}
329329

330330
func resultToMaterializedResult(ctx context.Context, r query.Result) (query.Result, error) {
@@ -354,7 +354,7 @@ func resultToMaterializedResult(ctx context.Context, r query.Result) (query.Resu
354354
rows = append(rows, row)
355355
}
356356

357-
resultSets = append(resultSets, NewMaterializedResultSet(rs.Columns(), rs.ColumnTypes(), rows))
357+
resultSets = append(resultSets, NewMaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows))
358358
}
359359

360360
return newMaterializedResult(resultSets), nil

‎internal/query/result_set.go

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,11 @@ var (
2323

2424
type (
2525
materializedResultSet struct {
26+
index int
2627
columnNames []string
2728
columnTypes []types.Type
2829
rows []query.Row
29-
idx int
30+
rowIndex int
3031
}
3132
resultSet struct {
3233
index int64
@@ -74,23 +75,33 @@ func (rs *resultSet) Columns() (columnNames []string) {
7475
}
7576

7677
func (rs *materializedResultSet) NextRow(ctx context.Context) (query.Row, error) {
77-
if rs.idx == len(rs.rows) {
78+
if rs.rowIndex == len(rs.rows) {
7879
return nil, xerrors.WithStackTrace(io.EOF)
7980
}
8081

8182
defer func() {
82-
rs.idx++
83+
rs.rowIndex++
8384
}()
8485

85-
return rs.rows[rs.idx], nil
86+
return rs.rows[rs.rowIndex], nil
87+
}
88+
89+
func (rs *materializedResultSet) Index() int {
90+
if rs == nil {
91+
return -1
92+
}
93+
94+
return rs.index
8695
}
8796

8897
func NewMaterializedResultSet(
98+
index int,
8999
columnNames []string,
90100
columnTypes []types.Type,
91101
rows []query.Row,
92102
) *materializedResultSet {
93103
return &materializedResultSet{
104+
index: index,
94105
columnNames: columnNames,
95106
columnTypes: columnTypes,
96107
rows: rows,
@@ -169,3 +180,11 @@ func (rs *resultSet) NextRow(ctx context.Context) (_ query.Row, err error) {
169180

170181
return rs.nextRow(ctx)
171182
}
183+
184+
func (rs *resultSet) Index() int {
185+
if rs == nil {
186+
return -1
187+
}
188+
189+
return int(rs.index)
190+
}

‎internal/query/result_set_range_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -681,7 +681,8 @@ func TestResultSetRange(t *testing.T) {
681681
}
682682
if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
683683
return nil, xerrors.WithStackTrace(fmt.Errorf(
684-
"critical violation of the logic: wrong result set index: %d != %d",
684+
"%w: %d != %d",
685+
errWrongNextResultSetIndex,
685686
resultSetIndex, 0,
686687
))
687688
}
@@ -771,7 +772,8 @@ func TestResultSetRange(t *testing.T) {
771772
}
772773
if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
773774
return nil, xerrors.WithStackTrace(fmt.Errorf(
774-
"critical violation of the logic: wrong result set index: %d != %d",
775+
"%w: %d != %d",
776+
errWrongNextResultSetIndex,
775777
resultSetIndex, 0,
776778
))
777779
}

‎internal/query/result_set_test.go

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ func TestResultSetNext(t *testing.T) {
5959

6060
return part, nil
6161
}, recv, nil)
62-
require.EqualValues(t, 0, rs.index)
62+
require.EqualValues(t, 0, rs.Index())
6363
{
6464
_, err := rs.nextRow(ctx)
6565
require.ErrorIs(t, err, io.EOF)
@@ -144,7 +144,7 @@ func TestResultSetNext(t *testing.T) {
144144

145145
return part, nil
146146
}, recv, nil)
147-
require.EqualValues(t, 0, rs.index)
147+
require.EqualValues(t, 0, rs.Index())
148148
{
149149
_, err := rs.nextRow(ctx)
150150
require.NoError(t, err)
@@ -303,7 +303,7 @@ func TestResultSetNext(t *testing.T) {
303303

304304
return part, nil
305305
}, recv, nil)
306-
require.EqualValues(t, 0, rs.index)
306+
require.EqualValues(t, 0, rs.Index())
307307
{
308308
_, err := rs.nextRow(ctx)
309309
require.NoError(t, err)
@@ -441,7 +441,7 @@ func TestResultSetNext(t *testing.T) {
441441

442442
return part, nil
443443
}, recv, nil)
444-
require.EqualValues(t, 0, rs.index)
444+
require.EqualValues(t, 0, rs.Index())
445445
{
446446
_, err := rs.nextRow(ctx)
447447
require.NoError(t, err)
@@ -544,7 +544,7 @@ func TestResultSetNext(t *testing.T) {
544544

545545
return part, nil
546546
}, recv, nil)
547-
require.EqualValues(t, 0, rs.index)
547+
require.EqualValues(t, 0, rs.Index())
548548
{
549549
_, err := rs.nextRow(childCtx)
550550
require.NoError(t, err)
@@ -629,14 +629,14 @@ func TestResultSetNext(t *testing.T) {
629629
}
630630
if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
631631
return nil, xerrors.WithStackTrace(fmt.Errorf(
632-
"critical violation of the logic: wrong result set index: %d != %d",
632+
"critical violation of the logic: wrong result set rowIndex: %d != %d",
633633
resultSetIndex, 0,
634634
))
635635
}
636636

637637
return part, nil
638638
}, recv, nil)
639-
require.EqualValues(t, 0, rs.index)
639+
require.EqualValues(t, 0, rs.Index())
640640
{
641641
_, err := rs.nextRow(ctx)
642642
require.NoError(t, err)
@@ -728,14 +728,14 @@ func TestResultSetNext(t *testing.T) {
728728
}
729729
if resultSetIndex := part.GetResultSetIndex(); resultSetIndex != 0 {
730730
return nil, xerrors.WithStackTrace(fmt.Errorf(
731-
"critical violation of the logic: wrong result set index: %d != %d",
731+
"critical violation of the logic: wrong result set rowIndex: %d != %d",
732732
resultSetIndex, 0,
733733
))
734734
}
735735

736736
return part, nil
737737
}, recv, nil)
738-
require.EqualValues(t, 0, rs.index)
738+
require.EqualValues(t, 0, rs.Index())
739739
{
740740
_, err := rs.nextRow(ctx)
741741
require.NoError(t, err)
@@ -886,7 +886,7 @@ func TestResultSetNext(t *testing.T) {
886886

887887
return part, nil
888888
}, recv, nil)
889-
require.EqualValues(t, 0, rs.index)
889+
require.EqualValues(t, 0, rs.Index())
890890
{
891891
_, err := rs.nextRow(ctx)
892892
require.NoError(t, err)
@@ -948,7 +948,7 @@ func TestResultSetNext(t *testing.T) {
948948

949949
return part, nil
950950
}, recv, nil)
951-
require.EqualValues(t, 0, rs.index)
951+
require.EqualValues(t, 0, rs.Index())
952952
t.Run("Columns", func(t *testing.T) {
953953
require.EqualValues(t, []string{"a", "b"}, rs.Columns())
954954
})

‎internal/query/result_test.go

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,7 @@ func TestResultNextResultSet(t *testing.T) {
352352
t.Log("nextResultSet")
353353
rs, err := r.nextResultSet(ctx)
354354
require.NoError(t, err)
355-
require.EqualValues(t, 0, rs.index)
355+
require.EqualValues(t, 0, rs.Index())
356356
{
357357
t.Log("next (row=1)")
358358
_, err := rs.nextRow(ctx)
@@ -393,13 +393,13 @@ func TestResultNextResultSet(t *testing.T) {
393393
t.Log("nextResultSet")
394394
rs, err := r.nextResultSet(ctx)
395395
require.NoError(t, err)
396-
require.EqualValues(t, 1, rs.index)
396+
require.EqualValues(t, 1, rs.Index())
397397
}
398398
{
399399
t.Log("nextResultSet")
400400
rs, err := r.nextResultSet(ctx)
401401
require.NoError(t, err)
402-
require.EqualValues(t, 2, rs.index)
402+
require.EqualValues(t, 2, rs.Index())
403403
{
404404
t.Log("next (row=1)")
405405
_, err := rs.nextRow(ctx)
@@ -442,8 +442,10 @@ func TestResultNextResultSet(t *testing.T) {
442442
}
443443
{
444444
t.Log("nextResultSet")
445-
_, err := r.nextResultSet(context.Background())
445+
rs, err := r.nextResultSet(context.Background())
446446
require.ErrorIs(t, err, errClosedResult)
447+
require.Nil(t, rs)
448+
require.Equal(t, -1, rs.Index())
447449
}
448450
t.Log("check final error")
449451
require.NoError(t, r.Err())
@@ -521,7 +523,7 @@ func TestResultNextResultSet(t *testing.T) {
521523
t.Log("nextResultSet")
522524
rs, err := r.nextResultSet(ctx)
523525
require.NoError(t, err)
524-
require.EqualValues(t, 0, rs.index)
526+
require.EqualValues(t, 0, rs.Index())
525527
{
526528
t.Log("next (row=1)")
527529
_, err := rs.nextRow(ctx)
@@ -840,7 +842,7 @@ func TestResultNextResultSet(t *testing.T) {
840842
t.Log("nextResultSet")
841843
rs, err := r.nextResultSet(ctx)
842844
require.NoError(t, err)
843-
require.EqualValues(t, 0, rs.index)
845+
require.EqualValues(t, 0, rs.Index())
844846
{
845847
t.Log("next (row=1)")
846848
_, err := rs.nextRow(ctx)
@@ -881,7 +883,7 @@ func TestResultNextResultSet(t *testing.T) {
881883
t.Log("nextResultSet")
882884
rs, err := r.nextResultSet(ctx)
883885
require.NoError(t, err)
884-
require.EqualValues(t, 2, rs.index)
886+
require.EqualValues(t, 2, rs.Index())
885887
}
886888
{
887889
t.Log("nextResultSet")

‎query/result.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ type (
2020
Range(ctx context.Context) xiter.Seq2[ResultSet, error]
2121
}
2222
ResultSet interface {
23+
Index() int
2324
Columns() []string
2425
ColumnTypes() []Type
2526
NextRow(ctx context.Context) (Row, error)

‎sugar/query_test.go

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,9 +140,7 @@ func TestUnmarshallResultSet(t *testing.T) {
140140
ID uint64 `sql:"id"`
141141
Str string `sql:"myStr"`
142142
}
143-
v, err := sugar.UnmarshallResultSet[myStruct](internalQuery.NewMaterializedResultSet(
144-
nil,
145-
nil,
143+
v, err := sugar.UnmarshallResultSet[myStruct](internalQuery.NewMaterializedResultSet(-1, nil, nil,
146144
[]query.Row{
147145
func() query.Row {
148146
row, err := internalQuery.NewRow(ctx, []*Ydb.Column{

0 commit comments

Comments
 (0)