Merge pull request #857 from ydb-platform/fix-scan-wide-results

fix unwrap raw values in wide results

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed panic when try to unwrap values with more than 127 columns with custom ydb unmarshaler 
+
 ## v3.53.1
 * Bumps `github.com/ydb-platform/ydb-go-genproto` for support `query` service
 * Bumps `golang.org/x/net` from `0.7.0` to `0.17.0`

internal/table/scanner/scanner.go

@@ -1175,15 +1175,15 @@ func (x item) isEmpty() bool {
 
 type scanStack struct {
 	v        []item
-	p        int8
+	p        int
 	scanItem item
 }
 
 func (s *scanStack) size() int {
 	if !s.scanItem.isEmpty() {
 		s.set(s.scanItem)
 	}
-	return int(s.p) + 1
+	return s.p + 1
 }
 
 func (s *scanStack) get(i int) item {
@@ -1212,7 +1212,7 @@ func (s *scanStack) leave() {
 }
 
 func (s *scanStack) set(v item) {
-	if int(s.p) == len(s.v) {
+	if s.p == len(s.v) {
 		s.v = append(s.v, v)
 	} else {
 		s.v[s.p] = v

tests/integration/table_regress_test.go

@@ -0,0 +1,95 @@
+//go:build integration
+// +build integration
+
+package integration
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/table"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table/result"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table/result/indexed"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table/result/named"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
+)
+
+type rawInt64 struct {
+	val int64
+}
+
+func (r *rawInt64) UnmarshalYDB(raw types.RawValue) error {
+	raw.Unwrap()
+	r.val = raw.Int64()
+	return nil
+}
+
+func TestIssueWideResultWithUnwrap(t *testing.T) {
+	const columnsCount = 200
+	var columns []string
+	for i := 0; i < columnsCount; i++ {
+		column := fmt.Sprintf("CAST(%v AS Int64?) as c%v", i, i)
+		columns = append(columns, column)
+	}
+
+	query := "SELECT " + strings.Join(columns, ", ")
+	t.Run("named", func(t *testing.T) {
+		scope := newScope(t)
+		var res result.Result
+		var err error
+		err = scope.Driver().Table().DoTx(scope.Ctx, func(ctx context.Context, tx table.TransactionActor) error {
+			res, err = tx.Execute(ctx, query, nil)
+			return err
+		})
+		require.NoError(t, err)
+
+		res.NextResultSet(scope.Ctx)
+		require.NoError(t, res.Err())
+		res.NextRow()
+		require.NoError(t, res.Err())
+
+		results := make([]rawInt64, columnsCount)
+		resultsPointers := make([]named.Value, columnsCount)
+		for i := range results {
+			resultsPointers[i] = named.Required("c"+strconv.Itoa(i), &results[i])
+		}
+		err = res.ScanNamed(resultsPointers...)
+		require.NoError(t, err)
+
+		for i, val := range results {
+			require.Equal(t, int64(i), val.val)
+		}
+	})
+	t.Run("indexed", func(t *testing.T) {
+		scope := newScope(t)
+		var res result.Result
+		var err error
+		err = scope.Driver().Table().DoTx(scope.Ctx, func(ctx context.Context, tx table.TransactionActor) error {
+			res, err = tx.Execute(ctx, query, nil)
+			return err
+		})
+		require.NoError(t, err)
+
+		res.NextResultSet(scope.Ctx)
+		require.NoError(t, res.Err())
+		res.NextRow()
+		require.NoError(t, res.Err())
+
+		results := make([]rawInt64, columnsCount)
+		resultsPointers := make([]indexed.RequiredOrOptional, columnsCount)
+		for i := range results {
+			resultsPointers[i] = &results[i]
+		}
+		err = res.Scan(resultsPointers...)
+		require.NoError(t, err)
+
+		for i, val := range results {
+			require.Equal(t, int64(i), val.val)
+		}
+	})
+}