{{/*
Copyright 2019-present Facebook Inc. All rights reserved.
This source code is licensed under the Apache 2.0 license found
in the LICENSE file in the root directory of this source tree.
*/}}

{{ define "tx" }}

{{ template "header" $ }}

import (
	"context"
	"sync"

	"github.com/facebookincubator/ent/dialect"
)

// Tx is a transactional client that is created by calling Client.Tx().
type Tx struct {
	config
	{{ range $_, $n := $.Nodes -}}
		// {{ $n.Name }} is the client for interacting with the {{ $n.Name }} builders.
		{{ $n.Name }} *{{ $n.Name }}Client
	{{ end }}

	// lazily loaded.
	client     *Client
	clientOnce sync.Once

	// completion callbacks.
	mu         sync.Mutex
	onCommit   []func(error)
	onRollback []func(error)
}

{{- range $func := list "Commit" "Rollback" }}
	{{- $onFuncs := print "on" $func }}
	// {{ $func }} {{ lower $func }}s the transaction.
	func (tx *Tx) {{ $func }}() error {
		err := tx.config.driver.(*txDriver).tx.{{ $func }}()
		tx.mu.Lock()
		defer tx.mu.Unlock()
		for _, f := range tx.{{ $onFuncs }} {
			f(err)
		}
		return err
	}

	// On{{ $func }} adds a function to call on {{ lower $func }}.
	func (tx *Tx) On{{ $func }}(f func(error)) {
		tx.mu.Lock()
		defer tx.mu.Unlock()
		tx.{{ $onFuncs }} = append(tx.{{ $onFuncs }}, f)
	}
{{- end }}

// Client returns a Client that binds to current transaction.
func (tx *Tx) Client() *Client {
	tx.clientOnce.Do(func() {
		tx.client = &Client{config: tx.config}
		tx.client.init()
	})
	return tx.client
}

func (tx *Tx) init() {
	{{ range $_, $n := $.Nodes -}}
    	tx.{{ $n.Name }} = New{{ $n.Name }}Client(tx.config)
	{{ end -}}
}

{{/* first node for doc example */}}
{{- $first := index $.Nodes 0 }}

// txDriver wraps the given dialect.Tx with a nop dialect.Driver implementation.
// The idea is to support transactions without adding any extra code to the builders.
// When a builder calls to driver.Tx(), it gets the same dialect.Tx instance.
// Commit and Rollback are nop for the internal builders and the user must call one
// of them in order to commit or rollback the transaction.
//
// If a closed transaction is embedded in one of the generated entities, and the entity
// applies a query, for example: {{ $first.Name }}.QueryXXX(), the query will be executed
// through the driver which created this transaction.
//
// Note that txDriver is not goroutine safe.
type txDriver struct {
	// the driver we started the transaction from.
	drv dialect.Driver
	// tx is the underlying transaction.
	tx dialect.Tx
}

// newTx creates a new transactional driver.
func newTx(ctx context.Context, drv dialect.Driver) (*txDriver, error) {
	tx, err := drv.Tx(ctx)
	if err != nil {
		return nil, err
	}
	return &txDriver{tx: tx, drv: drv}, nil
}

// Tx returns the transaction wrapper (txDriver) to avoid Commit or Rollback calls
// from the internal builders. Should be called only by the internal builders.
func (tx *txDriver) Tx(context.Context) (dialect.Tx, error) { return tx, nil }

// Dialect returns the dialect of the driver we started the transaction from.
func (tx *txDriver) Dialect() string { return tx.drv.Dialect() }

// Close is a nop close.
func (*txDriver) Close() error { return nil }

// Commit is a nop commit for the internal builders.
// User must call `Tx.Commit` in order to commit the transaction.
func (*txDriver) Commit() error { return nil }

// Rollback is a nop rollback for the internal builders.
// User must call `Tx.Rollback` in order to rollback the transaction.
func (*txDriver) Rollback() error { return nil }

// Exec calls tx.Exec.
func (tx *txDriver) Exec(ctx context.Context, query string, args, v interface{}) error {
	return tx.tx.Exec(ctx, query, args, v)
}

// Query calls tx.Query.
func (tx *txDriver) Query(ctx context.Context, query string, args, v interface{}) error {
	return tx.tx.Query(ctx, query, args, v)
}

var _ dialect.Driver = (*txDriver)(nil)

{{ end }}