BUGFIX: backup/restore rake tasks weren't working

lib/backup_restore.rb

@@ -22,7 +22,7 @@ module BackupRestore
   def self.rollback!
     raise BackupRestore::OperationRunningError if BackupRestore.is_operation_running?
     if can_rollback?
-      rename_schema("backup", "public")
+      move_tables_between_schemas("backup", "public")
       after_fork
     end
   end
@@ -75,15 +75,43 @@ module BackupRestore
     User.exec_sql("SELECT 1 FROM pg_namespace WHERE nspname = 'backup'").count > 0
   end
 
-  def self.rename_schema(old_name, new_name)
-    sql = <<-SQL
-      BEGIN;
-        DROP SCHEMA IF EXISTS #{new_name} CASCADE;
-        ALTER SCHEMA #{old_name} RENAME TO #{new_name};
-      COMMIT;
-    SQL
+  def self.move_tables_between_schemas(source, destination)
+    User.exec_sql(move_tables_between_schemas_sql(source, destination))
+  end
 
-    User.exec_sql(sql)
+  def self.move_tables_between_schemas_sql(source, destination)
+    # TODO: Postgres 9.3 has "CREATE SCHEMA schema IF NOT EXISTS;"
+    <<-SQL
+      DO $$DECLARE row record;
+      BEGIN
+        -- create "destination" schema if it does not exists already
+        -- NOTE: DROP & CREATE SCHEMA is easier, but we don't wont to drop the public schema
+        -- ortherwise extensions (like hstore & pg_trgm) won't work anymore
+        IF NOT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = '#{destination}')
+        THEN
+          CREATE SCHEMA #{destination};
+        END IF;
+        -- move all "source" tables to "destination" schema
+        FOR row IN SELECT tablename FROM pg_tables WHERE schemaname = '#{source}'
+        LOOP
+          EXECUTE 'ALTER TABLE #{source}.' || quote_ident(row.tablename) || ' SET SCHEMA #{destination};';
+        END LOOP;
+      END$$;
+    SQL
+  end
+
+  DatabaseConfiguration = Struct.new(:host, :username, :password, :database)
+
+  def self.database_configuration
+    if Rails.env.production?
+      conn = RailsMultisite::ConnectionManagement
+      db_conf = DatabaseConfiguration.new(conn.current_host, conn.current_username, conn.current_password, conn.current_db)
+    else
+      db = Rails.configuration.database_configuration[Rails.env]
+      db_conf = DatabaseConfiguration.new(db["host"], db["username"], db["password"], db["database"])
+    end
+    db_conf.username ||= ENV["USER"] || "postgres"
+    db_conf
   end
 
   private

lib/export/exporter.rb

@@ -152,14 +152,10 @@ module Export
     end
 
     def build_pg_dump_command
-      db_conf = Rails.configuration.database_configuration[Rails.env]
-      host = db_conf["host"]
-      password = db_conf["password"]
-      username = db_conf["username"] || ENV["USER"] || "postgres"
-      database = db_conf["database"]
+      db_conf = BackupRestore.database_configuration
 
-      password_argument = "PGPASSWORD=#{password}" if password.present?
-      host_argument     = "--host=#{host}"         if host.present?
+      password_argument = "PGPASSWORD=#{password}" if db_conf.password.present?
+      host_argument     = "--host=#{host}"         if db_conf.host.present?
 
       [ password_argument,                  # pass the password to pg_dump
         "pg_dump",                          # the pg_dump command
@@ -169,8 +165,8 @@ module Export
         "--no-privileges",                  # prevent dumping of access privileges
         "--verbose",                        # specifies verbose mode
         host_argument,                      # the hostname to connect to
-        "--username=#{username}",           # the username to connect as
-        database                            # the name of the database to dump
+        "--username=#{db_conf.username}",   # the username to connect as
+        db_conf.database                    # the name of the database to dump
       ].join(" ")
     end
 

lib/import/importer.rb

@@ -200,35 +200,30 @@ module Import
     end
 
     def build_psql_command
-      db_conf = Rails.configuration.database_configuration[Rails.env]
-      host = db_conf["host"]
-      password = db_conf["password"]
-      username = db_conf["username"] || ENV["USER"] || "postgres"
-      database = db_conf["database"]
+      db_conf = BackupRestore.database_configuration
 
-      password_argument = "PGPASSWORD=#{password}" if password.present?
-      host_argument     = "--host=#{host}"         if host.present?
+      password_argument = "PGPASSWORD=#{password}" if db_conf.password.present?
+      host_argument     = "--host=#{host}"         if db_conf.host.present?
 
       [ password_argument,                # pass the password to psql
         "psql",                           # the psql command
-        "--dbname='#{database}'",     # connect to database *dbname*
+        "--dbname='#{db_conf.database}'", # connect to database *dbname*
         "--file='#{@dump_filename}'",     # read the dump
         "--single-transaction",           # all or nothing (also runs COPY commands faster)
         host_argument,                    # the hostname to connect to
-        "--username=#{username}"      # the username to connect as
+        "--username=#{db_conf.username}"  # the username to connect as
       ].join(" ")
     end
 
     def switch_schema!
       log "Switching schemas..."
 
-      sql = <<-SQL
-        BEGIN;
-          DROP SCHEMA IF EXISTS backup CASCADE;
-          ALTER SCHEMA public RENAME TO backup;
-          ALTER SCHEMA restore RENAME TO public;
-        COMMIT;
-      SQL
+      sql = [
+        "BEGIN;",
+        BackupRestore.move_tables_between_schemas_sql("public", "backup"),
+        BackupRestore.move_tables_between_schemas_sql("restore", "public"),
+        "COMMIT;"
+      ].join("\n")
 
       User.exec_sql(sql)
     end
@@ -268,7 +263,7 @@ module Import
       log "Trying to rollback..."
       if BackupRestore.can_rollback?
         log "Rolling back..."
-        BackupRestore.rename_schema("backup", "public")
+        BackupRestore.move_tables_between_schemas("backup", "public")
       else
         log "There was no need to rollback"
       end

lib/tasks/export.rake

@@ -1,16 +1,27 @@
 desc 'export the database'
 task 'export', [:output_filename] => :environment do |t, args|
-  puts 'Starting export...'
-  output_filename = Jobs::Exporter.new.execute( format: :json, filename: args.output_filename )
-  puts 'Export done.'
-  puts "Output file is in: #{output_filename}", ''
+  require "backup_restore"
+  require "export/exporter"
+
+  puts "Starting export..."
+  backup = Export::Exporter.new(Discourse.system_user.id).run
+  if args.output_filename.present?
+    puts "Moving '#{backup}' to '#{filename}'"
+    FileUtils.mv(backup, args.output_filename)
+    backup = args.output_filename
+  end
+  puts "Export done."
+  puts "Output file is in: #{backup}", ""
 end
 
 desc 'import from an export file and replace the contents of the current database'
 task 'import', [:input_filename] => :environment do |t, args|
-  puts 'Starting import...'
+  require "backup_restore"
+  require "import/importer"
+
   begin
-    Jobs::Importer.new.execute( format: :json, filename: args.input_filename )
+    puts 'Starting import...'
+    Import::Importer.new(Discourse.system_user.id, args.input_filename).run
     puts 'Import done.'
   rescue Import::FilenameMissingError
     puts '', 'The filename argument was missing.', '', 'Usage:', ''
@@ -23,27 +34,10 @@ end
 
 desc 'After a successful import, restore the backup tables'
 task 'import:rollback' => :environment do |t|
-  num_backup_tables = Import::backup_tables_count
-
-  if User.exec_sql("select count(*) as count from information_schema.schemata where schema_name = 'backup'")[0]['count'].to_i <= 0
-    puts "Backup tables don't exist!  An import was never performed or the backup tables were dropped.", "Rollback cancelled."
-  elsif num_backup_tables != Export.models_included_in_export.size
-    puts "Expected #{Export.models_included_in_export.size} backup tables, but there are #{num_backup_tables}!", "Rollback cancelled."
-  else
-    puts 'Starting rollback..'
-    Jobs::Importer.new.rollback
-    puts 'Rollback done.'
-  end
-end
-
-desc 'After a successful import, drop the backup tables'
-task 'import:remove_backup' => :environment do |t|
-  if Import::backup_tables_count > 0
-    User.exec_sql("DROP SCHEMA IF EXISTS #{Jobs::Importer::BACKUP_SCHEMA} CASCADE")
-    puts "Backup tables dropped successfully."
-  else
-    puts "No backup found. Nothing was done."
-  end
+  puts 'Rolling back if needed..'
+  require "backup_restore"
+  BackupRestore.rollback!
+  puts 'Done.'
 end
 
 desc 'Allow imports'

script/discourse

@@ -39,16 +39,20 @@ WHERE table_schema='public' and (data_type like 'char%' or data_type like 'text%
   end
 
   desc "backup", "Backup a discourse forum"
-  def backup
+  def backup(filename = nil)
     load_rails
-
     require "backup_restore"
     require "export/exporter"
 
-    puts 'Starting export...'
-    output_filename = Export::Exporter.new(Discourse.system_user).run
-    puts 'Export done.'
-    puts "Output file is in: #{output_filename}", ''
+    puts "Starting export..."
+    backup = Export::Exporter.new(Discourse.system_user.id).run
+    if filename.present?
+      puts "Moving '#{backup}' to '#{filename}'"
+      FileUtils.mv(backup, filename)
+      backup = filename
+    end
+    puts "Export done."
+    puts "Output file is in: #{backup}", ""
   end
 
   desc "export", "Backup a Discourse forum"
@@ -59,20 +63,19 @@ WHERE table_schema='public' and (data_type like 'char%' or data_type like 'text%
   desc "restore", "Restore a Discourse backup"
   def restore(filename)
     load_rails
-
     require "backup_restore"
     require "import/importer"
 
     begin
-      puts "Starting import: #{filename}"
-      Import::Importer.new(Discourse.system_user, filename).run
-      puts 'Import done.'
+      puts "Starting restore: #{filename}"
+      Import::Importer.new(Discourse.system_user.id, filename).run
+      puts 'Restore done.'
     rescue Import::FilenameMissingError
       puts '', 'The filename argument was missing.', ''
       usage
     rescue Import::ImportDisabledError
-      puts '', 'Imports are not allowed.', 'An admin needs to set allow_restore to true in the site settings before imports can be run.', ''
-      puts 'Import cancelled.', ''
+      puts '', 'Restore are not allowed.', 'An admin needs to set allow_restore to true in the site settings before restores can be run.', ''
+      puts 'Restore cancelled.', ''
     end
   end
 

vendor/gems/rails_multisite/lib/rails_multisite/connection_management.rb

@@ -80,9 +80,20 @@ module RailsMultisite
 
     def self.current_hostname
       config = ActiveRecord::Base.connection_pool.spec.config
-      config[:host_names].nil? ? config[:host] : config[:host_names].first
+      config[:host_names].nil? ? current_host : config[:host_names].first
     end
 
+    def self.current_host
+      ActiveRecord::Base.connection_pool.spec.config[:host]
+    end
+
+    def self.current_username
+      ActiveRecord::Base.connection_pool.spec.config[:username]
+    end
+
+    def self.current_password
+      ActiveRecord::Base.connection_pool.spec.config[:password]
+    end
 
     def self.clear_settings!
       @@db_spec_cache = nil