feat: support breaker with sql statements (#3936)

core/stores/sqlx/stmt.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"time"
 
+	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/syncx"
 	"github.com/zeromicro/go-zero/core/timex"
@@ -18,6 +19,137 @@ var (
 	logSlowSql    = syncx.ForAtomicBool(true)
 )
 
+type (
+	// StmtSession interface represents a session that can be used to execute statements.
+	StmtSession interface {
+		Close() error
+		Exec(args ...any) (sql.Result, error)
+		ExecCtx(ctx context.Context, args ...any) (sql.Result, error)
+		QueryRow(v any, args ...any) error
+		QueryRowCtx(ctx context.Context, v any, args ...any) error
+		QueryRowPartial(v any, args ...any) error
+		QueryRowPartialCtx(ctx context.Context, v any, args ...any) error
+		QueryRows(v any, args ...any) error
+		QueryRowsCtx(ctx context.Context, v any, args ...any) error
+		QueryRowsPartial(v any, args ...any) error
+		QueryRowsPartialCtx(ctx context.Context, v any, args ...any) error
+	}
+
+	statement struct {
+		query  string
+		stmt   *sql.Stmt
+		brk    breaker.Breaker
+		accept breaker.Acceptable
+	}
+
+	stmtConn interface {
+		Exec(args ...any) (sql.Result, error)
+		ExecContext(ctx context.Context, args ...any) (sql.Result, error)
+		Query(args ...any) (*sql.Rows, error)
+		QueryContext(ctx context.Context, args ...any) (*sql.Rows, error)
+	}
+)
+
+func (s statement) Close() error {
+	return s.stmt.Close()
+}
+
+func (s statement) Exec(args ...any) (sql.Result, error) {
+	return s.ExecCtx(context.Background(), args...)
+}
+
+func (s statement) ExecCtx(ctx context.Context, args ...any) (result sql.Result, err error) {
+	ctx, span := startSpan(ctx, "Exec")
+	defer func() {
+		endSpan(span, err)
+	}()
+
+	err = s.brk.DoWithAcceptable(func() error {
+		result, err = execStmt(ctx, s.stmt, s.query, args...)
+		return err
+	}, func(err error) bool {
+		return s.accept(err)
+	})
+
+	return
+}
+
+func (s statement) QueryRow(v any, args ...any) error {
+	return s.QueryRowCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowCtx(ctx context.Context, v any, args ...any) (err error) {
+	ctx, span := startSpan(ctx, "QueryRow")
+	defer func() {
+		endSpan(span, err)
+	}()
+
+	return s.queryRows(ctx, func(v any, scanner rowsScanner) error {
+		return unmarshalRow(v, scanner, true)
+	}, v, args...)
+}
+
+func (s statement) QueryRowPartial(v any, args ...any) error {
+	return s.QueryRowPartialCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowPartialCtx(ctx context.Context, v any, args ...any) (err error) {
+	ctx, span := startSpan(ctx, "QueryRowPartial")
+	defer func() {
+		endSpan(span, err)
+	}()
+
+	return s.queryRows(ctx, func(v any, scanner rowsScanner) error {
+		return unmarshalRow(v, scanner, false)
+	}, v, args...)
+}
+
+func (s statement) QueryRows(v any, args ...any) error {
+	return s.QueryRowsCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowsCtx(ctx context.Context, v any, args ...any) (err error) {
+	ctx, span := startSpan(ctx, "QueryRows")
+	defer func() {
+		endSpan(span, err)
+	}()
+
+	return s.queryRows(ctx, func(v any, scanner rowsScanner) error {
+		return unmarshalRows(v, scanner, true)
+	}, v, args...)
+}
+
+func (s statement) QueryRowsPartial(v any, args ...any) error {
+	return s.QueryRowsPartialCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowsPartialCtx(ctx context.Context, v any, args ...any) (err error) {
+	ctx, span := startSpan(ctx, "QueryRowsPartial")
+	defer func() {
+		endSpan(span, err)
+	}()
+
+	return s.queryRows(ctx, func(v any, scanner rowsScanner) error {
+		return unmarshalRows(v, scanner, false)
+	}, v, args...)
+}
+
+func (s statement) queryRows(ctx context.Context, scanFn func(any, rowsScanner) error,
+	v any, args ...any) error {
+	var scanFailed bool
+	return s.brk.DoWithAcceptable(func() error {
+		return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
+			err := scanFn(v, rows)
+			if err != nil {
+				scanFailed = true
+			}
+			return err
+		}, s.query, args...)
+	}, func(err error) bool {
+		return scanFailed || s.accept(err)
+	})
+}
+
 // DisableLog disables logging of sql statements, includes info and slow logs.
 func DisableLog() {
 	logSql.Set(false)