Add ResultSet.Index method

Author: Roman Golov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `ResultSet.Index()` method
+
 ## v3.74.5
 * Fixed bug with reading empty result set parts.
 * Fixed nil pointer dereference when closing result set

internal/query/result.go

@@ -25,7 +25,7 @@ var (
 type (
 	materializedResult struct {
 		resultSets []query.ResultSet
-		idx        int
+		idx        int64
 	}
 	result struct {
 		stream         Ydb_Query_V1.QueryService_ExecuteQueryClient
@@ -51,7 +51,7 @@ func (r *materializedResult) Close(ctx context.Context) error {
 }
 
 func (r *materializedResult) NextResultSet(ctx context.Context) (query.ResultSet, error) {
-	if r.idx == len(r.resultSets) {
+	if r.idx == int64(len(r.resultSets)) {
 		return nil, xerrors.WithStackTrace(io.EOF)
 	}
 

internal/query/result_set.go

@@ -85,6 +85,10 @@ func (rs *materializedResultSet) NextRow(ctx context.Context) (query.Row, error)
 	return rs.rows[rs.idx], nil
 }
 
+func (rs *materializedResultSet) Index() int {
+	return rs.idx
+}
+
 func NewMaterializedResultSet(
 	columnNames []string,
 	columnTypes []types.Type,
@@ -169,3 +173,7 @@ func (rs *resultSet) NextRow(ctx context.Context) (_ query.Row, err error) {
 
 	return rs.nextRow(ctx)
 }
+
+func (rs *resultSet) Index() int {
+	return int(rs.index)
+}

internal/query/result_set_test.go

@@ -59,7 +59,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.ErrorIs(t, err, io.EOF)
@@ -144,7 +144,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.NoError(t, err)
@@ -303,7 +303,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.NoError(t, err)
@@ -441,7 +441,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.NoError(t, err)
@@ -544,7 +544,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(childCtx)
 				require.NoError(t, err)
@@ -636,7 +636,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.NoError(t, err)
@@ -735,7 +735,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.NoError(t, err)
@@ -886,7 +886,7 @@ func TestResultSetNext(t *testing.T) {
 
 				return part, nil
 			}, recv, nil)
-			require.EqualValues(t, 0, rs.index)
+			require.EqualValues(t, 0, rs.Index())
 			{
 				_, err := rs.nextRow(ctx)
 				require.NoError(t, err)
@@ -948,7 +948,7 @@ func TestResultSetNext(t *testing.T) {
 
 		return part, nil
 	}, recv, nil)
-	require.EqualValues(t, 0, rs.index)
+	require.EqualValues(t, 0, rs.Index())
 	t.Run("Columns", func(t *testing.T) {
 		require.EqualValues(t, []string{"a", "b"}, rs.Columns())
 	})

internal/query/result_test.go

@@ -352,7 +352,7 @@ func TestResultNextResultSet(t *testing.T) {
 				t.Log("nextResultSet")
 				rs, err := r.nextResultSet(ctx)
 				require.NoError(t, err)
-				require.EqualValues(t, 0, rs.index)
+				require.EqualValues(t, 0, rs.Index())
 				{
 					t.Log("next (row=1)")
 					_, err := rs.nextRow(ctx)
@@ -393,13 +393,13 @@ func TestResultNextResultSet(t *testing.T) {
 				t.Log("nextResultSet")
 				rs, err := r.nextResultSet(ctx)
 				require.NoError(t, err)
-				require.EqualValues(t, 1, rs.index)
+				require.EqualValues(t, 1, rs.Index())
 			}
 			{
 				t.Log("nextResultSet")
 				rs, err := r.nextResultSet(ctx)
 				require.NoError(t, err)
-				require.EqualValues(t, 2, rs.index)
+				require.EqualValues(t, 2, rs.Index())
 				{
 					t.Log("next (row=1)")
 					_, err := rs.nextRow(ctx)
@@ -521,7 +521,7 @@ func TestResultNextResultSet(t *testing.T) {
 				t.Log("nextResultSet")
 				rs, err := r.nextResultSet(ctx)
 				require.NoError(t, err)
-				require.EqualValues(t, 0, rs.index)
+				require.EqualValues(t, 0, rs.Index())
 				{
 					t.Log("next (row=1)")
 					_, err := rs.nextRow(ctx)
@@ -840,7 +840,7 @@ func TestResultNextResultSet(t *testing.T) {
 				t.Log("nextResultSet")
 				rs, err := r.nextResultSet(ctx)
 				require.NoError(t, err)
-				require.EqualValues(t, 0, rs.index)
+				require.EqualValues(t, 0, rs.Index())
 				{
 					t.Log("next (row=1)")
 					_, err := rs.nextRow(ctx)
@@ -881,7 +881,7 @@ func TestResultNextResultSet(t *testing.T) {
 				t.Log("nextResultSet")
 				rs, err := r.nextResultSet(ctx)
 				require.NoError(t, err)
-				require.EqualValues(t, 2, rs.index)
+				require.EqualValues(t, 2, rs.Index())
 			}
 			{
 				t.Log("nextResultSet")

query/result.go

@@ -23,6 +23,7 @@ type (
 		Columns() []string
 		ColumnTypes() []Type
 		NextRow(ctx context.Context) (Row, error)
+		Index() int
 
 		// Range is experimental API for range iterators available with Go version 1.22+ and flag `GOEXPERIMENT=rangefunc`.
 		Range(ctx context.Context) xiter.Seq2[Row, error]