refactoring + fixes

internal/tx/id.go

@@ -4,6 +4,7 @@ var _ Identifier = LazyID{}
 
 const (
 	LazyTxID = "LAZY_TX"
+	FakeTxID = "FAKE_TX"
 )
 
 type (

internal/xsql/conn.go

@@ -33,7 +33,8 @@ var (
 
 type (
 	connWrapper struct {
-		cc conn.Conn
+		cc        conn.Conn
+		currentTx *txWrapper
 
 		connector *Connector
 		lastUsage xsync.LastUsage
@@ -54,16 +55,22 @@ func (c *connWrapper) CheckNamedValue(value *driver.NamedValue) error {
 }
 
 func (c *connWrapper) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	if c.currentTx != nil {
+		return nil, xerrors.WithStackTrace(xerrors.AlreadyHasTx(c.currentTx.ID()))
+	}
+
 	tx, err := c.cc.BeginTx(ctx, opts)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return &txWrapper{
+	c.currentTx = &txWrapper{
 		conn: c,
 		ctx:  ctx,
 		tx:   tx,
-	}, nil
+	}
+
+	return c.currentTx, nil
 }
 
 func (c *connWrapper) Close() error {
@@ -179,6 +186,10 @@ func (c *connWrapper) QueryContext(ctx context.Context, sql string, args []drive
 		return rowByAstPlan(ast, plan), nil
 	}
 
+	if c.currentTx != nil {
+		return c.currentTx.tx.Query(ctx, sql, params)
+	}
+
 	return c.cc.Query(ctx, sql, params)
 }
 
@@ -191,6 +202,10 @@ func (c *connWrapper) ExecContext(ctx context.Context, sql string, args []driver
 		return nil, xerrors.WithStackTrace(err)
 	}
 
+	if c.currentTx != nil {
+		return c.currentTx.tx.Exec(ctx, sql, params)
+	}
+
 	return c.cc.Exec(ctx, sql, params)
 }
 

internal/xsql/conn/query/conn.go

@@ -32,6 +32,7 @@ type Conn struct {
 	session *query.Session
 	onClose []func()
 	closed  atomic.Bool
+	fakeTx  bool
 }
 
 func (c *Conn) Exec(ctx context.Context, sql string, params *params.Params) (
@@ -114,6 +115,10 @@ func (c *Conn) isReady() bool {
 }
 
 func (c *Conn) beginTx(ctx context.Context, txOptions driver.TxOptions) (tx conn.Tx, finalErr error) {
+	if c.fakeTx {
+		return beginTxFake(ctx, c), nil
+	}
+
 	tx, err := beginTx(ctx, c, txOptions)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)

internal/xsql/conn/query/options.go

@@ -7,3 +7,9 @@ func WithOnClose(onClose func()) Option {
 		c.onClose = append(c.onClose, onClose)
 	}
 }
+
+func WithFakeTx() Option {
+	return func(c *Conn) {
+		c.fakeTx = true
+	}
+}

internal/xsql/conn/query/tx_fake.go

@@ -0,0 +1,62 @@
+package query
+
+import (
+	"context"
+	"database/sql/driver"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn/table/badconn"
+)
+
+type txFake struct {
+	conn *Conn
+	ctx  context.Context //nolint:containedctx
+}
+
+func (t *txFake) Exec(ctx context.Context, sql string, params *params.Params) (driver.Result, error) {
+	result, err := t.conn.Exec(ctx, sql, params)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
+
+	return result, nil
+}
+
+func (t *txFake) Query(ctx context.Context, sql string, params *params.Params) (driver.RowsNextResultSet, error) {
+	rows, err := t.conn.Query(ctx, sql, params)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
+
+	return rows, nil
+}
+
+func (t *txFake) ID() string {
+	return tx.FakeTxID
+}
+
+func beginTxFake(ctx context.Context, c *Conn) conn.Tx {
+	return &txFake{
+		conn: c,
+		ctx:  ctx,
+	}
+}
+
+func (t *txFake) Commit(ctx context.Context) (err error) {
+	if !t.conn.isReady() {
+		return badconn.Map(xerrors.WithStackTrace(errNotReadyConn))
+	}
+
+	return nil
+}
+
+func (t *txFake) Rollback(ctx context.Context) (err error) {
+	if !t.conn.isReady() {
+		return badconn.Map(xerrors.WithStackTrace(errNotReadyConn))
+	}
+
+	return err
+}

internal/xsql/conn/table/errors.go

@@ -11,4 +11,5 @@ var (
 	ErrUnsupported     = driver.ErrSkip
 	errConnClosedEarly = xerrors.Retryable(errors.New("conn closed early"), xerrors.InvalidObject())
 	errNotReadyConn    = xerrors.Retryable(errors.New("conn not ready"), xerrors.InvalidObject())
+	ErrWrongQueryMode  = errors.New("wrong query mode")
 )

internal/xsql/conn/table/tx.go

@@ -26,15 +26,7 @@ func (tx *transaction) ID() string {
 func (tx *transaction) Exec(ctx context.Context, sql string, params *params.Params) (driver.Result, error) {
 	m := queryModeFromContext(ctx, tx.conn.defaultQueryMode)
 	if m != DataQueryMode {
-		return nil, badconn.Map(
-			xerrors.WithStackTrace(
-				xerrors.Retryable(
-					fmt.Errorf("wrong query mode: %s", m.String()),
-					xerrors.InvalidObject(),
-					xerrors.WithName("WRONG_QUERY_MODE"),
-				),
-			),
-		)
+		return nil, xerrors.WithStackTrace(fmt.Errorf("%q: %w", m.String(), ErrWrongQueryMode))
 	}
 	_, err := tx.tx.Execute(ctx, sql, params, tx.conn.dataQueryOptions(ctx)...)
 	if err != nil {

internal/xsql/conn/table/tx_fake.go

@@ -12,48 +12,49 @@ import (
 )
 
 type txFake struct {
-	tx.Identifier
-
 	conn *Conn
 	ctx  context.Context //nolint:containedctx
 }
 
-func (tx *txFake) Exec(ctx context.Context, sql string, params *params.Params) (driver.Result, error) {
-	result, err := tx.conn.Exec(ctx, sql, params)
+func (t *txFake) Exec(ctx context.Context, sql string, params *params.Params) (driver.Result, error) {
+	result, err := t.conn.Exec(ctx, sql, params)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	return result, nil
 }
 
-func (tx *txFake) Query(ctx context.Context, sql string, params *params.Params) (driver.RowsNextResultSet, error) {
-	rows, err := tx.conn.Query(ctx, sql, params)
+func (t *txFake) Query(ctx context.Context, sql string, params *params.Params) (driver.RowsNextResultSet, error) {
+	rows, err := t.conn.Query(ctx, sql, params)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
 	return rows, nil
 }
 
+func (t *txFake) ID() string {
+	return tx.FakeTxID
+}
+
 func beginTxFake(ctx context.Context, c *Conn) conn.Tx {
 	return &txFake{
-		Identifier: tx.ID("FAKE"),
-		conn:       c,
-		ctx:        ctx,
+		conn: c,
+		ctx:  ctx,
 	}
 }
 
-func (tx *txFake) Commit(ctx context.Context) (err error) {
-	if !tx.conn.isReady() {
+func (t *txFake) Commit(ctx context.Context) (err error) {
+	if !t.conn.isReady() {
 		return badconn.Map(xerrors.WithStackTrace(errNotReadyConn))
 	}
 
 	return nil
 }
 
-func (tx *txFake) Rollback(ctx context.Context) (err error) {
-	if !tx.conn.isReady() {
+func (t *txFake) Rollback(ctx context.Context) (err error) {
+	if !t.conn.isReady() {
 		return badconn.Map(xerrors.WithStackTrace(errNotReadyConn))
 	}
 

internal/xsql/connector.go

@@ -15,8 +15,8 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
-	query2 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn/query"
-	table2 "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn/table"
+	connOverQueryServiceClient "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn/query"
+	connOverTableServiceClient "github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/conn/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scheme"
@@ -38,8 +38,8 @@ type (
 
 		queryProcessor queryProcessor
 
-		TableOpts             []table2.Option
-		QueryOpts             []query2.Option
+		TableOpts             []connOverTableServiceClient.Option
+		QueryOpts             []connOverQueryServiceClient.Option
 		disableServerBalancer bool
 		onCLose               []func(*Connector)
 
@@ -122,9 +122,9 @@ func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
 		id := uuid.New()
 
 		conn := &connWrapper{
-			cc: query2.New(ctx, c, s, append(
+			cc: connOverQueryServiceClient.New(ctx, c, s, append(
 				c.QueryOpts,
-				query2.WithOnClose(func() {
+				connOverQueryServiceClient.WithOnClose(func() {
 					c.conns.Delete(id)
 				}))...,
 			),
@@ -145,8 +145,8 @@ func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
 		id := uuid.New()
 
 		conn := &connWrapper{
-			cc: table2.New(ctx, c, s, append(c.TableOpts,
-				table2.WithOnClose(func() {
+			cc: connOverTableServiceClient.New(ctx, c, s, append(c.TableOpts,
+				connOverTableServiceClient.WithOnClose(func() {
 					c.conns.Delete(id)
 				}))...,
 			),

internal/xsql/errors.go

@@ -12,5 +12,5 @@ var (
 	errDeprecated          = driver.ErrSkip
 	errAlreadyClosed       = errors.New("already closed")
 	errWrongQueryProcessor = errors.New("wrong query processor")
-	errNotReadyConn        = xerrors.Retryable(errors.New("connWrapper not ready"), xerrors.InvalidObject())
+	errNotReadyConn        = xerrors.Retryable(errors.New("conn not ready"), xerrors.InvalidObject())
 )

internal/xsql/options.go

@@ -170,6 +170,22 @@ func WithIdleThreshold(idleThreshold time.Duration) Option {
 	}
 }
 
+type mergedOptions []Option
+
+func (opts mergedOptions) Apply(c *Connector) error {
+	for _, opt := range opts {
+		if err := opt.Apply(c); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func Merge(opts ...Option) Option {
+	return mergedOptions(opts)
+}
+
 func WithTableOptions(opts ...table.Option) Option {
 	return tableQueryOptionsOption{
 		tableOps: opts,

internal/xsql/tx.go

@@ -28,6 +28,10 @@ var (
 )
 
 func (tx *txWrapper) Commit() (finalErr error) {
+	defer func() {
+		tx.conn.currentTx = nil
+	}()
+
 	var (
 		ctx    = tx.ctx
 		onDone = trace.DatabaseSQLOnTxCommit(tx.conn.connector.Trace(), &ctx,
@@ -47,6 +51,10 @@ func (tx *txWrapper) Commit() (finalErr error) {
 }
 
 func (tx *txWrapper) Rollback() (finalErr error) {
+	defer func() {
+		tx.conn.currentTx = nil
+	}()
+
 	var (
 		ctx    = tx.ctx
 		onDone = trace.DatabaseSQLOnTxRollback(tx.conn.connector.Trace(), &ctx,