new-api-custom/model/token_astra_migration.go

package model

import "gorm.io/gorm"

const astraAutoSwitchMigrationKey = "migration.token_astra_auto_switch.v1"

// MigrateTokenAstraAutoSwitch preserves the pre-token-policy behavior for
// tokens that existed before the user-level switch was introduced. It is
// guarded by an Option marker so newly-created tokens (which default to false)
// are never promoted on a later restart.
func MigrateTokenAstraAutoSwitch() error {
	return DB.Transaction(func(tx *gorm.DB) error {
		var marker Option
		if err := tx.Where(&Option{Key: astraAutoSwitchMigrationKey}).First(&marker).Error; err == nil {
			return nil
		} else if err != gorm.ErrRecordNotFound {
			return err
		}
		if err := tx.Model(&Token{}).Where("deleted_at IS NULL").Update("astra_auto_switch", true).Error; err != nil {
			return err
		}
		return tx.Create(&Option{Key: astraAutoSwitchMigrationKey, Value: "true"}).Error
	})
}

// MigrateTokenAstraDetection preserves observation for existing tokens without
// resetting an explicit opt-out on subsequent starts.
func MigrateTokenAstraDetection() error {
	return DB.Transaction(func(tx *gorm.DB) error {
		const key = "migration.token_astra_detection.v1"
		var marker Option
		if err := tx.Where(&Option{Key: key}).First(&marker).Error; err == nil {
			return nil
		} else if err != gorm.ErrRecordNotFound {
			return err
		}
		if err := tx.Model(&Token{}).Where("deleted_at IS NULL").Update("astra_detection_enabled", true).Error; err != nil {
			return err
		}
		return tx.Create(&Option{Key: key, Value: "true"}).Error
	})
}