internal/xsql/tx.go - master - githubmirror/ydb-go-sdk

internal/xsql/tx.go
160 строк4 KB
Aleksey Myasnikov
Moved the trace-generated code in `trace` package from the public API to the internal packages (#2172)
24 май 2026, 16:46
24 май 2026, 16:4686fc0c2
package xsql

import (
	"context"
	"database/sql/driver"

	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
	"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/internal/xsql/common"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/gtrace"
)

type Tx struct {
	conn *Conn
	tx   common.Tx
	ctx  context.Context //nolint:containedctx
}

func (tx *Tx) ID() string {
	if tx.tx == nil {
		return ""
	}

	return tx.tx.ID()
}

var (
	_ driver.Tx             = &Tx{}
	_ driver.ExecerContext  = &Tx{}
	_ driver.QueryerContext = &Tx{}
)

func (tx *Tx) Commit() (finalErr error) {
	defer func() {
		tx.conn.currentTx = nil
	}()

	var (
		ctx    = tx.ctx
		onDone = gtrace.DatabaseSQLOnTxCommit(tx.conn.connector.Trace(), &ctx,
			stack.FunctionID("database/sql.(*Tx).Commit" /*stack.Package("database/sql")*/),
			tx,
		)
	)
	defer func() {
		onDone(finalErr)
	}()

	if err := tx.tx.Commit(tx.ctx); err != nil {
		return xerrors.WithStackTrace(badconn.Map(err))
	}

	return nil
}

func (tx *Tx) Rollback() (finalErr error) {
	defer func() {
		tx.conn.currentTx = nil
	}()

	var (
		ctx    = tx.ctx
		onDone = gtrace.DatabaseSQLOnTxRollback(tx.conn.connector.Trace(), &ctx,
			stack.FunctionID("database/sql.(*Tx).Rollback" /*stack.Package("database/sql")*/),
			tx,
		)
	)
	defer func() {
		onDone(finalErr)
	}()

	err := tx.tx.Rollback(tx.ctx)
	if err != nil {
		return xerrors.WithStackTrace(badconn.Map(err))
	}

	return err
}

func (tx *Tx) QueryContext(ctx context.Context, sql string, args []driver.NamedValue) (
	_ driver.Rows, finalErr error,
) {
	onDone := gtrace.DatabaseSQLOnTxQuery(tx.conn.connector.Trace(), &ctx,
		stack.FunctionID("database/sql.(*Tx).QueryContext" /*stack.Package("database/sql")*/),
		tx.ctx, tx, sql,
	)
	defer func() {
		onDone(finalErr)
	}()

	sql, params, err := tx.conn.toYdb(sql, args...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	if isExplain(ctx) {
		ast, plan, err := tx.conn.cc.Explain(ctx, sql, params)
		if err != nil {
			return nil, xerrors.WithStackTrace(badconn.Map(err))
		}

		return newRows(ctx, rowByAstPlan(ast, plan)), nil
	}

	rows, err := tx.tx.Query(ctx, sql, params)
	if err != nil {
		return nil, xerrors.WithStackTrace(badconn.Map(err))
	}

	return newRows(ctx, rows), nil
}

func (tx *Tx) ExecContext(ctx context.Context, sql string, args []driver.NamedValue) (
	_ driver.Result, finalErr error,
) {
	onDone := gtrace.DatabaseSQLOnTxExec(tx.conn.connector.Trace(), &ctx,
		stack.FunctionID("database/sql.(*Tx).ExecContext" /*stack.Package("database/sql")*/),
		tx.ctx, tx, sql,
	)
	defer func() {
		onDone(finalErr)
	}()

	sql, params, err := tx.conn.toYdb(sql, args...)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	result, err := tx.tx.Exec(ctx, sql, params)
	if err != nil {
		return nil, xerrors.WithStackTrace(badconn.Map(err))
	}

	return result, nil
}

func (tx *Tx) PrepareContext(ctx context.Context, sql string) (_ driver.Stmt, finalErr error) {
	onDone := gtrace.DatabaseSQLOnTxPrepare(tx.conn.connector.Trace(), &ctx,
		stack.FunctionID("database/sql.(*Tx).PrepareContext" /*stack.Package("database/sql")*/),
		tx.ctx, tx, sql,
	)
	defer func() {
		onDone(finalErr)
	}()
	if !tx.conn.cc.IsValid() {
		return nil, badconn.Map(xerrors.WithStackTrace(xerrors.Retryable(errNotReadyConn,
			xerrors.Invalid(tx),
			xerrors.Invalid(tx.conn),
			xerrors.Invalid(tx.conn.cc),
		)))
	}

	return &Stmt{
		conn:      tx.conn,
		processor: tx.tx,
		ctx:       ctx,
		sql:       sql,
	}, nil
}