core/cmd/twhelp/cmd_db.go

package main

import (
	"context"
	"fmt"
	"log/slog"
	"strings"

	"gitea.dwysokinski.me/twhelp/core/internal/bun/migrations"
	"github.com/uptrace/bun/migrate"
	"github.com/urfave/cli/v2"
)

const migrationGoTemplate = `package %s

import (
	"context"
	"fmt"

	"github.com/uptrace/bun"
)

func init() {
	migrations.MustRegister(func(ctx context.Context, db *bun.DB) error {
		fmt.Print(" [up migration] ")
		return nil
	}, func(ctx context.Context, db *bun.DB) error {
		fmt.Print(" [down migration] ")
		return nil
	})
}
`

var cmdDB = &cli.Command{
	Name:  "db",
	Usage: "Manage the database",
	Subcommands: []*cli.Command{
		{
			Name:  "migrate",
			Usage: "Migrate database",
			Flags: dbFlags,
			Action: func(c *cli.Context) error {
				shutdownSignalCtx, stop := newShutdownSignalContext(c.Context)
				defer stop()

				logger := loggerFromCtx(c.Context)

				bunDB, err := newBunDBFromFlags(c)
				if err != nil {
					return err
				}
				defer closeBunDB(bunDB, logger)

				migrator := migrations.NewMigrator(bunDB)

				if err = migrator.Init(shutdownSignalCtx); err != nil {
					return fmt.Errorf("couldn't init migrator: %w", err)
				}

				// we use c.Context instead of shutdownSignalCtx here because we still want to unlock the db
				// even after one of the shutdown signals has been received
				defer unlockMigrations(c.Context, migrator, logger)

				if err = migrator.Lock(shutdownSignalCtx); err != nil {
					return fmt.Errorf("couldn't lock db: %w", err)
				}

				group, err := migrator.Migrate(shutdownSignalCtx)
				if err != nil {
					return fmt.Errorf("migration failed: %w", err)
				}

				if group.ID == 0 {
					logger.Info("there are no new migrations to run")
					return nil
				}

				logger.Info(
					"migrations have been successfully applied",
					slog.Int64("group.id", group.ID),
					slog.String("group.migrations", group.Migrations.String()),
				)

				return nil
			},
		},
		{
			Name:  "rollback",
			Usage: "Rollback the last migration group",
			Flags: dbFlags,
			Action: func(c *cli.Context) error {
				shutdownSignalCtx, stop := newShutdownSignalContext(c.Context)
				defer stop()

				logger := loggerFromCtx(c.Context)

				bunDB, err := newBunDBFromFlags(c)
				if err != nil {
					return err
				}
				defer closeBunDB(bunDB, logger)

				migrator := migrations.NewMigrator(bunDB)

				// we use c.Context instead of shutdownSignalCtx here because we still want to unlock the db
				// even after one of the shutdown signals has been received
				defer unlockMigrations(c.Context, migrator, logger)

				if err = migrator.Lock(shutdownSignalCtx); err != nil {
					return fmt.Errorf("couldn't lock db: %w", err)
				}

				group, err := migrator.Rollback(shutdownSignalCtx)
				if err != nil {
					return fmt.Errorf("couldn't rollback last migration: %w", err)
				}

				if group.ID == 0 {
					logger.Info("there are no groups to roll back")
					return nil
				}

				logger.Info(
					"last migration group has been rolled back",
					slog.Int64("group.id", group.ID),
					slog.String("group.migrations", group.Migrations.String()),
				)

				return nil
			},
		},
		{
			Name:  "create",
			Usage: "Create migration",
			Subcommands: []*cli.Command{
				{
					Name:  "go",
					Usage: "Create Go migration",
					Action: func(c *cli.Context) error {
						logger := loggerFromCtx(c.Context)

						migrator := migrations.NewMigrator(nil)

						mf, err := migrator.CreateGoMigration(
							c.Context,
							strings.Join(c.Args().Slice(), "_"),
							migrate.WithGoTemplate(migrationGoTemplate),
						)
						if err != nil {
							return fmt.Errorf("couldn't create Go migration: %w", err)
						}

						logger.Info(
							"Go migration has been created",
							slog.String("migration.path", mf.Path),
							slog.String("migration.name", mf.Name),
						)

						return nil
					},
				},
				{
					Name:  "sql",
					Usage: "Create SQL migration",
					Action: func(c *cli.Context) error {
						logger := loggerFromCtx(c.Context)

						migrator := migrations.NewMigrator(nil)

						files, err := migrator.CreateSQLMigrations(c.Context, strings.Join(c.Args().Slice(), "_"))
						if err != nil {
							return fmt.Errorf("couldn't create sql migration: %w", err)
						}

						for _, mf := range files {
							logger.Info(
								"SQL migration has been created",
								slog.String("migration.path", mf.Path),
								slog.String("migration.name", mf.Name),
							)
						}

						return nil
					},
				},
			},
		},
	},
}

func unlockMigrations(ctx context.Context, migrator *migrate.Migrator, logger *slog.Logger) {
	logger.Debug("unlocking db...")
	if unlockErr := migrator.Unlock(ctx); unlockErr != nil {
		logger.Warn("couldn't unlock db", slog.Any("error", unlockErr))
	} else {
		logger.Debug("db unlocked")
	}
}