Add rake task for easy migration of SQL dumps

3 files changed, 159 insertions, 27 deletions

diff --git a/lib/ci/migrate/database.rb b/lib/ci/migrate/database.rb
new file mode 100644
index 00000000000..74f592dcaea
--- /dev/null
+++ b/lib/ci/migrate/database.rb
@@ -0,0 +1,67 @@
+require 'yaml'
+
+module Ci
+  module Migrate
+    class Database
+      attr_reader :config
+
+      def initialize
+        @config = YAML.load_file(File.join(Rails.root, 'config', 'database.yml'))[Rails.env]
+      end
+
+      def restore(ci_dump)
+        puts 'Deleting all CI related data ... '
+        truncate_ci_tables
+
+        puts 'Restoring CI data ... '
+        case config["adapter"]
+          when /^mysql/ then
+            print "Restoring MySQL database #{config['database']} ... "
+            # Workaround warnings from MySQL 5.6 about passwords on cmd line
+            ENV['MYSQL_PWD'] = config["password"].to_s if config["password"]
+            system('mysql', *mysql_args, config['database'], in: ci_dump)
+          when "postgresql" then
+            puts "Restoring PostgreSQL database #{config['database']} ... "
+            pg_env
+            system('psql', config['database'], '-f', ci_dump)
+        end
+      end
+
+      protected
+
+      def truncate_ci_tables
+        c = ActiveRecord::Base.connection
+        c.tables.select { |t| t.start_with?('ci_') }.each do |table|
+          puts "Deleting data from #{table}..."
+          c.execute("DELETE FROM #{table}")
+        end
+      end
+
+      def mysql_args
+        args = {
+          'host' => '--host',
+          'port' => '--port',
+          'socket' => '--socket',
+          'username' => '--user',
+          'encoding' => '--default-character-set'
+        }
+        args.map { |opt, arg| "#{arg}=#{config[opt]}" if config[opt] }.compact
+      end
+
+      def pg_env
+        ENV['PGUSER'] = config["username"] if config["username"]
+        ENV['PGHOST'] = config["host"] if config["host"]
+        ENV['PGPORT'] = config["port"].to_s if config["port"]
+        ENV['PGPASSWORD'] = config["password"].to_s if config["password"]
+      end
+
+      def report_success(success)
+        if success
+          puts '[DONE]'.green
+        else
+          puts '[FAILED]'.red
+        end
+      end
+    end
+  end
+end
diff --git a/lib/ci/migrate/tags.rb b/lib/ci/migrate/tags.rb
new file mode 100644
index 00000000000..f4114c698d2
--- /dev/null
+++ b/lib/ci/migrate/tags.rb
@@ -0,0 +1,49 @@
+require 'yaml'
+
+module Ci
+  module Migrate
+    class Tags
+      def restore
+        puts 'Migrating tags for Runners... '
+        list_objects('Runner').each do |id|
+          putc '.'
+          runner = Ci::Runner.find_by_id(id)
+          if runner
+            tags = list_tags('Runner', id)
+            runner.update_attributes(tag_list: tags)
+          end
+        end
+        puts ''
+
+        puts 'Migrating tags for Builds... '
+        list_objects('Build').each do |id|
+          putc '.'
+          build = Ci::Build.find_by_id(id)
+          if build
+            tags = list_tags('Build', id)
+            build.update_attributes(tag_list: tags)
+          end
+        end
+        puts ''
+      end
+
+      protected
+
+      def list_objects(type)
+        ids = ActiveRecord::Base.connection.select_all(
+          "select distinct taggable_id from ci_taggings where taggable_type = #{ActiveRecord::Base::sanitize(type)}"
+        )
+        ids.map { |id| id['taggable_id'] }
+      end
+
+      def list_tags(type, id)
+        tags = ActiveRecord::Base.connection.select_all(
+          'select ci_tags.name from ci_tags ' +
+            'join ci_taggings on ci_tags.id = ci_taggings.tag_id ' +
+            "where taggable_type = #{ActiveRecord::Base::sanitize(type)} and taggable_id = #{ActiveRecord::Base::sanitize(id)} and context = \"tags\""
+        )
+        tags.map { |tag| tag['name'] }
+      end
+    end
+  end
+end
diff --git a/lib/tasks/ci/migrate.rake b/lib/tasks/ci/migrate.rake
index 7d99664dcf3..2760c503e22 100644
--- a/lib/tasks/ci/migrate.rake
+++ b/lib/tasks/ci/migrate.rake
@@ -1,38 +1,54 @@
 namespace :ci do
-  namespace :migrate do
-    def list_objects(type)
-      ids = ActiveRecord::Base.connection.select_all(
-        'select distinct taggable_id from ci_taggings where taggable_type = $1',
-        nil, [[nil, type]]
-      )
-      ids.map { |id| id['taggable_id'] }
+  desc 'GitLab | Import and migrate CI database'
+  task migrate: :environment do
+    unless ENV['force'] == 'yes'
+      puts "This will truncate all CI tables and restore it from provided backup."
+      puts "You will lose any previous CI data stored in the database."
+      ask_to_continue
+      puts ""
     end
 
-    def list_tags(type, id)
-      tags = ActiveRecord::Base.connection.select_all(
-        'select ci_tags.name from ci_tags ' +
-        'join ci_taggings on ci_tags.id = ci_taggings.tag_id ' +
-        'where taggable_type = $1 and taggable_id = $2 and context = $3',
-        nil, [[nil, type], [nil, id], [nil, 'tags']]
-      )
-      tags.map { |tag| tag['name'] }
+    Rake::Task["ci:migrate:db"].invoke
+    Rake::Task["ci:migrate:autoincrements"].invoke
+    Rake::Task["ci:migrate:tags"].invoke
+  end
+
+  namespace :migrate do
+    desc 'GitLab | Import CI database'
+    task db: :environment do
+      if ENV["CI_DUMP"].nil?
+        puts "No CI SQL dump specified:"
+        puts "rake gitlab:backup:restore CI_DUMP=ci_dump.sql"
+        exit 1
+      end
+
+      ci_dump = ENV["CI_DUMP"]
+      unless File.exists?(ci_dump)
+        puts "The specified sql dump doesn't exist!"
+        exit 1
+      end
+
+      ::Ci::Migrate::Database.new.restore(ci_dump)
     end
 
     desc 'GitLab | Migrate CI tags'
     task tags: :environment do
-      list_objects('Runner').each do |id|
-        runner = Ci::Runner.find_by_id(id)
-        if runner
-          tags = list_tags('Runner', id)
-          runner.update_attributes(tag_list: tags)
-        end
-      end
+      ::Ci::Migrate::Tags.new.restore
+    end
 
-      list_objects('Build').each do |id|
-        build = Ci::Build.find_by_id(id)
-        if build
-          tags = list_tags('Build', id)
-          build.update_attributes(tag_list: tags)
+    desc 'GitLab | Migrate CI auto-increments'
+    task autoincrements: :environment do
+      c = ActiveRecord::Base.connection
+      c.tables.select { |t| t.start_with?('ci_') }.each do |table|
+        result = c.select_one("SELECT id FROM #{table} ORDER BY id DESC LIMIT 1")
+        if result
+          ai_val = result['id'].to_i + 1
+          puts "Resetting auto increment ID for #{table} to #{ai_val}"
+          if c.adapter_name == 'PostgreSQL'
+            c.execute("ALTER SEQUENCE #{table}_id_seq RESTART WITH #{ai_val}")
+          else
+            c.execute("ALTER TABLE #{table} AUTO_INCREMENT = #{ai_val}")
+          end
         end
       end
     end