entc/gen: adding tx callback support (#414)

Signed-off-by: Alex Snast <alexsn@fb.com>

examples/edgeindex/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/edgeindex/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -19,16 +20,47 @@ type Tx struct {
 	City *CityClient
 	// Street is the client for interacting with the Street builders.
 	Street *StreetClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/entcpkg/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/entcpkg/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -17,16 +18,47 @@ type Tx struct {
 	config
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/m2m2types/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/m2m2types/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -19,16 +20,47 @@ type Tx struct {
 	Group *GroupClient
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/m2mbidi/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/m2mbidi/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -17,16 +18,47 @@ type Tx struct {
 	config
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/m2mrecur/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/m2mrecur/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -17,16 +18,47 @@ type Tx struct {
 	config
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/o2m2types/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/o2m2types/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -19,16 +20,47 @@ type Tx struct {
 	Pet *PetClient
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/o2mrecur/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/o2mrecur/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -17,16 +18,47 @@ type Tx struct {
 	config
 	// Node is the client for interacting with the Node builders.
 	Node *NodeClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/o2o2types/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/o2o2types/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -19,16 +20,47 @@ type Tx struct {
 	Card *CardClient
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/o2obidi/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/o2obidi/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -17,16 +18,47 @@ type Tx struct {
 	config
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/o2orecur/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/o2orecur/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -17,16 +18,47 @@ type Tx struct {
 	config
 	// Node is the client for interacting with the Node builders.
 	Node *NodeClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/start/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/start/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -21,16 +22,47 @@ type Tx struct {
 	Group *GroupClient
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.

examples/traversal/ent/context.go

@@ -10,15 +10,28 @@ import (
 	"context"
 )
 
-type contextKey struct{}
+type clientCtxKey struct{}
 
 // FromContext returns the Client stored in a context, or nil if there isn't one.
 func FromContext(ctx context.Context) *Client {
-	c, _ := ctx.Value(contextKey{}).(*Client)
+	c, _ := ctx.Value(clientCtxKey{}).(*Client)
 	return c
 }
 
 // NewContext returns a new context with the given Client attached.
 func NewContext(parent context.Context, c *Client) context.Context {
-	return context.WithValue(parent, contextKey{}, c)
+	return context.WithValue(parent, clientCtxKey{}, c)
+}
+
+type txCtxKey struct{}
+
+// TxFromContext returns the Tx stored in a context, or nil if there isn't one.
+func TxFromContext(ctx context.Context) *Tx {
+	tx, _ := ctx.Value(txCtxKey{}).(*Tx)
+	return tx
+}
+
+// NewTxContext returns a new context with the given Client attached.
+func NewTxContext(parent context.Context, tx *Tx) context.Context {
+	return context.WithValue(parent, txCtxKey{}, tx)
 }

examples/traversal/ent/tx.go

@@ -8,6 +8,7 @@ package ent
 
 import (
 	"context"
+	"sync"
 
 	"github.com/facebookincubator/ent/dialect"
 )
@@ -21,16 +22,47 @@ type Tx struct {
 	Pet *PetClient
 	// User is the client for interacting with the User builders.
 	User *UserClient
+
+	// completion callbacks.
+	mu         sync.Mutex
+	onCommit   []func(error)
+	onRollback []func(error)
 }
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	return tx.config.driver.(*txDriver).tx.Commit()
+	err := tx.config.driver.(*txDriver).tx.Commit()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onCommit {
+		f(err)
+	}
+	return err
+}
+
+// OnCommit adds a function to call on commit.
+func (tx *Tx) OnCommit(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onCommit = append(tx.onCommit, f)
 }
 
 // Rollback rollbacks the transaction.
 func (tx *Tx) Rollback() error {
-	return tx.config.driver.(*txDriver).tx.Rollback()
+	err := tx.config.driver.(*txDriver).tx.Rollback()
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	for _, f := range tx.onRollback {
+		f(err)
+	}
+	return err
+}
+
+// OnRollback adds a function to call on rollback.
+func (tx *Tx) OnRollback(f func(error)) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+	tx.onRollback = append(tx.onRollback, f)
 }
 
 // Client returns a Client that binds to current transaction.