guesstimate/internal/orm/version_resolver.go

package orm

import (
	"context"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/pkg/errors"
)

type VersionResolver interface {
	Current(context.Context) (string, error)
	Set(context.Context, string) error
}

type DBVersionResolver struct {
	db *gorm.DB
}

type DatabaseVersion struct {
	ID         uint   `gorm:"primary_key"`
	Version    string `gorm:"unique; not null"`
	MigratedAt time.Time
	IsCurrent  bool
}

func (r *DBVersionResolver) Current(ctx context.Context) (string, error) {
	var version string

	err := WithTx(ctx, r.db, func(ctx context.Context, tx *gorm.DB) error {
		dbVersion := &DatabaseVersion{}
		err := tx.Where("is_current = ?", true).First(dbVersion).Error
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil
		}

		if err != nil {
			return errors.WithStack(err)
		}

		version = dbVersion.Version

		return nil
	})

	if err != nil {
		return "", errors.Wrap(err, "could execute version resolver init transaction")
	}

	return version, nil
}

func (r *DBVersionResolver) Set(ctx context.Context, version string) error {
	err := WithTx(ctx, r.db, func(ctx context.Context, tx *gorm.DB) error {
		dbVersion := &DatabaseVersion{
			Version:    version,
			MigratedAt: time.Now(),
		}

		if version != "" {
			if err := tx.FirstOrCreate(dbVersion).Error; err != nil {
				return err
			}

			err := tx.Model(dbVersion).
				UpdateColumn("is_current", true).Error
			if err != nil {
				return errors.WithStack(err)
			}
		}

		err := tx.Model(&DatabaseVersion{}).
			Where("version <> ?", version).
			UpdateColumn("is_current", false).Error

		if err != nil {
			return errors.WithStack(err)
		}

		return err
	})

	if err != nil {
		return errors.Wrap(err, "could not update schema version")
	}

	return nil
}

func (r *DBVersionResolver) Init(ctx context.Context) error {
	err := WithTx(ctx, r.db, func(ctx context.Context, tx *gorm.DB) error {
		if err := tx.AutoMigrate(&DatabaseVersion{}).Error; err != nil {
			return errors.WithStack(err)
		}
		if err := tx.Model(&DatabaseVersion{}).AddUniqueIndex("idx_unique_version", "version").Error; err != nil {
			return errors.WithStack(err)
		}
		return nil
	})

	if err != nil {
		return errors.Wrap(err, "could execute version resolver init transaction")
	}

	return nil
}

func NewDBVersionResolver(db *gorm.DB) *DBVersionResolver {
	return &DBVersionResolver{
		db: db,
	}
}