Merge pull request #1416 from ydb-platform/execute-idempotent

make option `query.WithIdempotent()` as option for `query.Client.{Exec,Query,QueryResultSet,QueryRow}`

Author: asmyasnikov

CHANGELOG.md

@@ -8,6 +8,7 @@
 * Removed experimental methods `query.Session.ReadResultSet` and  `query.Session.ReadRows`
 * Removed experimental methods `query.TxActor.ReadResultSet` and  `query.TxActor.ReadRows`
 * Removed experimental method `query.Client.Stats`
+* Option `query.WithIdempotent()` allowed for `query.Client.{Exec,Query,QueryResultSet,QueryRow}` methods now
 * Added experimental support for operation service client through `db.Operation()` method (supports methods `Get`, `List`, `Cancel` and `Forget`)
 
 ## v3.76.6

example_test.go

@@ -1,9 +1,10 @@
+//go:build go1.23
+
 package ydb_test
 
 import (
 	"context"
 	"database/sql"
-	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -28,61 +29,45 @@ func Example_query() {
 	ctx := context.TODO()
 	db, err := ydb.Open(ctx, "grpc://localhost:2136/local")
 	if err != nil {
-		log.Fatal(err)
+		panic(err)
 	}
 	defer db.Close(ctx) // cleanup resources
 
-	err = db.Query().Do( // Do retry operation on errors with best effort
+	materializedResult, err := db.Query().Query( // Do retry operation on errors with best effort
 		ctx, // context manage exiting from Do
-		func(ctx context.Context, s query.Session) (err error) { // retry operation
-			streamResult, err := s.Query(ctx,
-				`SELECT $id as myId, $str as myStr`,
-				query.WithParameters(
-					ydb.ParamsBuilder().
-						Param("$id").Uint64(42).
-						Param("$str").Text("my string").
-						Build(),
-				),
-			)
+		`SELECT $id as myId, $str as myStr`,
+		query.WithParameters(
+			ydb.ParamsBuilder().
+				Param("$id").Uint64(42).
+				Param("$str").Text("my string").
+				Build(),
+		),
+		query.WithIdempotent(),
+	)
+	if err != nil {
+		panic(err)
+	}
+	defer func() { _ = materializedResult.Close(ctx) }() // cleanup resources
+	for rs, err := range materializedResult.ResultSets(ctx) {
+		if err != nil {
+			panic(err)
+		}
+		for row, err := range rs.Rows(ctx) {
 			if err != nil {
-				return err // for auto-retry with driver
+				panic(err)
 			}
-			defer func() { _ = streamResult.Close(ctx) }() // cleanup resources
-			for {                                          // iterate over result sets
-				rs, err := streamResult.NextResultSet(ctx)
-				if err != nil {
-					if errors.Is(err, io.EOF) {
-						break
-					}
-
-					return err
-				}
-				for { // iterate over rows
-					row, err := rs.NextRow(ctx)
-					if err != nil {
-						if errors.Is(err, io.EOF) {
-							break
-						}
-
-						return err
-					}
-					type myStruct struct {
-						ID  uint64 `sql:"id"`
-						Str string `sql:"myStr"`
-					}
-					var s myStruct
-					if err = row.ScanStruct(&s); err != nil {
-						return err // generally scan error not retryable, return it for driver check error
-					}
-				}
+			type myStruct struct {
+				ID  uint64 `sql:"id"`
+				Str string `sql:"myStr"`
 			}
-
-			return nil
-		},
-		query.WithIdempotent(),
-	)
+			var s myStruct
+			if err = row.ScanStruct(&s); err != nil {
+				panic(err)
+			}
+		}
+	}
 	if err != nil {
-		log.Printf("unexpected error: %v", err)
+		panic(err)
 	}
 }
 

internal/query/client.go

@@ -322,6 +322,10 @@ func do(
 
 		err := op(ctx, s)
 		if err != nil {
+			if xerrors.IsOperationError(err) {
+				s.setStatus(statusClosed)
+			}
+
 			return xerrors.WithStackTrace(err)
 		}
 
@@ -411,7 +415,7 @@ func clientQueryRow(
 		}
 
 		return nil
-	})
+	}, settings.RetryOpts()...)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -441,14 +445,20 @@ func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute
 }
 
 func clientExec(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (finalErr error) {
+	settings := options.ExecuteSettings(opts...)
 	err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
-		err = s.Exec(ctx, q, opts...)
+		_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, withTrace(s.cfg.Trace()))
+		if err != nil {
+			return xerrors.WithStackTrace(err)
+		}
+
+		err = readAll(ctx, r)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
 
 		return nil
-	})
+	}, settings.RetryOpts()...)
 	if err != nil {
 		return xerrors.WithStackTrace(err)
 	}
@@ -479,8 +489,15 @@ func (c *Client) Exec(ctx context.Context, q string, opts ...options.Execute) (f
 func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (
 	r query.Result, err error,
 ) {
+	settings := options.ExecuteSettings(opts...)
 	err = do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
-		streamResult, err := s.Query(ctx, q, opts...)
+		_, streamResult, err := execute(ctx, s.id, s.queryServiceClient, q,
+			options.ExecuteSettings(opts...), withTrace(s.cfg.Trace()),
+		)
+		if err != nil {
+			return xerrors.WithStackTrace(err)
+		}
+
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
@@ -494,7 +511,7 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
 		}
 
 		return nil
-	})
+	}, settings.RetryOpts()...)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -528,24 +545,16 @@ func clientQueryResultSet(
 	err := do(ctx, pool, func(ctx context.Context, s *Session) error {
 		_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, resultOpts...)
 		if err != nil {
-			if xerrors.IsOperationError(err) {
-				s.setStatus(statusClosed)
-			}
-
 			return xerrors.WithStackTrace(err)
 		}
 
 		rs, err = readMaterializedResultSet(ctx, r)
 		if err != nil {
-			if xerrors.IsOperationError(err) {
-				s.setStatus(statusClosed)
-			}
-
 			return xerrors.WithStackTrace(err)
 		}
 
 		return nil
-	})
+	}, settings.RetryOpts()...)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}

internal/query/execute_query.go

@@ -18,6 +18,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/stats"
 )
 
@@ -29,6 +30,7 @@ type executeSettings interface {
 	Syntax() options.Syntax
 	Params() *params.Parameters
 	CallOptions() []grpc.CallOption
+	RetryOpts() []retry.Option
 }
 
 type executeScriptConfig interface {

internal/query/options/execute.go

@@ -7,6 +7,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stats"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 )
 
 var (
@@ -33,6 +34,7 @@ type (
 		statsCallback func(queryStats stats.QueryStats)
 		callOptions   []grpc.CallOption
 		txControl     *tx.Control
+		retryOptions  []retry.Option
 	}
 
 	// Execute is an interface for execute method options
@@ -53,6 +55,10 @@ type (
 	execModeOption = ExecMode
 )
 
+func (s *executeSettings) RetryOpts() []retry.Option {
+	return s.retryOptions
+}
+
 func (s *executeSettings) StatsCallback() func(stats.QueryStats) {
 	return s.statsCallback
 }

internal/query/options/retry.go

@@ -44,6 +44,10 @@ type (
 	}
 )
 
+func (opts RetryOptionsOption) applyExecuteOption(s *executeSettings) {
+	s.retryOptions = append(s.retryOptions, opts...)
+}
+
 func (s *doSettings) Trace() *trace.Query {
 	return s.trace
 }

internal/query/session.go

@@ -39,26 +39,13 @@ func (s *Session) QueryResultSet(
 		onDone(finalErr)
 	}()
 
-	settings := options.ExecuteSettings(opts...)
-
-	resultOpts := []resultOption{
-		withTrace(s.cfg.Trace()),
-	}
-	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, resultOpts...)
+	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, options.ExecuteSettings(opts...), withTrace(s.cfg.Trace()))
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	rs, err = readResultSet(ctx, r)
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -70,19 +57,11 @@ func (s *Session) queryRow(
 ) (row query.Row, finalErr error) {
 	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, resultOpts...)
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	row, err = readRow(ctx, r)
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -98,10 +77,6 @@ func (s *Session) QueryRow(ctx context.Context, q string, opts ...options.Execut
 
 	row, err := s.queryRow(ctx, q, options.ExecuteSettings(opts...), withTrace(s.cfg.Trace()))
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return nil, xerrors.WithStackTrace(err)
 	}
 
@@ -316,27 +291,13 @@ func (s *Session) Exec(
 		onDone(finalErr)
 	}()
 
-	settings := options.ExecuteSettings(opts...)
-
-	resultOpts := []resultOption{
-		withTrace(s.cfg.Trace()),
-	}
-
-	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, resultOpts...)
+	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, options.ExecuteSettings(opts...), withTrace(s.cfg.Trace()))
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return xerrors.WithStackTrace(err)
 	}
 
 	err = readAll(ctx, r)
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return xerrors.WithStackTrace(err)
 	}
 
@@ -352,17 +313,8 @@ func (s *Session) Query(
 		onDone(finalErr)
 	}()
 
-	settings := options.ExecuteSettings(opts...)
-
-	resultOpts := []resultOption{
-		withTrace(s.cfg.Trace()),
-	}
-	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, resultOpts...)
+	_, r, err := execute(ctx, s.id, s.queryServiceClient, q, options.ExecuteSettings(opts...), withTrace(s.cfg.Trace()))
 	if err != nil {
-		if xerrors.IsOperationError(err) {
-			s.setStatus(statusClosed)
-		}
-
 		return nil, xerrors.WithStackTrace(err)
 	}