ims/util/db/mysql/migrate.go

package mysql

import (
	"errors"
	"ims/util/backoff"

	"gorm.io/gorm"
)

type Migrator struct {
	DB    *gorm.DB
	Retry *backoff.Options
	Log   func(format string, args ...any)
}

func (m Migrator) retry(fn func() error) error {
	if m.Retry == nil {
		return fn()
	}
	return backoff.Retry(fn, func(o *backoff.Options) {
		*o = *m.Retry
	})
}

func (m Migrator) acquireLock(tx *gorm.DB, database string) (func() error, error) {
	return AcquireLock(tx, database, m.Retry)
}

// MigrateTenantModels creates a database for a specific tenant and migrates the tenant tables.
func (m Migrator) MigrateTenantModels(tenantId uint, models []any) (err error) {
	m.Log("⏳ migrating tables for tenant %d", tenantId)

	if len(models) == 0 {
		err = errors.New("no tenant tables to migrate")
		return
	}

	tx := m.DB.Session(&gorm.Session{})
	database := TenantDatabase(tenantId)
	sql := "CREATE DATABASE IF NOT EXISTS " + tx.Statement.Quote(database)
	if err = tx.Exec(sql).Error; err != nil {
		m.Log("❌ failed to create database '%s': %w", database, err)
		return
	}

	unlock, lockErr := m.acquireLock(tx, database)
	if lockErr != nil {
		m.Log("❌ failed to acquire advisory lock for tenant %d: %w", tenantId, lockErr)
		return lockErr
	}
	defer unlock()

	err = tx.Transaction(func(tx *gorm.DB) error {
		reset, err := UseDatabase(tx, database)
		if err != nil {
			m.Log("❌ failed to switch to tenant database %d: %w", tenantId, err)
			return err
		}
		defer reset()

		if err = tx.AutoMigrate(models...); err != nil {
			m.Log("❌ failed to migrate tables for tenant %d: %w", tenantId, err)
			return err
		}
		m.Log("✅ private tables migrated for tenant %d", tenantId)
		return nil
	})

	return
}

// MigrateSharedModels migrates the shared tables in the database.
func (m Migrator) MigrateSharedModels(models []any) error {
	m.Log("⏳ migrating public tables")

	if len(models) == 0 {
		return errors.New("no public tables to migrate")
	}

	db := m.DB.Session(&gorm.Session{})
	database := PublicDatabase()
	sql := "CREATE DATABASE IF NOT EXISTS " + db.Statement.Quote(database)
	if err := db.Exec(sql).Error; err != nil {
		m.Log("❌ failed to create public database '%s': %w", database, err)
		return err
	}

	unlock, lockErr := m.acquireLock(m.DB, database)
	if lockErr != nil {
		m.Log("❌ failed to acquire advisory lock: %w", lockErr)
		return lockErr
	}
	defer unlock()

	tx := db.Begin()
	defer func() {
		if tx.Error == nil {
			tx.Commit()
			m.Log("✅ public tables migrated")
		} else {
			tx.Rollback()
		}
	}()

	sql = "USE " + tx.Statement.Quote(database)
	if err := tx.Exec(sql).Error; err != nil {
		m.Log("❌ failed to switch to public database '%s': %w", database, err)
		return err
	}

	if err := tx.AutoMigrate(models...); err != nil {
		m.Log("❌ failed to migrate public tables: %w", err)
		return err
	}

	return nil
}

// DropDatabaseForTenant drops the database for a specific tenant.
func (m Migrator) DropDatabaseForTenant(tenantId uint) error {
	m.Log("⏳ dropping database for tenant %d", tenantId)

	tx := m.DB.Session(&gorm.Session{})
	database := TenantDatabase(tenantId)

	return m.retry(func() error {
		sql := "DROP DATABASE IF EXISTS " + tx.Statement.Quote(database)
		if err := tx.Exec(sql).Error; err != nil {
			m.Log("❌ failed to drop database '%s' for tenant %d: %w", database, tenantId, err)
			return err
		}
		m.Log("✅ database dropped for tenant %d", tenantId)
		return nil
	})
}