refactored retries and open template

Author: LostImagin4tion

dialect/dialect.go

@@ -88,6 +88,11 @@ func DebugWithContext(d Driver, logger func(context.Context, ...any)) Driver {
 	return drv
 }
 
+// Log returns the log function used by this DebugDriver.
+func (d *DebugDriver) Log() func(context.Context, ...any) {
+	return d.log
+}
+
 // Exec logs its params and calls the underlying driver Exec method.
 func (d *DebugDriver) Exec(ctx context.Context, query string, args, v any) error {
 	d.log(ctx, fmt.Sprintf("driver.Exec: query=%v args=%v", query, args))

dialect/sql/driver.go

@@ -14,31 +14,75 @@ import (
 	"strings"
 
 	"entgo.io/ent/dialect"
+	"github.com/ydb-platform/ydb-go-sdk/v3"
 )
 
 // Driver is a dialect.Driver implementation for SQL based databases.
 type Driver struct {
 	Conn
-	dialect string
+	dialect       string
+	retryExecutor RetryExecutor
 }
 
 // NewDriver creates a new Driver with the given Conn and dialect.
-func NewDriver(dialect string, c Conn) *Driver {
-	return &Driver{dialect: dialect, Conn: c}
+func NewDriver(
+	dialect string,
+	c Conn,
+	retryExecutor RetryExecutor,
+) *Driver {
+	return &Driver{
+		Conn:          c,
+		dialect:       dialect,
+		retryExecutor: retryExecutor,
+	}
 }
 
 // Open wraps the database/sql.Open method and returns a dialect.Driver that implements the an ent/dialect.Driver interface.
-func Open(dialect, source string) (*Driver, error) {
-	db, err := sql.Open(dialect, source)
+func Open(sqlDialect, dsn string) (*Driver, error) {
+	var (
+		db  *sql.DB
+		err error
+	)
+
+	if sqlDialect == dialect.YDB {
+		nativeDriver, err := ydb.Open(context.Background(), dsn)
+		if err != nil {
+			return nil, err
+		}
+
+		conn, err := ydb.Connector(
+			nativeDriver,
+			ydb.WithAutoDeclare(),
+			ydb.WithTablePathPrefix(nativeDriver.Name()),
+			ydb.WithQueryService(true),
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		db = sql.OpenDB(conn)
+	} else {
+		db, err = sql.Open(sqlDialect, dsn)
+	}
+
 	if err != nil {
 		return nil, err
 	}
-	return NewDriver(dialect, Conn{db, dialect}), nil
+
+	return NewDriver(
+		sqlDialect,
+		Conn{db, sqlDialect},
+		NewRetryExecutor(sqlDialect, db),
+	), nil
 }
 
 // OpenDB wraps the given database/sql.DB method with a Driver.
-func OpenDB(dialect string, db *sql.DB) *Driver {
-	return NewDriver(dialect, Conn{db, dialect})
+func OpenDB(sqlDialect string, db *sql.DB) *Driver {
+	return NewDriver(
+		sqlDialect,
+		Conn{db, sqlDialect},
+		NewRetryExecutor(sqlDialect, db),
+	)
 }
 
 // DB returns the underlying *sql.DB instance.
@@ -74,6 +118,10 @@ func (d *Driver) BeginTx(ctx context.Context, opts *TxOptions) (dialect.Tx, erro
 	}, nil
 }
 
+func (d *Driver) RetryExecutor() RetryExecutor {
+	return d.retryExecutor
+}
+
 // Close closes the underlying connection.
 func (d *Driver) Close() error { return d.DB().Close() }
 

dialect/sql/retry.go

@@ -0,0 +1,236 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+
+	"entgo.io/ent/dialect"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
+)
+
+// RetryExecutor is an interface for database operations with automatic retries.
+type RetryExecutor interface {
+	// Do executes the given function within a retry loop without a transaction.
+	// The function receives a dialect.Driver that wraps the connection.
+	// opts are driver-specific retry options (e.g., ydb retry.Option).
+	Do(
+		ctx context.Context,
+		fn func(ctx context.Context, drv dialect.Driver) error,
+		opts ...any,
+	) error
+
+	// DoTx executes the given function within a retry loop with a transaction.
+	// The function receives a dialect.Driver that wraps the database/sql.Tx transaction.
+	// opts are driver-specific retry options (e.g., ydb retry.Option).
+	DoTx(
+		ctx context.Context,
+		fn func(ctx context.Context, drv dialect.Driver) error,
+		opts ...any,
+	) error
+}
+
+// NewRetryExecutor creates a new RetryExecutor with the given database connection
+func NewRetryExecutor(
+	sqlDialect string,
+	db *sql.DB,
+) RetryExecutor {
+	if sqlDialect == dialect.YDB && db != nil {
+		return &YDBRetryExecutor{db: db}
+	} else {
+		return nil
+	}
+}
+
+// RetryExecutorGetter is an optional interface that drivers can implement to provide
+// a RetryExecutor for automatic retry handling.
+// If a driver implements this interface,
+// sqlgraph operations will use the RetryExecutor for database operations.
+type RetryExecutorGetter interface {
+	// RetryExecutor returns the RetryExecutor for this driver.
+	// If nil is returned, no retry handling will be applied.
+	RetryExecutor() RetryExecutor
+}
+
+// GetRetryExecutor returns the RetryExecutor for the given driver if available.
+// If the driver is wrapped with a DebugDriver, the returned executor will preserve
+// debug logging by wrapping the driver passed to callback functions.
+func GetRetryExecutor(drv dialect.Driver) RetryExecutor {
+	var logFn func(context.Context, ...any)
+	if dd, ok := drv.(*dialect.DebugDriver); ok {
+		logFn = dd.Log()
+		drv = dd.Driver
+	}
+	if getter, ok := drv.(RetryExecutorGetter); ok {
+		executor := getter.RetryExecutor()
+		if executor == nil {
+			return nil
+		}
+		if logFn != nil {
+			return &debugRetryExecutor{
+				RetryExecutor: executor,
+				log:           logFn,
+			}
+		}
+		return executor
+	}
+	return nil
+}
+
+// debugRetryExecutor wraps a RetryExecutor to preserve debug logging.
+type debugRetryExecutor struct {
+	RetryExecutor
+	log func(context.Context, ...any)
+}
+
+// Do executes the operation with debug logging preserved.
+func (d *debugRetryExecutor) Do(
+	ctx context.Context,
+	fn func(ctx context.Context, drv dialect.Driver) error,
+	opts ...any,
+) error {
+	return d.RetryExecutor.Do(
+		ctx,
+		func(ctx context.Context, drv dialect.Driver) error {
+			return fn(ctx, dialect.DebugWithContext(drv, d.log))
+		},
+		opts...,
+	)
+}
+
+// DoTx executes the operation within a transaction with debug logging preserved.
+func (d *debugRetryExecutor) DoTx(
+	ctx context.Context,
+	fn func(ctx context.Context, drv dialect.Driver) error,
+	opts ...any,
+) error {
+	return d.RetryExecutor.DoTx(
+		ctx,
+		func(ctx context.Context, drv dialect.Driver) error {
+			return fn(ctx, dialect.DebugWithContext(drv, d.log))
+		},
+		opts...,
+	)
+}
+
+// RetryConfig holds retry configuration for sqlgraph operations.
+// This is used to pass retry options to the RetryExecutor.
+type RetryConfig struct {
+	// Options are driver-specific retry options.
+	Options []any
+}
+
+// YDBRetryExecutor implements sqlgraph.YDBRetryExecutor for YDB
+type YDBRetryExecutor struct {
+	db *sql.DB
+}
+
+// Do executes a read-only operation with retry support.
+// It uses ydb-go-sdk's retry.Do which handles YDB-specific retryable errors.
+// Options should be created using retry.WithIdempotent(), retry.WithLabel(), etc.
+func (r *YDBRetryExecutor) Do(
+	ctx context.Context,
+	fn func(ctx context.Context, drv dialect.Driver) error,
+	opts ...any,
+) error {
+	err := retry.Do(
+		ctx,
+		r.db,
+		func(ctx context.Context, conn *sql.Conn) error {
+			return fn(ctx, newConnRetryDriver(conn))
+		},
+		retry.WithDoRetryOptions(toRetryOptions(opts)...),
+	)
+	return unwrapYDBError(err)
+}
+
+// DoTx executes the operation within a transaction with retry support.
+// It uses ydb-go-sdk's retry.DoTx which handles YDB-specific retryable errors.
+// Options should be created using retry.WithIdempotent(), retry.WithLabel(), etc.
+func (r *YDBRetryExecutor) DoTx(
+	ctx context.Context,
+	fn func(ctx context.Context, drv dialect.Driver) error,
+	opts ...any,
+) error {
+	err := retry.DoTx(
+		ctx,
+		r.db,
+		func(ctx context.Context, tx *sql.Tx) error {
+			return fn(ctx, newTxRetryDriver(tx))
+		},
+		retry.WithDoTxRetryOptions(toRetryOptions(opts)...),
+	)
+	return unwrapYDBError(err)
+}
+
+// unwrapYDBError extracts the original error from YDB's error wrapping.
+// YDB SDK wraps errors with stack traces and retry context, which changes
+// the error message.
+func unwrapYDBError(err error) error {
+	if err == nil {
+		return nil
+	}
+	original := err
+	for {
+		unwrapped := errors.Unwrap(original)
+		if unwrapped == nil {
+			break
+		}
+		original = unwrapped
+	}
+	return original
+}
+
+// toRetryOptions converts a slice of any options to retry.Option slice
+func toRetryOptions(opts []any) []retry.Option {
+	retryOpts := make([]retry.Option, 0, len(opts))
+	for _, opt := range opts {
+		if ro, ok := opt.(retry.Option); ok {
+			retryOpts = append(retryOpts, ro)
+		}
+	}
+	return retryOpts
+}
+
+// retryDriver is designed for use only in sqlgraph,
+// specifically - in retry.DoTx callbacks
+type retryDriver struct {
+	Conn
+}
+
+var _ dialect.Driver = (*retryDriver)(nil)
+
+// newConnRetryDriver creates a new RetryDriver from a database connection.
+func newConnRetryDriver(conn *sql.Conn) *retryDriver {
+	return &retryDriver{
+		Conn: Conn{ExecQuerier: conn},
+	}
+}
+
+// newTxRetryDriver creates a new RetryDriver from a transaction.
+func newTxRetryDriver(tx *sql.Tx) *retryDriver {
+	return &retryDriver{
+		Conn: Conn{ExecQuerier: tx},
+	}
+}
+
+// sqlgraph creates nested transactions in several methods.
+// But YDB doesnt support nested transactions.
+// Therefore, this methods returns no-op tx
+func (d *retryDriver) Tx(ctx context.Context) (dialect.Tx, error) {
+	return dialect.NopTx(d), nil
+}
+
+// Close is a no-op for RetryDriver since retry.DoTx manages the transaction lifecycle.
+func (d *retryDriver) Close() error {
+	return nil
+}
+
+// Dialect returns the YDB dialect name.
+func (d *retryDriver) Dialect() string {
+	return dialect.YDB
+}