Backup database if a migration is needed (#415)

* Confirm before migrating database

Co-authored-by: WithoutPants <53250216+WithoutPants@users.noreply.github.com>

pkg/api/migrate.go

@@ -0,0 +1,66 @@
+package api
+
+import (
+	"fmt"
+	"html/template"
+	"net/http"
+
+	"github.com/stashapp/stash/pkg/database"
+)
+
+type migrateData struct {
+	ExistingVersion uint
+	MigrateVersion  uint
+	BackupPath      string
+}
+
+func getMigrateData() migrateData {
+	return migrateData{
+		ExistingVersion: database.Version(),
+		MigrateVersion:  database.AppSchemaVersion(),
+		BackupPath:      database.DatabaseBackupPath(),
+	}
+}
+
+func getMigrateHandler(w http.ResponseWriter, r *http.Request) {
+	if !database.NeedsMigration() {
+		http.Redirect(w, r, "/", 301)
+		return
+	}
+
+	data, _ := setupUIBox.Find("migrate.html")
+	templ, err := template.New("Migrate").Parse(string(data))
+	if err != nil {
+		http.Error(w, fmt.Sprintf("error: %s", err), 500)
+		return
+	}
+
+	err = templ.Execute(w, getMigrateData())
+	if err != nil {
+		http.Error(w, fmt.Sprintf("error: %s", err), 500)
+	}
+}
+
+func doMigrateHandler(w http.ResponseWriter, r *http.Request) {
+	err := r.ParseForm()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("error: %s", err), 500)
+	}
+
+	backupPath := r.Form.Get("backuppath")
+
+	// perform database backup
+	if backupPath != "" {
+		if err = database.Backup(backupPath); err != nil {
+			http.Error(w, fmt.Sprintf("error backing up database: %s", err), 500)
+			return
+		}
+	}
+
+	err = database.RunMigrations()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("error performing migration: %s", err), 500)
+		return
+	}
+	http.Redirect(w, r, "/", 301)
+}

pkg/api/server.go

@@ -20,6 +20,7 @@ import (
 	"github.com/gobuffalo/packr/v2"
 	"github.com/gorilla/websocket"
 	"github.com/rs/cors"
+	"github.com/stashapp/stash/pkg/database"
 	"github.com/stashapp/stash/pkg/logger"
 	"github.com/stashapp/stash/pkg/manager"
 	"github.com/stashapp/stash/pkg/manager/config"
@@ -83,6 +84,7 @@ func Start() {
 	r.Use(cors.AllowAll().Handler)
 	r.Use(BaseURLMiddleware)
 	r.Use(ConfigCheckMiddleware)
+	r.Use(DatabaseCheckMiddleware)
 
 	recoverFunc := handler.RecoverFunc(func(ctx context.Context, err interface{}) error {
 		logger.Error(err)
@@ -127,6 +129,10 @@ func Start() {
 		http.ServeFile(w, r, fn)
 	})
 
+	// Serve the migration UI
+	r.Get("/migrate", getMigrateHandler)
+	r.Post("/migrate", doMigrateHandler)
+
 	// Serve the setup UI
 	r.HandleFunc("/setup*", func(w http.ResponseWriter, r *http.Request) {
 		ext := path.Ext(r.URL.Path)
@@ -323,3 +329,17 @@ func ConfigCheckMiddleware(next http.Handler) http.Handler {
 		next.ServeHTTP(w, r)
 	})
 }
+
+func DatabaseCheckMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ext := path.Ext(r.URL.Path)
+		shouldRedirect := ext == "" && r.Method == "GET"
+		if shouldRedirect && database.NeedsMigration() {
+			if !strings.HasPrefix(r.URL.Path, "/migrate") {
+				http.Redirect(w, r, "/migrate", 301)
+				return
+			}
+		}
+		next.ServeHTTP(w, r)
+	})
+}

pkg/database/database.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"regexp"
+	"time"
 
 	"github.com/gobuffalo/packr/v2"
 	"github.com/golang-migrate/migrate/v4"
@@ -17,7 +18,9 @@ import (
 )
 
 var DB *sqlx.DB
+var dbPath string
 var appSchemaVersion uint = 4
+var databaseSchemaVersion uint
 
 const sqlite3Driver = "sqlite3_regexp"
 
@@ -27,7 +30,30 @@ func init() {
 }
 
 func Initialize(databasePath string) {
-	runMigrations(databasePath)
+	dbPath = databasePath
+
+	if err := getDatabaseSchemaVersion(); err != nil {
+		panic(err)
+	}
+
+	if databaseSchemaVersion == 0 {
+		// new database, just run the migrations
+		if err := RunMigrations(); err != nil {
+			panic(err)
+		}
+		// RunMigrations calls Initialise. Just return
+		return
+	} else {
+		if databaseSchemaVersion > appSchemaVersion {
+			panic(fmt.Sprintf("Database schema version %d is incompatible with required schema version %d", databaseSchemaVersion, appSchemaVersion))
+		}
+
+		// if migration is needed, then don't open the connection
+		if NeedsMigration() {
+			logger.Warnf("Database schema version %d does not match required schema version %d.", databaseSchemaVersion, appSchemaVersion)
+			return
+		}
+	}
 
 	// https://github.com/mattn/go-sqlite3
 	conn, err := sqlx.Open(sqlite3Driver, "file:"+databasePath+"?_fk=true")
@@ -55,34 +81,87 @@ func Reset(databasePath string) error {
 	return nil
 }
 
+// Backup the database
+func Backup(backupPath string) error {
+	db, err := sqlx.Connect(sqlite3Driver, "file:"+dbPath+"?_fk=true")
+	if err != nil {
+		return fmt.Errorf("Open database %s failed:%s", dbPath, err)
+	}
+	defer db.Close()
+
+	_, err = db.Exec(`VACUUM INTO "` + backupPath + `"`)
+	if err != nil {
+		return fmt.Errorf("Vacuum failed: %s", err)
+	}
+
+	return nil
+}
+
 // Migrate the database
-func runMigrations(databasePath string) {
+func NeedsMigration() bool {
+	return databaseSchemaVersion != appSchemaVersion
+}
+
+func AppSchemaVersion() uint {
+	return appSchemaVersion
+}
+
+func DatabaseBackupPath() string {
+	return fmt.Sprintf("%s.%d.%s", dbPath, databaseSchemaVersion, time.Now().Format("20060102_150405"))
+}
+
+func Version() uint {
+	return databaseSchemaVersion
+}
+
+func getMigrate() (*migrate.Migrate, error) {
 	migrationsBox := packr.New("Migrations Box", "./migrations")
 	packrSource := &Packr2Source{
 		Box:        migrationsBox,
 		Migrations: source.NewMigrations(),
 	}
 
-	databasePath = utils.FixWindowsPath(databasePath)
+	databasePath := utils.FixWindowsPath(dbPath)
 	s, _ := WithInstance(packrSource)
-	m, err := migrate.NewWithSourceInstance(
+	return migrate.NewWithSourceInstance(
 		"packr2",
 		s,
 		fmt.Sprintf("sqlite3://%s", "file:"+databasePath),
 	)
+}
+
+func getDatabaseSchemaVersion() error {
+	m, err := getMigrate()
+	if err != nil {
+		return err
+	}
+
+	databaseSchemaVersion, _, _ = m.Version()
+	m.Close()
+	return nil
+}
+
+// Migrate the database
+func RunMigrations() error {
+	m, err := getMigrate()
 	if err != nil {
 		panic(err.Error())
 	}
 
-	databaseSchemaVersion, _, _ := m.Version()
+	databaseSchemaVersion, _, _ = m.Version()
 	stepNumber := appSchemaVersion - databaseSchemaVersion
 	if stepNumber != 0 {
 		err = m.Steps(int(stepNumber))
 		if err != nil {
-			panic(err.Error())
+			return err
 		}
 	}
 	m.Close()
+
+	// re-initialise the database
+	Initialize(dbPath)
+
+	return nil
 }
 
 func registerRegexpFunc() {