fix(rollback): reimplement rollback feature [EE-6367] (#10721)

api/cli/confirm.go

@@ -9,7 +9,7 @@ import (
 
 // Confirm starts a rollback db cli application
 func Confirm(message string) (bool, error) {
-	fmt.Printf("%s [y/N]", message)
+	fmt.Printf("%s [y/N] ", message)
 
 	reader := bufio.NewReader(os.Stdin)
 

api/datastore/backup.go

@@ -1,187 +1,77 @@
 package datastore
 
 import (
-	"fmt"
 	"os"
 	"path"
-	"time"
 
-	"github.com/portainer/portainer/api/database/models"
 	"github.com/rs/zerolog/log"
 )
 
-var backupDefaults = struct {
+func (store *Store) Backup() (string, error) {
-	backupDir string
+	if err := store.createBackupPath(); err != nil {
-	commonDir string
+		return "", err
-}{
+	}
-	"backups",
+
-	"common",
+	backupFilename := store.backupFilename()
+	log.Info().Str("from", store.connection.GetDatabaseFilePath()).Str("to", backupFilename).Msgf("Backing up database")
+	err := store.fileService.Copy(store.connection.GetDatabaseFilePath(), backupFilename, true)
+	if err != nil {
+		log.Warn().Err(err).Msg("failed to create backup file")
+		return "", err
+	}
+
+	return backupFilename, nil
 }
 
-//
+func (store *Store) Restore() error {
-// Backup Helpers
+	backupFilename := store.backupFilename()
-//
+	return store.RestoreFromFile(backupFilename)
+}
 
-// createBackupFolders create initial folders for backups
+func (store *Store) RestoreFromFile(backupFilename string) error {
-func (store *Store) createBackupFolders() {
+	if exists, _ := store.fileService.FileExists(backupFilename); !exists {
-	// create common dir
+		log.Error().Str("backupFilename", backupFilename).Msg("backup file does not exist")
-	commonDir := store.commonBackupDir()
+		return os.ErrNotExist
-	if exists, _ := store.fileService.FileExists(commonDir); !exists {
+	}
-		if err := os.MkdirAll(commonDir, 0700); err != nil {
+
-			log.Error().Err(err).Msg("error while creating common backup folder")
+	if err := store.fileService.Copy(backupFilename, store.connection.GetDatabaseFilePath(), true); err != nil {
+		log.Error().Err(err).Msg("error while restoring backup.")
+		return err
+	}
+
+	log.Info().Str("from", store.connection.GetDatabaseFilePath()).Str("to", backupFilename).Msgf("database restored")
+
+	// determine the db version
+	store.Open()
+	version, err := store.VersionService.Version()
+
+	edition := "CE"
+	if version.Edition == 2 {
+		edition = "EE"
+	}
+
+	if err == nil {
+		log.Info().Str("version", version.SchemaVersion).Msgf("Restored database version: Portainer %s %s", edition, version.SchemaVersion)
+	}
+
+	return nil
+}
+
+func (store *Store) createBackupPath() error {
+	backupDir := path.Join(store.connection.GetStorePath(), "backups")
+	if exists, _ := store.fileService.FileExists(backupDir); !exists {
+		if err := os.MkdirAll(backupDir, 0700); err != nil {
+			log.Error().Err(err).Msg("error while creating backup folder")
+			return err
 		}
 	}
+	return nil
+}
+
+func (store *Store) backupFilename() string {
+	return path.Join(store.connection.GetStorePath(), "backups", store.connection.GetDatabaseFileName()+".bak")
 }
 
 func (store *Store) databasePath() string {
 	return store.connection.GetDatabaseFilePath()
 }
-
-func (store *Store) commonBackupDir() string {
-	return path.Join(store.connection.GetStorePath(), backupDefaults.backupDir, backupDefaults.commonDir)
-}
-
-func (store *Store) copyDBFile(from string, to string) error {
-	log.Info().Str("from", from).Str("to", to).Msg("copying DB file")
-
-	err := store.fileService.Copy(from, to, true)
-	if err != nil {
-		log.Error().Err(err).Msg("failed")
-	}
-
-	return err
-}
-
-// BackupOptions provide a helper to inject backup options
-type BackupOptions struct {
-	Version        string
-	BackupDir      string
-	BackupFileName string
-	BackupPath     string
-}
-
-// getBackupRestoreOptions returns options to store db at common backup dir location; used by:
-// - db backup prior to version upgrade
-// - db rollback
-func getBackupRestoreOptions(backupDir string) *BackupOptions {
-	return &BackupOptions{
-		BackupDir:      backupDir,
-		BackupFileName: beforePortainerVersionUpgradeBackup,
-	}
-}
-
-// Backup current database with default options
-func (store *Store) Backup(version *models.Version) (string, error) {
-	if version == nil {
-		return store.backupWithOptions(nil)
-	}
-
-	backupOptions := getBackupRestoreOptions(store.commonBackupDir())
-	backupOptions.Version = version.SchemaVersion
-	return store.backupWithOptions(backupOptions)
-}
-
-func (store *Store) setDefaultBackupOptions(options *BackupOptions) *BackupOptions {
-	if options == nil {
-		options = &BackupOptions{}
-	}
-	if options.Version == "" {
-		v, err := store.VersionService.Version()
-		if err != nil {
-			options.Version = ""
-		}
-		options.Version = v.SchemaVersion
-	}
-	if options.BackupDir == "" {
-		options.BackupDir = store.commonBackupDir()
-	}
-	if options.BackupFileName == "" {
-		options.BackupFileName = fmt.Sprintf("%s.%s.%s", store.connection.GetDatabaseFileName(), options.Version, time.Now().Format("20060102150405"))
-	}
-	if options.BackupPath == "" {
-		options.BackupPath = path.Join(options.BackupDir, options.BackupFileName)
-	}
-	return options
-}
-
-// BackupWithOptions backup current database with options
-func (store *Store) backupWithOptions(options *BackupOptions) (string, error) {
-	log.Info().Msg("creating DB backup")
-
-	store.createBackupFolders()
-
-	options = store.setDefaultBackupOptions(options)
-	dbPath := store.databasePath()
-
-	if err := store.Close(); err != nil {
-		return options.BackupPath, fmt.Errorf(
-			"error closing datastore before creating backup: %w",
-			err,
-		)
-	}
-
-	if err := store.copyDBFile(dbPath, options.BackupPath); err != nil {
-		return options.BackupPath, err
-	}
-
-	if _, err := store.Open(); err != nil {
-		return options.BackupPath, fmt.Errorf(
-			"error opening datastore after creating backup: %w",
-			err,
-		)
-	}
-
-	return options.BackupPath, nil
-}
-
-// RestoreWithOptions previously saved backup for the current Edition  with options
-// Restore strategies:
-// - default: restore latest from current edition
-// - restore a specific
-func (store *Store) restoreWithOptions(options *BackupOptions) error {
-	options = store.setDefaultBackupOptions(options)
-
-	// Check if backup file exist before restoring
-	_, err := os.Stat(options.BackupPath)
-	if os.IsNotExist(err) {
-		log.Error().Str("path", options.BackupPath).Err(err).Msg("backup file to restore does not exist")
-		return err
-	}
-
-	err = store.Close()
-	if err != nil {
-		log.Error().Err(err).Msg("error while closing store before restore")
-		return err
-	}
-
-	log.Info().Msg("restoring DB backup")
-	err = store.copyDBFile(options.BackupPath, store.databasePath())
-	if err != nil {
-		return err
-	}
-
-	_, err = store.Open()
-	return err
-}
-
-// RemoveWithOptions removes backup database based on supplied options
-func (store *Store) removeWithOptions(options *BackupOptions) error {
-	log.Info().Msg("removing DB backup")
-
-	options = store.setDefaultBackupOptions(options)
-	_, err := os.Stat(options.BackupPath)
-
-	if os.IsNotExist(err) {
-		log.Error().Str("path", options.BackupPath).Err(err).Msg("backup file to remove does not exist")
-		return err
-	}
-
-	log.Info().Str("path", options.BackupPath).Msg("removing DB file")
-	err = os.Remove(options.BackupPath)
-	if err != nil {
-		log.Error().Err(err).Msg("failed")
-		return err
-	}
-
-	return nil
-}

api/datastore/backup_test.go

@@ -2,106 +2,79 @@ package datastore
 
 import (
 	"fmt"
-	"os"
-	"path"
 	"testing"
 
 	portainer "github.com/portainer/portainer/api"
 	"github.com/portainer/portainer/api/database/models"
+
+	"github.com/rs/zerolog/log"
 )
 
-func TestCreateBackupFolders(t *testing.T) {
-	_, store := MustNewTestStore(t, true, true)
-
-	connection := store.GetConnection()
-	backupPath := path.Join(connection.GetStorePath(), backupDefaults.backupDir)
-
-	if isFileExist(backupPath) {
-		t.Error("Expect backups folder to not exist")
-	}
-
-	store.createBackupFolders()
-	if !isFileExist(backupPath) {
-		t.Error("Expect backups folder to exist")
-	}
-}
-
 func TestStoreCreation(t *testing.T) {
 	_, store := MustNewTestStore(t, true, true)
 	if store == nil {
-		t.Error("Expect to create a store")
+		t.Fatal("Expect to create a store")
 	}
 
-	if store.CheckCurrentEdition() != nil {
+	v, err := store.VersionService.Version()
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+
+	if portainer.SoftwareEdition(v.Edition) != portainer.PortainerCE {
 		t.Error("Expect to get CE Edition")
 	}
+
+	if v.SchemaVersion != portainer.APIVersion {
+		t.Error("Expect to get APIVersion")
+	}
 }
 
 func TestBackup(t *testing.T) {
 	_, store := MustNewTestStore(t, true, true)
-	connection := store.GetConnection()
+	backupFileName := store.backupFilename()
-
+	t.Run(fmt.Sprintf("Backup should create %s", backupFileName), func(t *testing.T) {
-	t.Run("Backup should create default db backup", func(t *testing.T) {
 		v := models.Version{
+			Edition:       int(portainer.PortainerCE),
 			SchemaVersion: portainer.APIVersion,
 		}
 		store.VersionService.UpdateVersion(&v)
-		store.backupWithOptions(nil)
+		store.Backup()
 
-		backupFileName := path.Join(connection.GetStorePath(), "backups", "common", fmt.Sprintf("portainer.edb.%s.*", portainer.APIVersion))
-		if !isFileExist(backupFileName) {
-			t.Errorf("Expect backup file to be created %s", backupFileName)
-		}
-	})
-
-	t.Run("BackupWithOption should create a name specific backup at common path", func(t *testing.T) {
-		store.backupWithOptions(&BackupOptions{
-			BackupFileName: beforePortainerVersionUpgradeBackup,
-			BackupDir:      store.commonBackupDir(),
-		})
-		backupFileName := path.Join(connection.GetStorePath(), "backups", "common", beforePortainerVersionUpgradeBackup)
 		if !isFileExist(backupFileName) {
 			t.Errorf("Expect backup file to be created %s", backupFileName)
 		}
 	})
 }
 
-func TestRemoveWithOptions(t *testing.T) {
+func TestRestore(t *testing.T) {
-	_, store := MustNewTestStore(t, true, true)
+	_, store := MustNewTestStore(t, true, false)
 
-	t.Run("successfully removes file if existent", func(t *testing.T) {
+	t.Run(fmt.Sprintf("Basic Restore"), func(t *testing.T) {
-		store.createBackupFolders()
+		// override and set initial db version and edition
-		options := &BackupOptions{
+		updateEdition(store, portainer.PortainerCE)
-			BackupDir:      store.commonBackupDir(),
+		updateVersion(store, "2.4")
-			BackupFileName: "test.txt",
-		}
 
-		filePath := path.Join(options.BackupDir, options.BackupFileName)
+		store.Backup()
-		f, err := os.Create(filePath)
+		updateVersion(store, "2.16")
-		if err != nil {
+		testVersion(store, "2.16", t)
-			t.Fatalf("file should be created; err=%s", err)
+		store.Restore()
-		}
-		f.Close()
 
-		err = store.removeWithOptions(options)
+		// check if the restore is successful and the version is correct
-		if err != nil {
+		testVersion(store, "2.4", t)
-			t.Errorf("RemoveWithOptions should successfully remove file; err=%v", err)
-		}
-
-		if isFileExist(f.Name()) {
-			t.Errorf("RemoveWithOptions should successfully remove file; file=%s", f.Name())
-		}
 	})
 
-	t.Run("fails to removes file if non-existent", func(t *testing.T) {
+	t.Run(fmt.Sprintf("Basic Restore After Multiple Backups"), func(t *testing.T) {
-		options := &BackupOptions{
+		// override and set initial db version and edition
-			BackupDir:      store.commonBackupDir(),
+		updateEdition(store, portainer.PortainerCE)
-			BackupFileName: "test.txt",
+		updateVersion(store, "2.4")
-		}
+		store.Backup()
+		updateVersion(store, "2.14")
+		updateVersion(store, "2.16")
+		testVersion(store, "2.16", t)
+		store.Restore()
 
-		err := store.removeWithOptions(options)
+		// check if the restore is successful and the version is correct
-		if err == nil {
+		testVersion(store, "2.4", t)
-			t.Error("RemoveWithOptions should fail for non-existent file")
-		}
 	})
 }

api/datastore/datastore.go

@@ -31,8 +31,14 @@ func (store *Store) Open() (newStore bool, err error) {
 	}
 
 	if encryptionReq {
+		backupFilename, err := store.Backup()
+		if err != nil {
+			return false, fmt.Errorf("failed to backup database prior to encrypting: %w", err)
+		}
+
 		err = store.encryptDB()
 		if err != nil {
+			store.RestoreFromFile(backupFilename) // restore from backup if encryption fails
 			return false, err
 		}
 	}

api/datastore/helpers_test.go

@@ -0,0 +1,68 @@
+package datastore
+
+import (
+	"path/filepath"
+	"testing"
+
+	portainer "github.com/portainer/portainer/api"
+
+	"github.com/rs/zerolog/log"
+)
+
+// isFileExist is helper function to check for file existence
+func isFileExist(path string) bool {
+	matches, err := filepath.Glob(path)
+	if err != nil {
+		return false
+	}
+	return len(matches) > 0
+}
+
+func updateVersion(store *Store, v string) {
+	version, err := store.VersionService.Version()
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+
+	version.SchemaVersion = v
+
+	err = store.VersionService.UpdateVersion(version)
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+}
+
+func updateEdition(store *Store, edition portainer.SoftwareEdition) {
+	version, err := store.VersionService.Version()
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+
+	version.Edition = int(edition)
+
+	err = store.VersionService.UpdateVersion(version)
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+}
+
+// testVersion is a helper which tests current store version against wanted version
+func testVersion(store *Store, versionWant string, t *testing.T) {
+	v, err := store.VersionService.Version()
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+	if v.SchemaVersion != versionWant {
+		t.Errorf("Expect store version to be %s but was %s", versionWant, v.SchemaVersion)
+	}
+}
+
+func testEdition(store *Store, editionWant portainer.SoftwareEdition, t *testing.T) {
+	v, err := store.VersionService.Version()
+	if err != nil {
+		log.Fatal().Err(err).Msg("")
+	}
+	if portainer.SoftwareEdition(v.Edition) != editionWant {
+		t.Errorf("Expect store edition to be %s but was %s", editionWant.GetEditionLabel(), portainer.SoftwareEdition(v.Edition).GetEditionLabel())
+	}
+}

api/datastore/migrate_data.go

@@ -16,8 +16,6 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-const beforePortainerVersionUpgradeBackup = "portainer.db.bak"
-
 func (store *Store) MigrateData() error {
 	updating, err := store.VersionService.IsUpdating()
 	if err != nil {
@@ -42,7 +40,7 @@ func (store *Store) MigrateData() error {
 	}
 
 	// before we alter anything in the DB, create a backup
-	backupPath, err := store.Backup(version)
+	_, err = store.Backup()
 	if err != nil {
 		return errors.Wrap(err, "while backing up database")
 	}
@@ -52,9 +50,9 @@ func (store *Store) MigrateData() error {
 		err = errors.Wrap(err, "failed to migrate database")
 
 		log.Warn().Err(err).Msg("migration failed, restoring database to previous version")
-		restorErr := store.restoreWithOptions(&BackupOptions{BackupPath: backupPath})
+		restoreErr := store.Restore()
-		if restorErr != nil {
+		if restoreErr != nil {
-			return errors.Wrap(restorErr, "failed to restore database")
+			return errors.Wrap(restoreErr, "failed to restore database")
 		}
 
 		log.Info().Msg("database restored to previous version")
@@ -141,9 +139,7 @@ func (store *Store) connectionRollback(force bool) error {
 		}
 	}
 
-	options := getBackupRestoreOptions(store.commonBackupDir())
+	err := store.Restore()
-
-	err := store.restoreWithOptions(options)
 	if err != nil {
 		return err
 	}

api/datastore/migrate_data_test.go

@@ -2,35 +2,25 @@ package datastore
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 
+	"github.com/Masterminds/semver"
 	portainer "github.com/portainer/portainer/api"
 	"github.com/portainer/portainer/api/database/boltdb"
 	"github.com/portainer/portainer/api/database/models"
+	"github.com/portainer/portainer/api/datastore/migrator"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/rs/zerolog/log"
-	"github.com/segmentio/encoding/json"
 )
 
-// testVersion is a helper which tests current store version against wanted version
-func testVersion(store *Store, versionWant string, t *testing.T) {
-	v, err := store.VersionService.Version()
-	if err != nil {
-		t.Errorf("Expect store version to be %s but was %s with error: %s", versionWant, v.SchemaVersion, err)
-	}
-	if v.SchemaVersion != versionWant {
-		t.Errorf("Expect store version to be %s but was %s", versionWant, v.SchemaVersion)
-	}
-}
-
 func TestMigrateData(t *testing.T) {
-	snapshotTests := []struct {
+	tests := []struct {
 		testName           string
 		srcPath            string
 		wantPath           string
@@ -43,7 +33,7 @@ func TestMigrateData(t *testing.T) {
 			overrideInstanceId: true,
 		},
 	}
-	for _, test := range snapshotTests {
+	for _, test := range tests {
 		t.Run(test.testName, func(t *testing.T) {
 			err := migrateDBTestHelper(t, test.srcPath, test.wantPath, test.overrideInstanceId)
 			if err != nil {
@@ -58,7 +48,6 @@ func TestMigrateData(t *testing.T) {
 
 	t.Run("MigrateData for New Store & Re-Open Check", func(t *testing.T) {
 		newStore, store := MustNewTestStore(t, true, false)
-
 		if !newStore {
 			t.Error("Expect a new DB")
 		}
@@ -72,75 +61,14 @@ func TestMigrateData(t *testing.T) {
 		}
 	})
 
-	tests := []struct {
-		version         string
-		expectedVersion string
-	}{
-		{version: "1.24.1", expectedVersion: portainer.APIVersion},
-		{version: "2.0.0", expectedVersion: portainer.APIVersion},
-	}
-	for _, tc := range tests {
-		_, store := MustNewTestStore(t, true, true)
-
-		// Setup data
-		v := models.Version{SchemaVersion: tc.version, Edition: int(portainer.PortainerCE)}
-		store.VersionService.UpdateVersion(&v)
-
-		// Required roles by migrations 22.2
-		store.RoleService.Create(&portainer.Role{ID: 1})
-		store.RoleService.Create(&portainer.Role{ID: 2})
-		store.RoleService.Create(&portainer.Role{ID: 3})
-		store.RoleService.Create(&portainer.Role{ID: 4})
-
-		t.Run(fmt.Sprintf("MigrateData for version %s", tc.version), func(t *testing.T) {
-			store.MigrateData()
-			testVersion(store, tc.expectedVersion, t)
-		})
-
-		t.Run(fmt.Sprintf("Restoring DB after migrateData for version %s", tc.version), func(t *testing.T) {
-			store.Rollback(true)
-			store.Open()
-			testVersion(store, tc.version, t)
-		})
-	}
-
-	t.Run("Error in MigrateData should restore backup before MigrateData", func(t *testing.T) {
-		_, store := MustNewTestStore(t, false, false)
-
-		v := models.Version{SchemaVersion: "1.24.1", Edition: int(portainer.PortainerCE)}
-		store.VersionService.UpdateVersion(&v)
-
-		store.MigrateData()
-
-		testVersion(store, v.SchemaVersion, t)
-	})
-
 	t.Run("MigrateData should create backup file upon update", func(t *testing.T) {
-		_, store := MustNewTestStore(t, false, false)
+		_, store := MustNewTestStore(t, true, false)
-
+		store.VersionService.UpdateVersion(&models.Version{SchemaVersion: "1.0", Edition: int(portainer.PortainerCE)})
-		v := models.Version{SchemaVersion: "0.0.0", Edition: int(portainer.PortainerCE)}
-		store.VersionService.UpdateVersion(&v)
-
 		store.MigrateData()
 
-		options := store.setDefaultBackupOptions(getBackupRestoreOptions(store.commonBackupDir()))
+		backupfilename := store.backupFilename()
-
+		if exists, _ := store.fileService.FileExists(backupfilename); !exists {
-		if !isFileExist(options.BackupPath) {
+			t.Errorf("Expect backup file to be created %s", backupfilename)
-			t.Errorf("Backup file should exist; file=%s", options.BackupPath)
-		}
-	})
-
-	t.Run("MigrateData should fail to create backup if database file is set to updating", func(t *testing.T) {
-		_, store := MustNewTestStore(t, false, false)
-
-		store.VersionService.StoreIsUpdating(true)
-
-		store.MigrateData()
-
-		options := store.setDefaultBackupOptions(getBackupRestoreOptions(store.commonBackupDir()))
-
-		if isFileExist(options.BackupPath) {
-			t.Errorf("Backup file should not exist for dirty database; file=%s", options.BackupPath)
 		}
 	})
 
@@ -150,50 +78,101 @@ func TestMigrateData(t *testing.T) {
 		version := "2.15"
 		_, store := MustNewTestStore(t, true, false)
 		store.VersionService.UpdateVersion(&models.Version{SchemaVersion: version, Edition: int(portainer.PortainerCE)})
-		err := store.MigrateData()
+		store.MigrateData()
-		if err == nil {
-			t.Errorf("Expect migration to fail")
-		}
 
 		store.Open()
 		testVersion(store, version, t)
 	})
-}
 
-func Test_getBackupRestoreOptions(t *testing.T) {
+	t.Run("MigrateData should fail to create backup if database file is set to updating", func(t *testing.T) {
-	_, store := MustNewTestStore(t, false, true)
+		_, store := MustNewTestStore(t, true, false)
+		store.VersionService.StoreIsUpdating(true)
+		store.MigrateData()
 
-	options := getBackupRestoreOptions(store.commonBackupDir())
+		// If you get an error, it usually means that the backup folder doesn't exist (no backups). Expected!
+		// If the backup file is not blank, then it means a backup was created.  We don't want that because we
+		// only create a backup when the version changes.
+		backupfilename := store.backupFilename()
+		if exists, _ := store.fileService.FileExists(backupfilename); exists {
+			t.Errorf("Backup file should not exist for dirty database")
+		}
+	})
 
-	wantDir := store.commonBackupDir()
+	t.Run("MigrateData should not create backup on startup if portainer version matches db", func(t *testing.T) {
-	if !strings.HasSuffix(options.BackupDir, wantDir) {
+		_, store := MustNewTestStore(t, true, false)
-		log.Fatal().Str("got", options.BackupDir).Str("want", wantDir).Msg("incorrect backup dir")
-	}
 
-	wantFilename := "portainer.db.bak"
+		// Set migrator the count to match our migrations array (simulate no changes).
-	if options.BackupFileName != wantFilename {
+		// Should not create a backup
-		log.Fatal().Str("got", options.BackupFileName).Str("want", wantFilename).Msg("incorrect backup file")
+		v, err := store.VersionService.Version()
-	}
+		if err != nil {
+			t.Errorf("Unable to read version from db: %s", err)
+			t.FailNow()
+		}
+
+		migratorParams := store.newMigratorParameters(v)
+		m := migrator.NewMigrator(migratorParams)
+		latestMigrations := m.LatestMigrations()
+
+		if latestMigrations.Version.Equal(semver.MustParse(portainer.APIVersion)) {
+			v.MigratorCount = len(latestMigrations.MigrationFuncs)
+			store.VersionService.UpdateVersion(v)
+		}
+
+		store.MigrateData()
+
+		// If you get an error, it usually means that the backup folder doesn't exist (no backups). Expected!
+		// If the backup file is not blank, then it means a backup was created.  We don't want that because we
+		// only create a backup when the version changes.
+		backupfilename := store.backupFilename()
+		if exists, _ := store.fileService.FileExists(backupfilename); exists {
+			t.Errorf("Backup file should not exist for dirty database")
+		}
+	})
+
+	t.Run("MigrateData should create backup on startup if portainer version matches db and migrationFuncs counts differ", func(t *testing.T) {
+		_, store := MustNewTestStore(t, true, false)
+
+		// Set migrator count very large to simulate changes
+		// Should not create a backup
+		v, err := store.VersionService.Version()
+		if err != nil {
+			t.Errorf("Unable to read version from db: %s", err)
+			t.FailNow()
+		}
+
+		v.MigratorCount = 1000
+		store.VersionService.UpdateVersion(v)
+		store.MigrateData()
+
+		// If you get an error, it usually means that the backup folder doesn't exist (no backups). Expected!
+		// If the backup file is not blank, then it means a backup was created.  We don't want that because we
+		// only create a backup when the version changes.
+		backupfilename := store.backupFilename()
+		if exists, _ := store.fileService.FileExists(backupfilename); !exists {
+			t.Errorf("DB backup should exist and there should be no error")
+		}
+	})
 }
 
 func TestRollback(t *testing.T) {
 	t.Run("Rollback should restore upgrade after backup", func(t *testing.T) {
-		version := models.Version{SchemaVersion: "2.4.0"}
+		version := "2.11"
-		_, store := MustNewTestStore(t, true, false)
 
-		err := store.VersionService.UpdateVersion(&version)
+		v := models.Version{
-		if err != nil {
+			SchemaVersion: version,
-			t.Errorf("Failed updating version: %v", err)
 		}
 
-		_, err = store.backupWithOptions(getBackupRestoreOptions(store.commonBackupDir()))
+		_, store := MustNewTestStore(t, false, false)
+		store.VersionService.UpdateVersion(&v)
+
+		_, err := store.Backup()
 		if err != nil {
 			log.Fatal().Err(err).Msg("")
 		}
 
-		// Change the current version
+		v.SchemaVersion = "2.14"
-		version2 := models.Version{SchemaVersion: "2.6.0"}
+		// Change the current edition
-		err = store.VersionService.UpdateVersion(&version2)
+		err = store.VersionService.UpdateVersion(&v)
 		if err != nil {
 			log.Fatal().Err(err).Msg("")
 		}
@@ -205,26 +184,45 @@ func TestRollback(t *testing.T) {
 			return
 		}
 
-		_, err = store.Open()
+		store.Open()
+		testVersion(store, version, t)
+	})
+
+	t.Run("Rollback should restore upgrade after backup", func(t *testing.T) {
+		version := "2.15"
+
+		v := models.Version{
+			SchemaVersion: version,
+			Edition:       int(portainer.PortainerCE),
+		}
+
+		_, store := MustNewTestStore(t, true, false)
+		store.VersionService.UpdateVersion(&v)
+
+		_, err := store.Backup()
 		if err != nil {
-			t.Logf("Open failed: %s", err)
+			log.Fatal().Err(err).Msg("")
+		}
+
+		v.SchemaVersion = "2.14"
+		// Change the current edition
+		err = store.VersionService.UpdateVersion(&v)
+		if err != nil {
+			log.Fatal().Err(err).Msg("")
+		}
+
+		err = store.Rollback(true)
+		if err != nil {
+			t.Logf("Rollback failed: %s", err)
 			t.Fail()
 			return
 		}
 
-		testVersion(store, version.SchemaVersion, t)
+		store.Open()
+		testVersion(store, version, t)
 	})
 }
 
-// isFileExist is helper function to check for file existence
-func isFileExist(path string) bool {
-	matches, err := filepath.Glob(path)
-	if err != nil {
-		return false
-	}
-	return len(matches) > 0
-}
-
 // migrateDBTestHelper loads a json representation of a bolt database from srcPath,
 // parses it into a database, runs a migration on that database, and then
 // compares it with an expected output database.