From 6b3e28216ca1975ecd01672b495bb0cb31383b63 Mon Sep 17 00:00:00 2001
From: Gerhard Schlager <gerhard.schlager@discourse.org>
Date: Mon, 16 Dec 2024 12:50:08 +0100
Subject: [PATCH] FEATURE: Allow pausing of restore before DB migration and
 uploads are restored (#30269)

This can be helpful if you need to fix problems in the DB before the DB gets migrated as well as before uploads are restored.
---
 lib/backup_restore/database_restorer.rb | 13 ++++++++++-
 lib/backup_restore/restorer.rb          | 30 ++++++++++++++++++++++---
 script/discourse                        | 10 ++++++++-
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/lib/backup_restore/database_restorer.rb b/lib/backup_restore/database_restorer.rb
index d8ce2dfb1ec..c210269f2ac 100644
--- a/lib/backup_restore/database_restorer.rb
+++ b/lib/backup_restore/database_restorer.rb
@@ -16,7 +16,7 @@ module BackupRestore
       @current_db = current_db
     end
 
-    def restore(db_dump_path)
+    def restore(db_dump_path, interactive = false)
       BackupRestore.move_tables_between_schemas(MAIN_SCHEMA, BACKUP_SCHEMA)
 
       @db_dump_path = db_dump_path
@@ -24,6 +24,7 @@ module BackupRestore
 
       create_missing_discourse_functions
       restore_dump
+      pause_before_migration if interactive
       migrate_database
       reconnect_database
 
@@ -136,6 +137,16 @@ module BackupRestore
       ].compact.join(" ")
     end
 
+    def pause_before_migration
+      puts ""
+      puts "Attention! Pausing restore before migrating database.".red.bold
+      puts "You can work on the restored database in a separate Rails console."
+      puts ""
+      puts "Press any key to continue with the restore.".bold
+      puts ""
+      STDIN.getch
+    end
+
     def migrate_database
       log "Migrating the database..."
 
diff --git a/lib/backup_restore/restorer.rb b/lib/backup_restore/restorer.rb
index 9aa8513d31d..8d0ae8a7d30 100644
--- a/lib/backup_restore/restorer.rb
+++ b/lib/backup_restore/restorer.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "colored2"
+
 module BackupRestore
   RestoreDisabledError = Class.new(RuntimeError)
   FilenameMissingError = Class.new(RuntimeError)
@@ -9,12 +11,20 @@ module BackupRestore
 
     attr_reader :success
 
-    def initialize(user_id:, filename:, factory:, disable_emails: true, location: nil)
+    def initialize(
+      user_id:,
+      filename:,
+      factory:,
+      disable_emails: true,
+      location: nil,
+      interactive: false
+    )
       @user_id = user_id
       @filename = filename
       @factory = factory
       @logger = factory.logger
       @disable_emails = disable_emails
+      @interactive = interactive
 
       ensure_restore_is_enabled
       ensure_we_have_a_user
@@ -48,7 +58,7 @@ module BackupRestore
       @system.flush_redis
       @system.clear_sidekiq_queues
 
-      @database_restorer.restore(db_dump_path)
+      @database_restorer.restore(db_dump_path, @interactive)
 
       reload_site_settings
 
@@ -58,7 +68,7 @@ module BackupRestore
       clear_stats
       reload_translations
 
-      @uploads_restorer.restore(@tmp_directory)
+      restore_uploads
 
       clear_emoji_cache
       clear_theme_cache
@@ -143,6 +153,20 @@ module BackupRestore
       TranslationOverride.reload_all_overrides!
     end
 
+    def restore_uploads
+      if @interactive
+        puts ""
+        puts "Attention! Pausing restore before uploads.".red.bold
+        puts "You can work on the restored database in a separate Rails console."
+        puts ""
+        puts "Press any key to continue with the restore.".bold
+        puts ""
+        STDIN.getch
+      end
+
+      @uploads_restorer.restore(@tmp_directory)
+    end
+
     def notify_user
       return if @success && @user_id == Discourse::SYSTEM_USER_ID
 
diff --git a/script/discourse b/script/discourse
index 6473bbddda7..c1246297acf 100755
--- a/script/discourse
+++ b/script/discourse
@@ -154,7 +154,14 @@ class DiscourseCLI < Thor
   end
 
   desc "restore", "Restore a Discourse backup"
-  option :disable_emails, type: :boolean, default: true
+  option :disable_emails,
+         type: :boolean,
+         default: true,
+         desc: "Disable outgoing emails for non-staff users after restore"
+  option :pause,
+         type: :boolean,
+         default: false,
+         desc: "Pause before migrating database and restoring uploads"
   option :location, type: :string, enum: %w[local s3], desc: "Override the backup location"
   def restore(filename = nil)
     load_rails
@@ -179,6 +186,7 @@ class DiscourseCLI < Thor
           disable_emails: options[:disable_emails],
           location: options[:location],
           factory: BackupRestore::Factory.new(user_id: Discourse.system_user.id),
+          interactive: options[:pause],
         )
       restorer.run
       puts "Restore done."