getting rowsAffected from YDB QueryStats (#1887)

Co-authored-by: Aleksey Myasnikov <[email protected]>

Author: kprokopenko

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added support for `Result.RowsAffected()` for YDB `database/sql` driver
+* Upgraded minimal version of Go to 1.23.9
 * Fixed race in `readerReconnector`
 
 ## v3.117.1

examples/go.mod

@@ -1,8 +1,6 @@
 module examples
 
-go 1.23.0
-
-toolchain go1.23.6
+go 1.23.9
 
 require (
 	github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40

go.mod

@@ -1,6 +1,6 @@
 module github.com/ydb-platform/ydb-go-sdk/v3
 
-go 1.22.5
+go 1.23.9
 
 require (
 	github.com/golang-jwt/jwt/v4 v4.5.2

internal/query/result.go

@@ -159,10 +159,6 @@ func newResult(
 
 		r.lastPart = part
 
-		if part.GetExecStats() != nil && r.statsCallback != nil {
-			r.statsCallback(stats.FromQueryStats(part.GetExecStats()))
-		}
-
 		return &r, nil
 	}
 }
@@ -211,6 +207,10 @@ func (r *streamResult) nextPart(ctx context.Context) (
 			}
 		}
 
+		if part.GetExecStats() != nil && r.statsCallback != nil {
+			r.statsCallback(stats.FromQueryStats(part.GetExecStats()))
+		}
+
 		return part, nil
 	}
 }
@@ -286,9 +286,6 @@ func (r *streamResult) nextResultSet(ctx context.Context) (_ *resultSet, err err
 			if err != nil {
 				return nil, xerrors.WithStackTrace(err)
 			}
-			if part.GetExecStats() != nil && r.statsCallback != nil {
-				r.statsCallback(stats.FromQueryStats(part.GetExecStats()))
-			}
 			if part.GetResultSetIndex() < r.resultSetIndex {
 				r.closer.Close(nil)
 
@@ -326,9 +323,6 @@ func (r *streamResult) nextPartFunc(
 				return nil, xerrors.WithStackTrace(err)
 			}
 			r.lastPart = part
-			if part.GetExecStats() != nil && r.statsCallback != nil {
-				r.statsCallback(stats.FromQueryStats(part.GetExecStats()))
-			}
 			if part.GetResultSetIndex() > nextResultSetIndex {
 				return nil, xerrors.WithStackTrace(fmt.Errorf(
 					"result set (index=%d) receive part (index=%d) for next result set: %w (%w)",

internal/xsql/xquery/conn.go

@@ -15,13 +15,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/common"
 )
 
-type resultNoRows struct{}
-
-func (resultNoRows) LastInsertId() (int64, error) { return 0, ErrUnsupported }
-func (resultNoRows) RowsAffected() (int64, error) { return 0, ErrUnsupported }
-
-var _ driver.Result = resultNoRows{}
-
 type Parent interface {
 	Query() *query.Client
 }
@@ -39,7 +32,7 @@ func (c *Conn) NodeID() uint32 {
 }
 
 func (c *Conn) Exec(ctx context.Context, sql string, params *params.Params) (
-	result driver.Result, finalErr error,
+	driver.Result, error,
 ) {
 	if !c.IsValid() {
 		return nil, xerrors.WithStackTrace(xerrors.Retryable(errNotReadyConn,
@@ -63,12 +56,15 @@ func (c *Conn) Exec(ctx context.Context, sql string, params *params.Params) (
 		opts = append(opts, options.WithTxControl(txControl))
 	}
 
+	r := &resultWithStats{}
+	opts = append(opts, options.WithStatsMode(options.StatsModeBasic, r.onQueryStats))
+
 	err := c.session.Exec(ctx, sql, opts...)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return resultNoRows{}, nil
+	return r, nil
 }
 
 func (c *Conn) Query(ctx context.Context, sql string, params *params.Params) (

internal/xsql/xquery/results_with_stats.go

@@ -0,0 +1,36 @@
+package xquery
+
+import (
+	"database/sql/driver"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+)
+
+type (
+	resultWithStats struct {
+		rowsAffected *uint64
+	}
+)
+
+var _ driver.Result = &resultWithStats{}
+
+func (r *resultWithStats) onQueryStats(qs query.Stats) {
+	var rowsAffected uint64
+	for queryPhase := range qs.QueryPhases() {
+		for tableAccess := range queryPhase.TableAccess() {
+			rowsAffected += tableAccess.Deletes.Rows + tableAccess.Updates.Rows
+		}
+	}
+	// last stats always contains the full stats of query
+	r.rowsAffected = &rowsAffected
+}
+
+func (r *resultWithStats) RowsAffected() (int64, error) {
+	if r.rowsAffected == nil {
+		return 0, ErrUnsupported
+	}
+
+	return int64(*r.rowsAffected), nil
+}
+
+func (r *resultWithStats) LastInsertId() (int64, error) { return 0, ErrUnsupported }