Repository: ariga/atlas
Pull request #837 (2022-06-01 08:38:33)

sql/sqlclient: add support for transactions


Patch:
diff --git a/sql/sqlclient/client.go b/sql/sqlclient/client.go
index 436679c6..bed6dcb2 100644
--- a/sql/sqlclient/client.go
+++ b/sql/sqlclient/client.go
@@ -45,6 +45,15 @@ type (
 		schemahcl.Evaluator
 	}
 
+	// TxClient is returned by calling Client.Tx. It behaves the same as Client,
+	// but wraps all operations within a transaction.
+	TxClient struct {
+		*Client
+
+		// The transaction this Client wraps.
+		Tx *sql.Tx
+	}
+
 	// URL extends the standard url.URL with additional
 	// connection information attached by the Opener (if any).
 	URL struct {
@@ -58,6 +67,35 @@ type (
 	}
 )
 
+// Tx returns a transactional client.
+func (c *Client) Tx(ctx context.Context, opts *sql.TxOptions) (*TxClient, error) {
+	tx, err := c.DB.BeginTx(ctx, opts)
+	if err != nil {
+		return nil, fmt.Errorf("sql/sqlclient: starting transaction: %w", err)
+	}
+	v, ok := drivers.Load(c.Name)
+	if !ok {
+		return nil, fmt.Errorf("sql/sqlclient: unexpected missing opener %q", c.Name)
+	}
+	drv, err := v.(*driver).openDriver(tx)
+	if err != nil {
+		return nil, fmt.Errorf("sql/sqlclient: opengin atlas driver: %w", err)
+	}
+	ic := *c
+	ic.Driver = drv
+	return &TxClient{Client: &ic, Tx: tx}, nil
+}
+
+// Commit the transaction.
+func (c *TxClient) Commit() error {
+	return c.Tx.Commit()
+}
+
+// Rollback the transaction.
+func (c *TxClient) Rollback() error {
+	return c.Tx.Rollback()
+}
+
 // AddClosers adds list of closers to close at the end of the client lifetime.
 func (c *Client) AddClosers(closers ...io.Closer) {
 	c.closers = append(c.closers, closers...)
@@ -101,8 +139,9 @@ type (
 
 	driver struct {
 		Opener
-		name   string
-		parser URLParser
+		name       string
+		parser     URLParser
+		openDriver func(schema.ExecQuerier) (migrate.Driver, error)
 	}
 )
 
@@ -154,18 +193,18 @@ func OpenURL(ctx context.Context, u *url.URL, opts ...OpenOption) (*Client, erro
 	}
 	// If there is a schema given and the driver allows to change the schema for the url, do it.
 	if cfg.schema != "" {
-		sc, ok := v.(driver).parser.(SchemaChanger)
+		sc, ok := v.(*driver).parser.(SchemaChanger)
 		if !ok {
 			return nil, ErrUnsupported
 		}
 		u = sc.ChangeSchema(u, cfg.schema)
 	}
-	client, err := v.(driver).Open(ctx, u)
+	client, err := v.(*driver).Open(ctx, u)
 	if err != nil {
 		return nil, err
 	}
 	if client.URL == nil {
-		client.URL = v.(driver).parser.ParseURL(u)
+		client.URL = v.(*driver).parser.ParseURL(u)
 	}
 	return client, nil
 }
@@ -181,9 +220,10 @@ func OpenSchema(s string) OpenOption {
 
 type (
 	registerOptions struct {
-		parser   URLParser
-		flavours []string
-		codec    interface {
+		openDriver func(schema.ExecQuerier) (migrate.Driver, error)
+		parser     URLParser
+		flavours   []string
+		codec      interface {
 			schemaspec.Marshaler
 			schemahcl.Evaluator
 		}
@@ -223,6 +263,14 @@ func RegisterCodec(m schemaspec.Marshaler, e schemahcl.Evaluator) RegisterOption
 	}
 }
 
+// RegisterDriverOpener registers a func to create a migrate.Driver from a schema.ExecQuerier.
+// Registering this function is implicitly done when using DriverOpener.
+func RegisterDriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) RegisterOption {
+	return func(opts *registerOptions) {
+		opts.openDriver = open
+	}
+}
+
 // DriverOpener is a helper Opener creator for sharing between all drivers.
 func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener {
 	return OpenerFunc(func(ctx context.Context, u *url.URL) (*Client, error) {
@@ -230,8 +278,11 @@ func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener
 		if !ok {
 			return nil, fmt.Errorf("sql/sqlclient: unexpected missing opener %q", u.Scheme)
 		}
-		ur := v.(driver).parser.ParseURL(u)
-		db, err := sql.Open(v.(driver).name, ur.DSN)
+		if v.(*driver).openDriver == nil {
+			v.(*driver).openDriver = open
+		}
+		ur := v.(*driver).parser.ParseURL(u)
+		db, err := sql.Open(v.(*driver).name, ur.DSN)
 		if err != nil {
 			return nil, err
 		}
@@ -243,7 +294,7 @@ func DriverOpener(open func(schema.ExecQuerier) (migrate.Driver, error)) Opener
 			return nil, err
 		}
 		return &Client{
-			Name:   v.(driver).name,
+			Name:   v.(*driver).name,
 			DB:     db,
 			URL:    ur,
 			Driver: drv,
@@ -278,10 +329,11 @@ func Register(name string, opener Opener, opts ...RegisterOption) {
 		if _, ok := drivers.Load(f); ok {
 			panic("sql/sqlclient: Register called twice for " + f)
 		}
-		drivers.Store(f, driver{
-			name:   name,
-			parser: opt.parser,
-			Opener: opener,
+		drivers.Store(f, &driver{
+			name:       name,
+			parser:     opt.parser,
+			openDriver: opt.openDriver,
+			Opener:     opener,
 		})
 	}
 }


Test patch:
diff --git a/sql/sqlclient/client_test.go b/sql/sqlclient/client_test.go
index 4834100b..971de93b 100644
--- a/sql/sqlclient/client_test.go
+++ b/sql/sqlclient/client_test.go
@@ -10,6 +10,9 @@ import (
 	"net/url"
 	"testing"
 
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
+	"github.com/DATA-DOG/go-sqlmock"
 	"github.com/stretchr/testify/require"
 
 	"ariga.io/atlas/sql/sqlclient"
@@ -71,3 +74,55 @@ func TestClient_AddClosers(t *testing.T) {
 type closerFunc func() error
 
 func (f closerFunc) Close() error { return f() }
+
+func TestClient_Tx(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	defer db.Close()
+
+	const stmt = "create database `test`"
+	mock.ExpectBegin()
+	mock.ExpectExec(stmt).WillReturnResult(sqlmock.NewResult(0, 1))
+	mock.ExpectCommit()
+	mock.ExpectBegin()
+	mock.ExpectExec(stmt).WillReturnResult(sqlmock.NewResult(0, 1))
+	mock.ExpectRollback()
+
+	sqlclient.Register(
+		"tx",
+		sqlclient.OpenerFunc(func(context.Context, *url.URL) (*sqlclient.Client, error) {
+			return &sqlclient.Client{Name: "tx", DB: db, Driver: &mockDriver{db: db}}, nil
+		}),
+		sqlclient.RegisterDriverOpener(func(db schema.ExecQuerier) (migrate.Driver, error) {
+			return &mockDriver{db: db}, nil
+		}),
+	)
+
+	c, err := sqlclient.Open(context.Background(), "tx://")
+	require.NoError(t, err)
+
+	// Commit works.
+	tx, err := c.Tx(context.Background(), nil)
+	require.NoError(t, err)
+	_, err = tx.ExecContext(context.Background(), stmt)
+	require.NoError(t, err)
+	require.NoError(t, tx.Commit())
+
+	// Rollback works as well.
+	tx, err = c.Tx(context.Background(), nil)
+	require.NoError(t, err)
+	_, err = tx.ExecContext(context.Background(), stmt)
+	require.NoError(t, err)
+	require.NoError(t, tx.Rollback())
+
+	require.NoError(t, mock.ExpectationsWereMet())
+}
+
+type mockDriver struct {
+	migrate.Driver
+	db schema.ExecQuerier
+}
+
+func (m *mockDriver) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	return m.db.ExecContext(ctx, query, args...)
+}