buncomponent/bun.go

package buncomponent

import (
	"context"
	"database/sql"
	"fmt"
	"strings"

	"github.com/golang-migrate/migrate/v4"
	migratePostgres "github.com/golang-migrate/migrate/v4/database/postgres"
	_ "github.com/golang-migrate/migrate/v4/source/file"

	pgxLogrus "github.com/jackc/pgx-logrus"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/stdlib"
	"github.com/jackc/pgx/v5/tracelog"
	"github.com/uptrace/bun"
	"github.com/uptrace/bun/dialect/pgdialect"
	"github.com/uptrace/bun/extra/bundebug"
	"go-micro.dev/v4/errors"

	"github.com/urfave/cli/v2"

	"jochum.dev/jo-micro/components"
	"jochum.dev/jo-micro/logruscomponent"
)

const Name = "bun"

type BunComponent struct {
	initialized bool

	sqlDB *sql.DB
	bun   *bun.DB
}

func New() *BunComponent {
	return &BunComponent{initialized: false}
}

func MustReg(cReg *components.Registry) *BunComponent {
	return cReg.Must(Name).(*BunComponent)
}

func (c *BunComponent) Name() string {
	return Name
}

func (c *BunComponent) Priority() int {
	return 10
}

func (c *BunComponent) Initialized() bool {
	return c.initialized
}

func (c *BunComponent) Init(r *components.Registry, cli *cli.Context) error {
	if c.initialized {
		return nil
	}

	if cli.String(fmt.Sprintf("%s_database_url", strings.ToLower(r.FlagPrefix()))) == "" {
		return errors.InternalServerError("DATABASE_URL_NOT_GIVEN", "%s_DATABASE_URL is required", strings.ToUpper(r.FlagPrefix()))
	} else if strings.HasPrefix(cli.String(fmt.Sprintf("%s_database_url", strings.ToLower(r.FlagPrefix()))), "postgres://") {
		config, err := pgx.ParseConfig(cli.String(fmt.Sprintf("%s_database_url", strings.ToLower(r.FlagPrefix()))))
		if err != nil {
			return err
		}

		if logruscomponent.MustReg(r).Initialized() {
			config.Tracer = &tracelog.TraceLog{Logger: pgxLogrus.NewLogger(logruscomponent.MustReg(r).Logger()), LogLevel: tracelog.LogLevelInfo}
		}

		connStr := stdlib.RegisterConnConfig(config)
		c.sqlDB, _ = sql.Open("pgx", connStr)
		driver, err := migratePostgres.WithInstance(c.sqlDB, &migratePostgres.Config{MigrationsTable: cli.String(fmt.Sprintf("%s_migrations_table", strings.ToLower(r.FlagPrefix())))})
		if err != nil {
			return err
		}

		m, err := migrate.NewWithDatabaseInstance(
			fmt.Sprintf("file://%s/postgres", cli.String(fmt.Sprintf("%s_migrations_dir", strings.ToLower(r.FlagPrefix())))), "postgres", driver)
		if err != nil {
			return errors.InternalServerError("internal/ibun.Start|migrate.NewWithDatabaseInstance", fmt.Sprintf("%s", err))
		}
		if err := m.Up(); err != migrate.ErrNoChange && err != nil {
			return errors.InternalServerError("internal/ibun.Start|migrate.Up", fmt.Sprintf("%s", err))
		}

		c.bun = bun.NewDB(c.sqlDB, pgdialect.New())
		if c.bun == nil {
			return errors.InternalServerError("internal/ibun.Start|bun.NewDB", "failed to create bun")
		}

		if cli.Bool(fmt.Sprintf("%s_database_debug", strings.ToLower(r.FlagPrefix()))) {
			// Print all queries to stdout.
			c.bun.AddQueryHook(bundebug.NewQueryHook(bundebug.WithVerbose(true)))
		}
	} else {
		return errors.InternalServerError("internal/ibun.Start|sqltype", "%s_DATABASE_URL has a unknown type", strings.ToUpper(r.FlagPrefix()))
	}

	c.initialized = true
	return nil
}

func (c *BunComponent) Stop() error {
	c.initialized = false
	return nil
}

func (c *BunComponent) Flags(r *components.Registry) []cli.Flag {
	return []cli.Flag{
		&cli.StringFlag{
			Name:    fmt.Sprintf("%s_database_url", strings.ToLower(r.FlagPrefix())),
			Usage:   "bun Database URL",
			EnvVars: []string{fmt.Sprintf("%s_DATABASE_URL", strings.ToUpper(r.FlagPrefix()))},
		},
		&cli.BoolFlag{
			Name:        fmt.Sprintf("%s_database_debug", strings.ToLower(r.FlagPrefix())),
			Usage:       "Set it to the debug the database queries",
			EnvVars:     []string{fmt.Sprintf("%s_DATABASE_DEBUG", strings.ToUpper(r.FlagPrefix()))},
			DefaultText: "false",
			Value:       false,
		},
		&cli.StringFlag{
			Name:    fmt.Sprintf("%s_migrations_table", strings.ToLower(r.FlagPrefix())),
			Value:   "schema_migrations",
			Usage:   "Table to store migrations info",
			EnvVars: []string{fmt.Sprintf("%s_MIGRATIONS_TABLE", strings.ToUpper(r.FlagPrefix()))},
		},
		&cli.StringFlag{
			Name:    fmt.Sprintf("%s_migrations_dir", strings.ToLower(r.FlagPrefix())),
			Value:   "/migrations",
			Usage:   "Folder which contains migrations",
			EnvVars: []string{fmt.Sprintf("%s_MIGRATIONS_DIR", strings.ToUpper(r.FlagPrefix()))},
		},
	}
}

func (c *BunComponent) Health(context context.Context) error {
	return nil
}

func (c *BunComponent) Bun() *bun.DB {
	return c.bun
}