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 }