From 3a6c3ee65d9d13da088fb3e9113bc187e171fa8e Mon Sep 17 00:00:00 2001
From: Neil Lalonde <neillalonde@gmail.com>
Date: Fri, 31 Jan 2014 14:16:15 -0500
Subject: [PATCH] Add two rake tasks: db:rebuild_indexes and
 import:remove_backup

---
 lib/import/import.rb  |  4 +++
 lib/tasks/db.rake     | 70 +++++++++++++++++++++++++++++++++++++++++++
 lib/tasks/export.rake | 12 +++++++-
 3 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/lib/import/import.rb b/lib/import/import.rb
index b9b81151067..d17c65c7a57 100644
--- a/lib/import/import.rb
+++ b/lib/import/import.rb
@@ -27,6 +27,10 @@ module Import
     $redis.del import_running_key
   end
 
+  def self.backup_tables_count
+    User.exec_sql("select count(*) as count from information_schema.tables where table_schema = '#{Jobs::Importer::BACKUP_SCHEMA}'")[0]['count'].to_i
+  end
+
 
   def self.clear_adapters
     @adapters = {}
diff --git a/lib/tasks/db.rake b/lib/tasks/db.rake
index ce531412f8f..f6b88ccc49d 100644
--- a/lib/tasks/db.rake
+++ b/lib/tasks/db.rake
@@ -6,3 +6,73 @@ end
 task 'test:prepare' => 'environment' do
   SeedFu.seed
 end
+
+desc 'Rebuild indexes'
+task 'db:rebuild_indexes' => 'environment' do
+  if Import::backup_tables_count > 0
+    raise "Backup from a previous import exists. Drop them before running this job with rake import:remove_backup, or move them to another schema."
+  end
+
+  Discourse.enable_maintenance_mode
+  backup_schema = Jobs::Importer::BACKUP_SCHEMA
+  table_names = User.exec_sql("select table_name from information_schema.tables where table_schema = 'public'").map do |row|
+    row['table_name']
+  end
+
+  begin
+    # Move all tables to the backup schema:
+    User.exec_sql("DROP SCHEMA IF EXISTS #{backup_schema} CASCADE")
+    User.exec_sql("CREATE SCHEMA #{backup_schema}")
+    table_names.each do |table_name|
+      User.exec_sql("ALTER TABLE public.#{table_name} SET SCHEMA #{backup_schema}")
+    end
+
+    # Create a new empty db
+    Rake::Task["db:migrate"].invoke
+
+    # Fetch index definitions from the new db
+    index_definitions = {}
+    table_names.each do |table_name|
+      index_definitions[table_name] = User.exec_sql("SELECT indexdef FROM pg_indexes WHERE tablename = '#{table_name}' and schemaname = 'public';").map { |x| x['indexdef'] }
+    end
+
+    # Drop the new tables
+    table_names.each do |table_name|
+      User.exec_sql("DROP TABLE public.#{table_name}")
+    end
+
+    # Move the old tables back to the public schema
+    table_names.each do |table_name|
+      User.exec_sql("ALTER TABLE #{backup_schema}.#{table_name} SET SCHEMA public")
+    end
+
+    # Drop their indexes
+    index_names = User.exec_sql("SELECT indexname FROM pg_indexes WHERE schemaname = 'public' AND tablename IN ('#{table_names.join("', '")}')").map { |x| x['indexname'] }
+    index_names.each do |index_name|
+      begin
+        puts index_name
+        User.exec_sql("DROP INDEX public.#{index_name}")
+      rescue ActiveRecord::StatementInvalid => e
+        # It's this:
+        # PG::Error: ERROR:  cannot drop index category_users_pkey because constraint category_users_pkey on table category_users requires it
+        # HINT:  You can drop constraint category_users_pkey on table category_users instead.
+      end
+    end
+
+    # Create the indexes
+    table_names.each do |table_name|
+      index_definitions[table_name].each do |index_def|
+        begin
+          User.exec_sql(index_def)
+        rescue ActiveRecord::StatementInvalid => e
+          # Trying to recreate a primary key
+        end
+      end
+    end
+  rescue
+    # Can we roll this back?
+    raise
+  ensure
+    Discourse.disable_maintenance_mode
+  end
+end
\ No newline at end of file
diff --git a/lib/tasks/export.rake b/lib/tasks/export.rake
index b19c8acf865..f0d7a079fa9 100644
--- a/lib/tasks/export.rake
+++ b/lib/tasks/export.rake
@@ -23,7 +23,7 @@ end
 
 desc 'After a successful import, restore the backup tables'
 task 'import:rollback' => :environment do |t|
-  num_backup_tables = User.exec_sql("select count(*) as count from information_schema.tables where table_schema = 'backup'")[0]['count'].to_i
+  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."
@@ -36,6 +36,16 @@ task 'import:rollback' => :environment do |t|
   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
+end
+
 desc 'Allow imports'
 task 'import:enable' => :environment do |t|
   SiteSetting.allow_import = true