daddy/internal/orm/migration.go

package orm

import (
	"context"

	"github.com/jinzhu/gorm"

	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/middleware/container"
)

type MigrationFunc func(ctx context.Context, tx *gorm.DB) error

type Migration interface {
	Version() string
	Up(context.Context) error
	Down(context.Context) error
}

type DBMigration struct {
	version string
	up      MigrationFunc
	down    MigrationFunc
}

func (m *DBMigration) Version() string {
	return m.version
}

func (m *DBMigration) Up(ctx context.Context) error {
	db, err := m.getDatabase(ctx)
	if err != nil {
		return err
	}

	err = WithTx(ctx, db, func(ctx context.Context, tx *gorm.DB) error {
		return m.up(ctx, tx)
	})

	if err != nil {
		return errors.Wrap(err, "could not apply up migration")
	}

	return nil
}

func (m *DBMigration) Down(ctx context.Context) error {
	db, err := m.getDatabase(ctx)
	if err != nil {
		return err
	}

	err = WithTx(ctx, db, func(ctx context.Context, tx *gorm.DB) error {
		return m.down(ctx, tx)
	})

	if err != nil {
		return errors.Wrap(err, "could not apply down migration")
	}

	return nil
}

func (m *DBMigration) getDatabase(ctx context.Context) (*gorm.DB, error) {
	ctn, err := container.From(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "could not retrieve service container")
	}

	orm, err := From(ctn)
	if err != nil {
		return nil, errors.Wrap(err, "could not retrieve orm service")
	}

	return orm.DB(), nil
}

func NewDBMigration(version string, up, down MigrationFunc) *DBMigration {
	return &DBMigration{
		version: version,
		up:      up,
		down:    down,
	}
}