Repository: ariga/atlas
Pull request #2746 (2024-05-04 08:34:58)

sql/sqlclient: support function based hooks


Patch:
diff --git a/sql/sqlclient/client.go b/sql/sqlclient/client.go
index 81f51716..1171fbae 100644
--- a/sql/sqlclient/client.go
+++ b/sql/sqlclient/client.go
@@ -45,6 +45,7 @@ type (
 		// Functions registered by the drivers and used for opening transactions and their clients.
 		openDriver func(schema.ExecQuerier) (migrate.Driver, error)
 		openTx     TxOpener
+		hooks      []*Hook
 	}
 
 	// TxClient is returned by calling Client.Tx. It behaves the same as Client,
@@ -53,7 +54,8 @@ type (
 		*Client
 
 		// The transaction this Client wraps.
-		Tx *Tx
+		Tx    *Tx
+		hooks []*Hook
 	}
 
 	// URL extends the standard url.URL with additional
@@ -67,6 +69,20 @@ type (
 		// The Schema this client is connected to.
 		Schema string
 	}
+
+	// Hook groups all possible hooks in
+	// connection and transaction lifecycle.
+	Hook struct {
+		Conn struct {
+			AfterOpen   func(context.Context, *Client) error
+			BeforeClose func(*Client) error
+		}
+		Tx struct {
+			AfterBegin func(context.Context, *TxClient) error
+			BeforeCommit,
+			BeforeRollback func(*TxClient) error
+		}
+	}
 )
 
 // Tx returns a transactional client.
@@ -95,17 +111,23 @@ func (c *Client) Tx(ctx context.Context, opts *sql.TxOptions) (*TxClient, error)
 	}
 	ic := *c
 	ic.Driver = drv
-	return &TxClient{Client: &ic, Tx: tx}, nil
+	tc := &TxClient{Client: &ic, Tx: tx, hooks: c.hooks}
+	if len(tc.hooks) > 0 {
+		if err := tc.afterBegin(ctx); err != nil {
+			return nil, errors.Join(err, tx.Rollback())
+		}
+	}
+	return tc, nil
 }
 
 // Commit the transaction.
 func (c *TxClient) Commit() error {
-	return c.Tx.Commit()
+	return errors.Join(c.beforeCommit(), c.Tx.Commit())
 }
 
 // Rollback the transaction.
 func (c *TxClient) Rollback() error {
-	return c.Tx.Rollback()
+	return errors.Join(c.beforeRollback(), c.Tx.Rollback())
 }
 
 // AddClosers adds list of closers to close at the end of the client lifetime.
@@ -115,16 +137,88 @@ func (c *Client) AddClosers(closers ...io.Closer) {
 
 // Close closes the underlying database connection and the migration
 // driver in case it implements the io.Closer interface.
-func (c *Client) Close() (err error) {
+func (c *Client) Close() error {
+	err := c.beforeClose()
 	for _, closer := range append(c.closers, c.DB) {
-		if cerr := closer.Close(); cerr != nil {
-			if err != nil {
-				cerr = fmt.Errorf("%v: %v", err, cerr)
+		err = errors.Join(err, closer.Close())
+	}
+	return err
+}
+
+type hookCtxKey struct{}
+
+// hookCtx marks the context as being in a hook.
+func hookCtx(ctx context.Context) context.Context {
+	if ctx.Value(hookCtxKey{}) == nil {
+		return context.WithValue(ctx, hookCtxKey{}, true)
+	}
+	return ctx
+}
+
+// afterOpen calls the AfterOpen hooks.
+func (c *Client) afterOpen(ctx context.Context) error {
+	if ctx.Value(hookCtxKey{}) != nil {
+		return errors.New("sql/sqlclient: cannot open a connection inside a hook")
+	}
+	for _, h := range c.hooks {
+		if f := h.Conn.AfterOpen; f != nil {
+			if err := f(hookCtx(ctx), c); err != nil {
+				return err
 			}
-			err = cerr
 		}
 	}
-	return err
+	return nil
+}
+
+// beforeClose calls the BeforeClose hooks.
+func (c *Client) beforeClose() error {
+	for _, h := range c.hooks {
+		if f := h.Conn.BeforeClose; f != nil {
+			if err := f(c); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// afterBegin calls the AfterBegin hooks.
+func (c *TxClient) afterBegin(ctx context.Context) error {
+	if ctx.Value(hookCtxKey{}) != nil {
+		return errors.New("sql/sqlclient: cannot begin a transaction inside a hook")
+	}
+	for _, h := range c.hooks {
+		if f := h.Tx.AfterBegin; f != nil {
+			if err := f(hookCtx(ctx), c); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// beforeCommit calls the BeforeCommit hooks.
+func (c *TxClient) beforeCommit() error {
+	for _, h := range c.hooks {
+		if f := h.Tx.BeforeCommit; f != nil {
+			if err := f(c); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// beforeRollback calls the BeforeRollback hooks.
+func (c *TxClient) beforeRollback() error {
+	for _, h := range c.hooks {
+		if f := h.Tx.BeforeRollback; f != nil {
+			if err := f(c); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
 }
 
 type (
@@ -186,8 +280,8 @@ type (
 	// openOptions holds additional configuration values for opening a Client.
 	openOptions struct {
 		schema *string
+		hooks  []*Hook
 	}
-
 	// OpenOption allows to configure a openOptions using functional arguments.
 	OpenOption func(*openOptions) error
 )
@@ -238,6 +332,12 @@ func OpenURL(ctx context.Context, u *url.URL, opts ...OpenOption) (*Client, erro
 	if client.openTx == nil && drv.txOpener != nil {
 		client.openTx = drv.txOpener
 	}
+	if len(cfg.hooks) > 0 {
+		client.hooks = cfg.hooks
+		if err := client.afterOpen(ctx); err != nil {
+			return nil, errors.Join(err, client.DB.Close())
+		}
+	}
 	return client, nil
 }
 
@@ -250,6 +350,15 @@ func OpenSchema(s string) OpenOption {
 	}
 }
 
+// OpenWithHooks returns an OpenOption that sets
+// the hooks for the client after opening.
+func OpenWithHooks(hks ...*Hook) OpenOption {
+	return func(c *openOptions) error {
+		c.hooks = append(c.hooks, hks...)
+		return nil
+	}
+}
+
 type (
 	registerOptions struct {
 		openDriver func(schema.ExecQuerier) (migrate.Driver, error)


Test patch:
diff --git a/sql/sqlclient/client_test.go b/sql/sqlclient/client_test.go
index 353ead77..e6100beb 100644
--- a/sql/sqlclient/client_test.go
+++ b/sql/sqlclient/client_test.go
@@ -7,6 +7,7 @@ package sqlclient_test
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"net/url"
 	"testing"
 
@@ -168,3 +169,109 @@ type mockDriver struct {
 func (m *mockDriver) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
 	return m.db.ExecContext(ctx, query, args...)
 }
+
+func TestClientHooks(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	sqlclient.Register(
+		"hook",
+		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
+		}),
+	)
+	var (
+		calls [5]int
+		hk    = &sqlclient.Hook{}
+	)
+	hk.Conn.AfterOpen = func(context.Context, *sqlclient.Client) error {
+		calls[0]++
+		return nil
+	}
+	hk.Conn.BeforeClose = func(*sqlclient.Client) error {
+		calls[1]++
+		return nil
+	}
+	hk.Tx.AfterBegin = func(context.Context, *sqlclient.TxClient) error {
+		calls[2]++
+		return nil
+	}
+	hk.Tx.BeforeCommit = func(*sqlclient.TxClient) error {
+		calls[3]++
+		return nil
+	}
+	hk.Tx.BeforeRollback = func(*sqlclient.TxClient) error {
+		calls[4]++
+		return nil
+	}
+	mock.ExpectClose()
+	oc, err := sqlclient.Open(context.Background(), "hook://", sqlclient.OpenWithHooks(hk))
+	require.NoError(t, err)
+	require.NoError(t, oc.Close())
+	require.Equal(t, [5]int{1, 1, 0, 0, 0}, calls)
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	db, mock, err = sqlmock.New()
+	require.NoError(t, err)
+	mock.ExpectBegin()
+	mock.ExpectCommit()
+	mock.ExpectClose()
+	oc, err = sqlclient.Open(context.Background(), "hook://", sqlclient.OpenWithHooks(hk))
+	require.NoError(t, err)
+	tc, err := oc.Tx(context.Background(), nil)
+	require.NoError(t, err)
+	require.NoError(t, tc.Commit())
+	require.NoError(t, oc.Close())
+	require.Equal(t, [5]int{2, 2, 1, 1, 0}, calls)
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	db, mock, err = sqlmock.New()
+	require.NoError(t, err)
+	mock.ExpectBegin()
+	mock.ExpectRollback()
+	mock.ExpectClose()
+	oc, err = sqlclient.Open(context.Background(), "hook://", sqlclient.OpenWithHooks(hk))
+	require.NoError(t, err)
+	tc, err = oc.Tx(context.Background(), nil)
+	require.NoError(t, err)
+	require.NoError(t, tc.Rollback())
+	require.NoError(t, oc.Close())
+	require.Equal(t, [5]int{3, 3, 2, 1, 1}, calls)
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	// Open hook failed.
+	hk.Conn.AfterOpen = func(context.Context, *sqlclient.Client) error {
+		calls[0]++
+		return errors.New("open failed")
+	}
+	db, mock, err = sqlmock.New()
+	require.NoError(t, err)
+	mock.ExpectClose()
+	oc, err = sqlclient.Open(context.Background(), "hook://", sqlclient.OpenWithHooks(hk))
+	require.EqualError(t, err, "open failed")
+	require.Equal(t, [5]int{4, 3, 2, 1, 1}, calls, "close hooks should not be called")
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	// After begin hook failed.
+	hk.Conn.AfterOpen = func(context.Context, *sqlclient.Client) error {
+		calls[0]++
+		return nil
+	}
+	hk.Tx.AfterBegin = func(context.Context, *sqlclient.TxClient) error {
+		calls[2]++
+		return errors.New("after begin failed")
+	}
+	db, mock, err = sqlmock.New()
+	require.NoError(t, err)
+	mock.ExpectBegin()
+	mock.ExpectRollback()
+	mock.ExpectClose()
+	oc, err = sqlclient.Open(context.Background(), "hook://", sqlclient.OpenWithHooks(hk))
+	require.NoError(t, err)
+	tc, err = oc.Tx(context.Background(), nil)
+	require.EqualError(t, err, "after begin failed")
+	require.NoError(t, oc.Close())
+	require.Equal(t, [5]int{5, 4, 3, 1, 1}, calls, "rollback hooks should not be called")
+}