1 files changed, 33 insertions, 11 deletions
diff --git a/lib/tasks/gitlab/db.rake b/lib/tasks/gitlab/db.rake
index efb0e1ef1e1..6d4af9d166f 100644
--- a/lib/tasks/gitlab/db.rake
+++ b/lib/tasks/gitlab/db.rake
@@ -6,23 +6,32 @@ namespace :gitlab do
   namespace :db do
     desc 'GitLab | DB | Manually insert schema migration version'
     task :mark_migration_complete, [:version] => :environment do |_, args|
-      unless args[:version]
-        puts "Must specify a migration version as an argument".color(:red)
-        exit 1
+      mark_migration_complete(args[:version])
+    end
+
+    namespace :mark_migration_complete do
+      ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+        desc "Gitlab | DB | Manually insert schema migration version on #{name} database"
+        task name, [:version] => :environment do |_, args|
+          mark_migration_complete(args[:version], database: name)
+        end
       end
+    end
 
-      version = args[:version].to_i
-      if version == 0
-        puts "Version '#{args[:version]}' must be a non-zero integer".color(:red)
+    def mark_migration_complete(version, database: nil)
+      if version.to_i == 0
+        puts 'Must give a version argument that is a non-zero integer'.color(:red)
         exit 1
       end
 
-      sql = "INSERT INTO schema_migrations (version) VALUES (#{version})"
-      begin
-        ActiveRecord::Base.connection.execute(sql)
-        puts "Successfully marked '#{version}' as complete".color(:green)
+      Gitlab::Database.database_base_models.each do |name, model|
+        next if database && database.to_s != name
+
+        model.connection.execute("INSERT INTO schema_migrations (version) VALUES (#{model.connection.quote(version)})")
+
+        puts "Successfully marked '#{version}' as complete on database #{name}".color(:green)
       rescue ActiveRecord::RecordNotUnique
-        puts "Migration version '#{version}' is already marked complete".color(:yellow)
+        puts "Migration version '#{version}' is already marked complete on database #{name}".color(:yellow)
       end
     end
 
@@ -261,6 +270,19 @@ namespace :gitlab do
       end
     end
 
+    desc 'Run migration as gitlab non-superuser'
+    task :reset_as_non_superuser, [:username] => :environment do |_, args|
+      username = args.fetch(:username, 'gitlab')
+      puts "Migrate using username #{username}"
+      Rake::Task['db:drop'].invoke
+      Rake::Task['db:create'].invoke
+      ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env).each do |db_config|
+        ActiveRecord::Base.establish_connection(db_config.configuration_hash.merge(username: username)) # rubocop: disable Database/EstablishConnection
+        Gitlab::Database.check_for_non_superuser
+        Rake::Task['db:migrate'].invoke
+      end
+    end
+
     # Only for development environments,
     # we execute pending data migrations inline for convenience.
     Rake::Task['db:migrate'].enhance do