Dropdb should be forced to close any active connections and not fail

.ruby-version

@@ -1 +1 @@
-2.7.2
+3.4.1

lib/parity/backup.rb

@@ -11,6 +11,7 @@ def initialize(args)
       @from, @to = args.values_at(:from, :to)
       @additional_args = args[:additional_args] || BLANK_ARGUMENTS
       @parallelize = args[:parallelize] || false
+      @backup_id = args[:backup_id]
     end
 
     def restore
@@ -25,31 +26,44 @@ def restore
 
     private
 
-    attr_reader :additional_args, :from, :to, :parallelize
+    attr_reader :additional_args, :from, :to, :parallelize, :backup_id
 
-    alias :parallelize? :parallelize
+    alias_method :parallelize?, :parallelize
+
+    def log_restore_info
+      if backup_id
+        puts "Restoring from #{from} backup ID: #{backup_id} to #{to}"
+      else
+        puts "Restoring from #{from} (latest backup) to #{to}"
+      end
+      puts "Starting backup restoration process..."
+    end
 
     def restore_from_development
+      log_restore_info
       reset_remote_database
       Kernel.system(
-        "heroku pg:push #{development_db} DATABASE_URL --remote #{to} "\
-          "#{additional_args}",
+        "heroku pg:push #{development_db} DATABASE_URL --remote #{to} " \
+          "#{additional_args}"
       )
+      puts "Backup restoration to #{to} completed successfully!"
     end
 
     def restore_to_development
+      log_restore_info
       ensure_temp_directory_exists
       download_remote_backup
       wipe_development_database
       create_heroku_ext_schema
       restore_from_local_temp_backup
       delete_local_temp_backup
       delete_rails_production_environment_settings
+      puts "Backup restoration to #{to} completed successfully!"
     end
 
     def wipe_development_database
       Kernel.system(
-        "dropdb --if-exists #{development_db} && createdb #{development_db}",
+        "dropdb --if-exists #{development_db} --force && createdb #{development_db}"
       )
     end
 
@@ -63,8 +77,8 @@ def create_heroku_ext_schema
 
     def reset_remote_database
       Kernel.system(
-        "heroku pg:reset --remote #{to} #{additional_args} "\
-          "--confirm #{heroku_app_name}",
+        "heroku pg:reset --remote #{to} #{additional_args} " \
+          "--confirm #{heroku_app_name}"
       )
     end
 
@@ -77,21 +91,34 @@ def ensure_temp_directory_exists
     end
 
     def download_remote_backup
-      Kernel.system(
-        "curl -o tmp/latest.backup \"$(heroku pg:backups:url --remote #{from})\"",
-      )
+      if backup_id
+        puts "Downloading backup #{backup_id} from #{from}..."
+        Kernel.system(
+          "curl -o tmp/#{backup_id}.backup \"$(heroku pg:backups:url #{backup_id} --remote #{from})\""
+        )
+      else
+        puts "Downloading latest backup from #{from}..."
+        Kernel.system(
+          "curl -o tmp/latest.backup \"$(heroku pg:backups:url --remote #{from})\""
+        )
+      end
     end
 
     def restore_from_local_temp_backup
+      puts "Restoring backup to #{development_db}..."
+      # Filter out --backup-id from additional_args as it's not needed for pg_restore
+      filtered_args = additional_args.gsub(/--backup-id\s+\S+/, "").strip
+      backup_filename = backup_id ? "#{backup_id}.backup" : "latest.backup"
       Kernel.system(
-        "pg_restore tmp/latest.backup --verbose --no-acl --no-owner "\
-          "--dbname #{development_db} --jobs=#{processor_cores} "\
-          "#{additional_args}",
+        "pg_restore tmp/#{backup_filename} --verbose --no-acl --no-owner " \
+          "--dbname #{development_db} --jobs=#{processor_cores} " \
+          "#{filtered_args}"
       )
     end
 
     def delete_local_temp_backup
-      Kernel.system("rm tmp/latest.backup")
+      backup_filename = backup_id ? "#{backup_id}.backup" : "latest.backup"
+      Kernel.system("rm tmp/#{backup_filename}")
     end
 
     def delete_rails_production_environment_settings
@@ -101,28 +128,41 @@ def delete_rails_production_environment_settings
     end
 
     def restore_to_remote_environment
+      log_restore_info
       reset_remote_database
+      # Filter out --backup-id from additional_args as it's handled separately
+      filtered_args = additional_args.gsub(/--backup-id\s+\S+/, "").strip
       Kernel.system(
-        "heroku pg:backups:restore #{backup_from} --remote #{to} "\
-          "#{additional_args}",
+        "heroku pg:backups:restore #{backup_from} --remote #{to} " \
+          "#{filtered_args}"
       )
+      puts "Backup restoration to #{to} completed successfully!"
     end
 
     def backup_from
       "`#{remote_db_backup_url}` DATABASE"
     end
 
     def remote_db_backup_url
-      "heroku pg:backups:url --remote #{from}"
+      if backup_id
+        "heroku pg:backups:url #{backup_id} --remote #{from}"
+      else
+        "heroku pg:backups:url --remote #{from}"
+      end
     end
 
     def development_db
-      YAML.safe_load(database_yaml_file, aliases: true).
-        fetch(DEVELOPMENT_ENVIRONMENT_KEY_NAME).
-        fetch(DATABASE_KEY_NAME)
+      YAML.safe_load(database_yaml_file, aliases: true)
+        .fetch(DEVELOPMENT_ENVIRONMENT_KEY_NAME)
+        .fetch(DATABASE_KEY_NAME)
     end
 
     def database_yaml_file
+      # Load Rails environment if Rails is not already loaded
+      # This is needed when database.yml uses Rails.application.credentials
+      unless defined?(Rails)
+        require File.expand_path("config/environment", Dir.pwd)
+      end
       ERB.new(IO.read(DATABASE_YML_RELATIVE_PATH)).result
     end
 

lib/parity/environment.rb

@@ -67,12 +67,15 @@ def restore
         $stdout.puts "Parity does not support restoring backups into your "\
           "production environment. Use `--force` to override."
       else
-        Backup.new(
+        backup_args = {
           from: arguments.first,
           to: environment,
           parallelize: parallelize?,
           additional_args: additional_restore_arguments,
-        ).restore
+        }
+        backup_args[:backup_id] = backup_id? if backup_id?
+
+        Backup.new(backup_args).restore
       end
     end
 
@@ -90,9 +93,19 @@ def parallelize?
       arguments.include?("--parallelize")
     end
 
+    def backup_id?
+      backup_id_index = arguments.index { |arg| arg == "--backup-id" }
+      if backup_id_index && backup_id_index + 1 < arguments.length
+        arguments[backup_id_index + 1]
+      end
+    end
+
     def additional_restore_arguments
-      (arguments.drop(1) - ["--force", "--parallelize"] +
-        [restore_confirmation_argument]).compact.join(" ")
+      # Filter out special flags that are handled separately
+      filtered_args = arguments.drop(1) - ["--force", "--parallelize"]
+      # Filter out --backup-id as it's handled separately by the Backup class
+      filtered_args = filtered_args.reject { |arg| arg.start_with?("--backup-id") }
+      (filtered_args + [restore_confirmation_argument]).compact.join(" ")
     end
 
     def restore_confirmation_argument

spec/parity/backup_spec.rb

@@ -33,6 +33,37 @@
         with(delete_local_temp_backup_command)
     end
 
+    it "restores from a specific backup ID when restoring to development" do
+      allow(IO).to receive(:read).and_return(database_fixture)
+      allow(Kernel).to receive(:system)
+      allow(Etc).to receive(:nprocessors).and_return(number_of_processes)
+
+      Parity::Backup.new(
+        from: "production",
+        to: "development",
+        backup_id: "b001",
+      ).restore
+
+      expect(Kernel).
+        to have_received(:system).
+        with(make_temp_directory_command)
+      expect(Kernel).
+        to have_received(:system).
+        with(specific_backup_id_download_command)
+      expect(Kernel).
+        to have_received(:system).
+        with(drop_development_database_drop_command)
+      expect(Kernel).
+        to have_received(:system).
+        with(create_heroku_ext_schema_command)
+      expect(Kernel).
+        to have_received(:system).
+        with(specific_backup_id_restore_from_local_temp_backup_command(cores: 1))
+      expect(Kernel).
+        to have_received(:system).
+        with(specific_backup_id_delete_local_temp_backup_command)
+    end
+
     it "restores backups to development with Rubies that do not support Etc.nprocessors" do
       allow(IO).to receive(:read).and_return(database_fixture)
       allow(Kernel).to receive(:system)
@@ -217,6 +248,20 @@
       to have_received(:system).with(additional_argument_pass_through)
   end
 
+  it "restores from a specific backup ID when provided" do
+    stub_heroku_app_name
+    allow(Kernel).to receive(:system)
+
+    Parity::Backup.new(
+      from: "production",
+      to: "staging",
+      backup_id: "b001",
+    ).restore
+
+    expect(Kernel).
+      to have_received(:system).with(specific_backup_id_restore_command)
+  end
+
   def stub_heroku_app_name
     heroku_app_name =
       instance_double(Parity::HerokuAppName, to_s: "parity-staging")
@@ -239,7 +284,7 @@ def fixture_path(filename)
   end
 
   def drop_development_database_drop_command(db_name: default_db_name)
-    "dropdb --if-exists #{db_name} && createdb #{db_name}"
+    "dropdb --if-exists #{db_name} --force && createdb #{db_name}"
   end
 
   def create_heroku_ext_schema_command(db_name: default_db_name)
@@ -258,11 +303,20 @@ def download_remote_database_command
     'curl -o tmp/latest.backup "$(heroku pg:backups:url --remote production)"'
   end
 
+  def specific_backup_id_download_command
+    'curl -o tmp/b001.backup "$(heroku pg:backups:url b001 --remote production)"'
+  end
+
   def restore_from_local_temp_backup_command(cores: number_of_processes)
     "pg_restore tmp/latest.backup --verbose --no-acl --no-owner "\
       "--dbname #{default_db_name} --jobs=#{cores} "
   end
 
+  def specific_backup_id_restore_from_local_temp_backup_command(cores: number_of_processes)
+    "pg_restore tmp/b001.backup --verbose --no-acl --no-owner "\
+      "--dbname #{default_db_name} --jobs=#{cores} "
+  end
+
   def number_of_processes
     2
   end
@@ -271,6 +325,10 @@ def delete_local_temp_backup_command
     "rm tmp/latest.backup"
   end
 
+  def specific_backup_id_delete_local_temp_backup_command
+    "rm tmp/b001.backup"
+  end
+
   def heroku_development_to_staging_passthrough(db_name: default_db_name)
     "heroku pg:push #{db_name} DATABASE_URL --remote staging "
   end
@@ -290,6 +348,11 @@ def additional_argument_pass_through
       "--confirm thisismyapp-staging"
   end
 
+  def specific_backup_id_restore_command
+    "heroku pg:backups:restore `heroku pg:backups:url "\
+      "b001 --remote production` DATABASE --remote staging "
+  end
+
   def default_db_name
     "parity_development"
   end