Ent custom driver

customdriver.go

package customdriver

import (
	"context"

	"entgo.io/ent/dialect"
)

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

type CustomDriver struct {
	dialect.Driver
}

func (c CustomDriver) Exec(ctx context.Context, query string, args, v interface{}) error {
	// create the span here
	return c.Driver.Exec(ctx, query, args, v)
}

func (c CustomDriver) Query(ctx context.Context, query string, args, v interface{}) error {
	// create the span here
	return c.Driver.Exec(ctx, query, args, v)
}

func (c CustomDriver) Tx(ctx context.Context) (dialect.Tx, error) {
	tx, err := c.Driver.Tx(ctx)
	if err != nil {
		return nil, err
	}
	return &CustomTx{tx, ctx}, nil
}

func (c CustomDriver) Close() error {
	return c.Driver.Close()
}

func (c CustomDriver) Dialect() string {
	return c.Driver.Dialect()
}

var _ dialect.Tx = (*CustomTx)(nil)

type CustomTx struct {
	dialect.Tx
	ctx context.Context
}

func (c CustomTx) Exec(ctx context.Context, query string, args, v interface{}) error {
	// create the span here
	return c.Tx.Exec(ctx, query, args, v)
}

func (c CustomTx) Query(ctx context.Context, query string, args, v interface{}) error {
	// create the span here
	return c.Tx.Query(ctx, query, args, v)
}

func (c CustomTx) Commit() error {
	return c.Tx.Commit()
}

func (c CustomTx) Rollback() error {
	return c.Tx.Rollback()
}