pkg/framework/entutils/transaction.go

package entutils

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"
	"sync"

	"entgo.io/ent/dialect"

	"github.com/openmeterio/openmeter/pkg/framework/transaction"
)

type RawEntConfig struct {
	// driver used for executing database requests.
	Driver dialect.Driver
	// debug enable a debug logging.
	Debug bool
	// log used for logging on debug mode.
	Log func(...any)

	// Hooks and interceptors are excluded in transaction handling
	// due to differing types.
	//
	// TODO: implement them in the templating when creating the new transactional client
	// from this RawEntConfig.

	// // hooks to execute on mutations.
	// hooks *hooks
	// // interceptors to execute on queries.
	// inters *inters
}

type Transactable interface {
	Commit() error
	Rollback() error
	SavePoint(name string) error
	RollbackTo(name string) error
	Release(name string) error
}

type TxHijacker interface {
	HijackTx(ctx context.Context, opts *sql.TxOptions) (context.Context, *RawEntConfig, Transactable, error)
}

func NewTxDriver(driver Transactable, cfg *RawEntConfig) *TxDriver {
	return &TxDriver{
		driver: driver,
		cfg:    cfg,
	}
}

type txSavepoint int

const (
	txSavepointNone txSavepoint = 0
)

func (sp txSavepoint) Next() txSavepoint {
	return sp + 1
}

func (sp txSavepoint) Prev() txSavepoint {
	if sp == txSavepointNone {
		return txSavepointNone
	}

	return sp - 1
}

func (sp txSavepoint) String() string {
	return "s" + strconv.Itoa(int(sp))
}

type TxDriver struct {
	driver Transactable
	// db.config is nominally different but structurally identical for all generations of entgo,
	// so we represent it as an interface{} here
	cfg *RawEntConfig

	mu   sync.Mutex
	once sync.Once

	currentSavepoint txSavepoint

	err error
}

var _ transaction.Driver = &TxDriver{}

func (t *TxDriver) GetConfig() *RawEntConfig {
	return t.cfg
}

// Commit commits the (complete) transaction.
func (t *TxDriver) Commit() error {
	// lock so we don't use the driver twice
	t.mu.Lock()
	defer t.mu.Unlock()

	// If there was an error before, we don't do anything
	if t.err != nil {
		return t.err
	}

	if t.currentSavepoint != txSavepointNone {
		// If we're not at the top level, we release the savepoint
		if err := t.driver.Release(t.currentSavepoint.String()); err == nil {
			t.currentSavepoint = t.currentSavepoint.Prev()
		} else {
			t.err = err
		}
	} else {
		// If we're at the top level, we commit the transaction
		t.err = t.driver.Commit()
	}

	return t.err
}

// Rollback rolls back the (complete) transaction.
func (t *TxDriver) Rollback() error {
	// lock so we don't use the driver twice
	t.mu.Lock()
	defer t.mu.Unlock()

	// If there was an error before, we don't do anything
	if t.err != nil {
		return t.err
	}

	if t.currentSavepoint != txSavepointNone {
		// If we're not at the top level, we rollback to the savepoint
		if err := t.driver.RollbackTo(t.currentSavepoint.String()); err == nil {
			t.currentSavepoint = t.currentSavepoint.Prev()
		} else {
			t.err = err
		}
	} else {
		// If we're at the top level, we rollback the transaction
		t.err = t.driver.Rollback()
	}

	return t.err
}

func (t *TxDriver) SavePoint() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	skipSavePoint := false

	t.once.Do(func() {
		// As savePoint() is called each time we use the wrapper (including the first)
		// we don't want to create a savepoint for the first call, otherwise the transaction itself
		// would never be closed.
		skipSavePoint = true
	})

	if !skipSavePoint {
		next := t.currentSavepoint.Next()

		err := t.driver.SavePoint(next.String())
		if err != nil {
			return err
		}

		t.currentSavepoint = next
	}

	return nil
}

// Able to start a new transaction
type TxCreator = transaction.Creator

// Able to use an existing transaction
type TxUser[T any] interface {
	// Creates a new instance of the adapter using the provided transaction.
	// Example:
	//
	// type dbAdapter struct {
	//     db *db.Client
	// }
	//
	// func (d *dbAdapter) WithTx(ctx context.Context, tx *entutils.TxDriver) SomeDB[db1.Example1] {
	//     // NewTxClientFromRawConfig gets generated when using expose.tpl
	//     txClient := db.NewTxClientFromRawConfig(ctx, *tx.GetConfig())
	//     res := &db1Adapter{db: txClient.Client()}
	//     return res
	// }
	WithTx(ctx context.Context, tx *TxDriver) T
	Self() T
}

// TransactingRepo is a helper that can be used inside repository methods.
// It uses any preexisting transaction in the context if exists.
func TransactingRepo[R, T any](
	ctx context.Context,
	repo interface {
		TxUser[T]
		TxCreator
	},
	cb func(ctx context.Context, rep T) (R, error),
) (R, error) {
	var def R
	tx, err := GetDriverFromContext(ctx)
	if err != nil {
		// For all other errors, we return the error
		if _, ok := err.(*transaction.DriverNotFoundError); !ok {
			return def, err
		}

		// If we're not in a transaction, we just use the repo
		return cb(ctx, repo.Self())
	}

	// If we're in a transaction, we use it
	return cb(ctx, repo.WithTx(ctx, tx))
}

// TransactingRepoWithNoValue is a helper that can be used inside repository methods.
func TransactingRepoWithNoValue[T any](
	ctx context.Context,
	repo interface {
		TxUser[T]
		TxCreator
	},
	cb func(ctx context.Context, rep T) error,
) error {
	_, err := TransactingRepo(ctx, repo, func(ctx context.Context, rep T) (interface{}, error) {
		return nil, cb(ctx, rep)
	})
	return err
}

func asEntDriver(drv transaction.Driver) (*TxDriver, error) {
	entTxDriver, ok := drv.(*TxDriver)
	if !ok {
		return nil, fmt.Errorf("tx driver is not ent tx driver")
	}
	return entTxDriver, nil
}

// Only use for direct interacton with the Ent driver implementation
func GetDriverFromContext(ctx context.Context) (*TxDriver, error) {
	driver, err := transaction.GetDriverFromContext(ctx)
	if err != nil {
		return nil, err
	}
	return asEntDriver(driver)
}