Merge pull request #570 from ydb-platform/Kikimr17104

* Added wrapping of errors from unary and stream results

Author: asmyasnikov

CHANGELOG.md

@@ -5,6 +5,10 @@
 * Refactored of applying grpc dial options with defaults
 * Added `trace.Driver.{OnBalancerDialEntrypoint,OnBalancerClusterDiscoveryAttempt}` trace events
 * Fixed compilation of package `internal/xresolver` with `google.golang.org/[email protected]`
+* Fixed returning `io.EOF` on `rows.Next` and `rows.NextResultSet`
+* Added wrapping of errors from unary and stream results
+* Added error throw on `database/sql.Conn.BeginTx()`, `*sql.Tx.ExecContext` and `*sql.Tx.QueryContext` if query mode is not `ydb.DataQueryMode`
+* Added test for `database/sql` scan-query
 
 ## v3.42.8
 * Fixed `internal/scheme/helpers/IsDirectoryExists(..)` recursive bug

internal/table/scanner/result.go

@@ -33,13 +33,31 @@ type streamResult struct {
 	close func(error) error
 }
 
+// Err returns error caused Scanner to be broken.
+func (r *streamResult) Err() error {
+	err := r.scanner.Err()
+	if err != nil {
+		return xerrors.WithStackTrace(err)
+	}
+	return nil
+}
+
 type unaryResult struct {
 	baseResult
 
 	sets    []*Ydb.ResultSet
 	nextSet int
 }
 
+// Err returns error caused Scanner to be broken.
+func (r *unaryResult) Err() error {
+	err := r.scanner.Err()
+	if err != nil {
+		return xerrors.WithStackTrace(err)
+	}
+	return nil
+}
+
 // Close closes the result, preventing further iteration.
 func (r *unaryResult) Close() error {
 	if atomic.CompareAndSwapUint32(&r.closed, 0, 1) {

internal/xsql/conn.go

@@ -342,6 +342,10 @@ func (c *conn) BeginTx(ctx context.Context, txOptions driver.TxOptions) (_ drive
 	defer func() {
 		onDone(transaction, err)
 	}()
+	m := queryModeFromContext(ctx, c.defaultQueryMode)
+	if m != DataQueryMode {
+		return nil, badconn.Map(xerrors.WithStackTrace(fmt.Errorf("wrong query mode: %s", m.String())))
+	}
 	if !c.isReady() {
 		return nil, badconn.Map(xerrors.WithStackTrace(errNotReadyConn))
 	}

internal/xsql/rows.go

@@ -67,16 +67,32 @@ func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
 	return yqlTypes[index]
 }
 
-func (r *rows) NextResultSet() error {
+func (r *rows) NextResultSet() (err error) {
+	defer func() {
+		if err != nil && xerrors.Is(err, io.EOF) {
+			// database/sql checks io.EOF with "==", not errors.Is(err, io.EOF)
+			err = io.EOF
+		}
+	}()
 	r.nextSet.Do(func() {})
-	return r.result.NextResultSetErr(context.Background())
+	err = r.result.NextResultSetErr(context.Background())
+	if err != nil {
+		return badconn.Map(xerrors.WithStackTrace(err))
+	}
+	return nil
 }
 
 func (r *rows) HasNextResultSet() bool {
 	return r.result.HasNextResultSet()
 }
 
 func (r *rows) Next(dst []driver.Value) (err error) {
+	defer func() {
+		if err != nil && xerrors.Is(err, io.EOF) {
+			// database/sql checks io.EOF with "==", not errors.Is(err, io.EOF)
+			err = io.EOF
+		}
+	}()
 	r.nextSet.Do(func() {
 		err = r.result.NextResultSetErr(context.Background())
 	})

internal/xsql/tx.go

@@ -3,6 +3,7 @@ package xsql
 import (
 	"context"
 	"database/sql/driver"
+	"fmt"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
@@ -70,6 +71,10 @@ func (tx *tx) QueryContext(ctx context.Context, query string, args []driver.Name
 	defer func() {
 		onDone(err)
 	}()
+	m := queryModeFromContext(ctx, tx.conn.defaultQueryMode)
+	if m != DataQueryMode {
+		return nil, badconn.Map(xerrors.WithStackTrace(fmt.Errorf("wrong query mode: %s", m.String())))
+	}
 	var res result.Result
 	res, err = tx.tx.Execute(ctx,
 		query,
@@ -93,6 +98,10 @@ func (tx *tx) ExecContext(ctx context.Context, query string, args []driver.Named
 	defer func() {
 		onDone(err)
 	}()
+	m := queryModeFromContext(ctx, tx.conn.defaultQueryMode)
+	if m != DataQueryMode {
+		return nil, badconn.Map(xerrors.WithStackTrace(fmt.Errorf("wrong query mode: %s", m.String())))
+	}
 	_, err = tx.tx.Execute(ctx,
 		query,
 		toQueryParams(args),

tests/database_sql_regression_test.go

@@ -7,16 +7,24 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"fmt"
 	"math/rand"
+	"net/url"
 	"os"
+	"strconv"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
+	"github.com/ydb-platform/ydb-go-sdk/v3/sugar"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table"
+	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
 )
 
 func TestRegressionCloud109307(t *testing.T) {
@@ -71,3 +79,166 @@ func TestRegressionCloud109307(t *testing.T) {
 		}
 	}
 }
+
+func TestRegressionKikimr17104(t *testing.T) {
+	tablePath := "/database/sql/kikimr/17104/stream_query"
+	if dsn, has := os.LookupEnv("YDB_CONNECTION_STRING"); !has {
+		t.Errorf("expected YDB_CONNECTION_STRING environment variable")
+	} else {
+		u, err := url.Parse(dsn)
+		require.NoError(t, err)
+		tablePath = u.Path + tablePath
+	}
+
+	var (
+		upsertRowsCount = 100000
+		upsertChecksum  uint64
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 42*time.Second)
+	defer cancel()
+
+	t.Run("data", func(t *testing.T) {
+		t.Run("prepare", func(t *testing.T) {
+			var db *sql.DB
+			defer func() {
+				if db != nil {
+					_ = db.Close()
+				}
+			}()
+			t.Run("connect", func(t *testing.T) {
+				var err error
+				db, err = sql.Open("ydb", os.Getenv("YDB_CONNECTION_STRING"))
+				require.NoError(t, err)
+			})
+			t.Run("scheme", func(t *testing.T) {
+				var cc ydb.Connection
+				t.Run("unwrap", func(t *testing.T) {
+					var err error
+					cc, err = ydb.Unwrap(db)
+					require.NoError(t, err)
+				})
+				var tableExists bool
+				t.Run("check_exists", func(t *testing.T) {
+					var err error
+					tableExists, err = sugar.IsTableExists(ctx, cc.Scheme(), tablePath)
+					require.NoError(t, err)
+				})
+				if tableExists {
+					t.Run("drop", func(t *testing.T) {
+						err := retry.Do(ydb.WithQueryMode(ctx, ydb.SchemeQueryMode), db,
+							func(ctx context.Context, cc *sql.Conn) (err error) {
+								_, err = cc.ExecContext(ctx,
+									fmt.Sprintf("DROP TABLE `%s`", tablePath),
+								)
+								if err != nil {
+									return err
+								}
+								return nil
+							}, retry.WithDoRetryOptions(retry.WithIdempotent(true)),
+						)
+						require.NoError(t, err)
+					})
+				}
+				t.Run("create", func(t *testing.T) {
+					err := retry.Do(ydb.WithQueryMode(ctx, ydb.SchemeQueryMode), db,
+						func(ctx context.Context, cc *sql.Conn) (err error) {
+							_, err = cc.ExecContext(ctx,
+								fmt.Sprintf("CREATE TABLE `%s` (val Int32, PRIMARY KEY (val))", tablePath),
+							)
+							if err != nil {
+								return err
+							}
+							return nil
+						}, retry.WithDoRetryOptions(retry.WithIdempotent(true)),
+					)
+					require.NoError(t, err)
+				})
+			})
+			t.Run("upsert", func(t *testing.T) {
+				if v, ok := os.LookupEnv("UPSERT_ROWS_COUNT"); ok {
+					var vv int
+					vv, err := strconv.Atoi(v)
+					require.NoError(t, err)
+					upsertRowsCount = vv
+				}
+				// - upsert data
+				fmt.Printf("> preparing values to upsert...\n")
+				values := make([]types.Value, 0, upsertRowsCount)
+				for i := 0; i < upsertRowsCount; i++ {
+					upsertChecksum += uint64(i)
+					values = append(values,
+						types.StructValue(
+							types.StructFieldValue("val", types.Int32Value(int32(i))),
+						),
+					)
+				}
+				fmt.Printf("> upsert data\n")
+				err := retry.Do(ydb.WithQueryMode(ctx, ydb.DataQueryMode), db,
+					func(ctx context.Context, cc *sql.Conn) (err error) {
+						values := table.NewQueryParameters(table.ValueParam("$values", types.ListValue(values...)))
+						declares, err := sugar.GenerateDeclareSection(values)
+						require.NoError(t, err)
+						_, err = cc.ExecContext(ctx,
+							declares+fmt.Sprintf("UPSERT INTO `%s` SELECT val FROM AS_TABLE($values);", tablePath),
+							values,
+						)
+						if err != nil {
+							return err
+						}
+						return nil
+					}, retry.WithDoRetryOptions(retry.WithIdempotent(true)),
+				)
+				require.NoError(t, err)
+			})
+		})
+		t.Run("scan", func(t *testing.T) {
+			var db *sql.DB
+			defer func() {
+				if db != nil {
+					_ = db.Close()
+				}
+			}()
+			t.Run("connect", func(t *testing.T) {
+				var err error
+				cc, err := ydb.Open(ctx, os.Getenv("YDB_CONNECTION_STRING"))
+				require.NoError(t, err)
+				connector, err := ydb.Connector(cc, ydb.WithDefaultQueryMode(ydb.ScanQueryMode))
+				require.NoError(t, err)
+				db = sql.OpenDB(connector)
+			})
+			t.Run("query", func(t *testing.T) {
+				var (
+					rowsCount int
+					checkSum  uint64
+				)
+				err := retry.Do(ydb.WithQueryMode(ctx, ydb.ScanQueryMode), db,
+					func(ctx context.Context, cc *sql.Conn) (err error) {
+						var rows *sql.Rows
+						rowsCount = 0
+						checkSum = 0
+						rows, err = cc.QueryContext(ctx, fmt.Sprintf("SELECT val FROM `%s`", tablePath))
+						if err != nil {
+							return err
+						}
+						for rows.NextResultSet() {
+							for rows.Next() {
+								rowsCount++
+								var val uint64
+								err = rows.Scan(&val)
+								if err != nil {
+									return err
+								}
+								checkSum += val
+							}
+						}
+						return rows.Err()
+					}, retry.WithDoRetryOptions(retry.WithIdempotent(true)),
+				)
+				require.NoError(t, err)
+				require.Equal(t, upsertRowsCount, rowsCount)
+				require.Equal(t, upsertChecksum, checkSum)
+			})
+		})
+	})
+}

tests/table_test.go

@@ -646,7 +646,7 @@ func testTable(t testing.TB) {
 	); err != nil {
 		t.Fatalf("create table failed: %v\n", err)
 	}
-	fmt.Printf("> table stream_query openSessions\n")
+	fmt.Printf("> table stream_query upsert data\n")
 	var (
 		upsertRowsCount = 100000
 		sum             uint64