Merge pull request #339 from ydb-platform/rm-tx

make conn as tx

Author: asmyasnikov

internal/xsql/conn.go

@@ -55,7 +55,35 @@ type conn struct {
 	dataOpts []options.ExecuteDataQueryOption
 	scanOpts []options.ExecuteScanQueryOption
 
-	currentTx *tx
+	transaction table.Transaction
+}
+
+func (c *conn) Commit() error {
+	if c.isClosed() {
+		return errClosedConn
+	}
+	defer func() {
+		c.transaction = nil
+	}()
+	_, err := c.transaction.CommitTx(context.Background())
+	if err != nil {
+		return c.checkClosed(err)
+	}
+	return nil
+}
+
+func (c *conn) Rollback() error {
+	if c.isClosed() {
+		return errClosedConn
+	}
+	defer func() {
+		c.transaction = nil
+	}()
+	err := c.transaction.Rollback(context.Background())
+	if err != nil {
+		return c.checkClosed(err)
+	}
+	return err
 }
 
 var (
@@ -119,7 +147,6 @@ func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, e
 	}
 	return &stmt{
 		conn:   c,
-		tx:     c.currentTx,
 		params: internal.Params(s),
 		query:  query,
 	}, nil
@@ -129,38 +156,38 @@ func (c *conn) BeginTx(ctx context.Context, txOptions driver.TxOptions) (driver.
 	if c.isClosed() {
 		return nil, errClosedConn
 	}
-	if c.currentTx != nil {
+	if c.transaction != nil {
 		return nil, xerrors.WithStackTrace(
-			fmt.Errorf("conn already have an opened tx: %s", c.currentTx.transaction.ID()),
+			fmt.Errorf("conn already have an opened tx: %s", c.transaction.ID()),
 		)
 	}
 	txSettings, err := isolation.ToYDB(txOptions)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
-	t, err := c.session.BeginTransaction(ctx, txSettings)
+	c.transaction, err = c.session.BeginTransaction(ctx, txSettings)
 	if err != nil {
 		return nil, c.checkClosed(err)
 	}
-	c.currentTx = &tx{
-		conn:        c,
-		transaction: t,
-	}
-	return c.currentTx, nil
+	return c, nil
 }
 
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
 	if c.isClosed() {
 		return nil, errClosedConn
 	}
 	m := queryModeFromContext(ctx, c.defaultQueryMode)
-	if c.currentTx != nil {
+	if c.transaction != nil {
 		if m != DataQueryMode {
 			return nil, xerrors.WithStackTrace(
 				fmt.Errorf("query mode `%s` not supported with transaction", m.String()),
 			)
 		}
-		return c.currentTx.ExecContext(ctx, query, args)
+		_, err := c.transaction.Execute(ctx, query, toQueryParams(args), dataQueryOptions(ctx)...)
+		if err != nil {
+			return nil, c.checkClosed(err)
+		}
+		return c, nil
 	}
 	switch m {
 	case DataQueryMode:
@@ -199,13 +226,22 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		return nil, errClosedConn
 	}
 	m := queryModeFromContext(ctx, c.defaultQueryMode)
-	if c.currentTx != nil {
+	if c.transaction != nil {
 		if m != DataQueryMode {
 			return nil, xerrors.WithStackTrace(
 				fmt.Errorf("query mode `%s` not supported with transaction", m.String()),
 			)
 		}
-		return c.currentTx.QueryContext(ctx, query, args)
+		res, err := c.transaction.Execute(ctx, query, toQueryParams(args), dataQueryOptions(ctx)...)
+		if err != nil {
+			return nil, c.checkClosed(err)
+		}
+		if err = res.Err(); err != nil {
+			return nil, c.checkClosed(err)
+		}
+		return &rows{
+			result: res,
+		}, nil
 	}
 	switch m {
 	case DataQueryMode:

internal/xsql/stmt.go

@@ -6,16 +6,13 @@ import (
 	"fmt"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
-
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 )
 
 type stmt struct {
 	nopResult
 	namedValueChecker
 
 	conn   *conn
-	tx     *tx
 	params map[string]*Ydb.Type
 	query  string
 }
@@ -32,67 +29,23 @@ func (s *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driv
 	if s.conn.isClosed() {
 		return nil, errClosedConn
 	}
-	m := queryModeFromContext(ctx, s.conn.defaultQueryMode)
-	if s.tx != nil {
-		if m != DataQueryMode {
-			return nil, xerrors.WithStackTrace(
-				fmt.Errorf("query mode `%s` not supported with prepared statement", m.String()),
-			)
-		}
-		return s.tx.QueryContext(withKeepInCache(ctx), s.query, args)
-	}
-	switch m {
+	switch m := queryModeFromContext(withKeepInCache(ctx), s.conn.defaultQueryMode); m {
 	case DataQueryMode:
-		_, res, err := s.conn.session.Execute(ctx,
-			txControl(ctx, s.conn.defaultTxControl),
-			s.query,
-			toQueryParams(args),
-			dataQueryOptions(withKeepInCache(ctx))...,
-		)
-		if err != nil {
-			return nil, s.conn.checkClosed(err)
-		}
-		if err = res.Err(); err != nil {
-			return nil, s.conn.checkClosed(res.Err())
-		}
-		return &rows{
-			result: res,
-		}, nil
+		return s.conn.QueryContext(ctx, s.query, args)
 	default:
-		return nil, fmt.Errorf("unsupported query mode '%s' for execute statement query", m)
+		return nil, fmt.Errorf("unsupported query mode '%s' for execute query on prepared statement", m)
 	}
 }
 
 func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
 	if s.conn.isClosed() {
 		return nil, errClosedConn
 	}
-	m := queryModeFromContext(ctx, s.conn.defaultQueryMode)
-	if s.tx != nil {
-		if m != DataQueryMode {
-			return nil, xerrors.WithStackTrace(
-				fmt.Errorf("query mode `%s` not supported with prepared statement", m.String()),
-			)
-		}
-		return s.tx.ExecContext(withKeepInCache(ctx), s.query, args)
-	}
-	switch m {
+	switch m := queryModeFromContext(withKeepInCache(ctx), s.conn.defaultQueryMode); m {
 	case DataQueryMode:
-		_, res, err := s.conn.session.Execute(ctx,
-			txControl(ctx, s.conn.defaultTxControl),
-			s.query,
-			toQueryParams(args),
-			dataQueryOptions(withKeepInCache(ctx))...,
-		)
-		if err != nil {
-			return nil, s.conn.checkClosed(err)
-		}
-		if err = res.Err(); err != nil {
-			return nil, s.conn.checkClosed(res.Err())
-		}
-		return s, nil
+		return s.conn.ExecContext(ctx, s.query, args)
 	default:
-		return nil, fmt.Errorf("unsupported query mode '%s' for execute query", m)
+		return nil, fmt.Errorf("unsupported query mode '%s' for execute query on prepared statement", m)
 	}
 }