Use ActiveRecord::Tasks::DatabaseTasks directly in geo rake tasks

Instead of piggybacking on the existing `db:` rake tasks, redefine
them and correctly set all the variables.

Closes gitlab-org/gitlab-ee#3543

changelogs/unreleased-ee/tc-geo-use-database-tasks.yml

+---
+title: Rewrite Geo database rake tasks so they operate on the correct database
+merge_request: 3019
+author:
+type: fixed

lib/tasks/geo.rake

@@ -2,64 +2,63 @@ task spec: ['geo:db:test:prepare']
 
 namespace :geo do
   namespace :db do |ns|
-    {
-      drop: 'Drops the Geo tracking database from config/database_geo.yml for the current RAILS_ENV.',
-      create: 'Creates the Geo tracking database from config/database_geo.yml for the current RAILS_ENV.',
-      setup: 'Create the Geo tracking database, load the schema, and initialize with the seed data.',
-      migrate: 'Migrate the Geo tracking database (options: VERSION=x, VERBOSE=false, SCOPE=blog).',
-      rollback: 'Rolls the schema back to the previous version (specify steps w/ STEP=n).',
-      seed: 'Load the seed data from db/geo/seeds.rb',
-      version: 'Retrieves the current schema version number.',
-      reset: 'Drops and recreates the database from db/geo/schema.rb for the current environment and loads the seeds.'
-    }.each do |task_name, task_desc|
-      desc task_desc
-      task task_name do
-        Rake::Task["db:#{task_name}"].invoke
+    include ActiveRecord::Tasks
+
+    desc 'Drops the Geo tracking database from config/database_geo.yml for the current RAILS_ENV.'
+    task :drop do
+      with_geo_db do
+        DatabaseTasks.drop_current
       end
     end
 
-    namespace :schema do
-      {
-        load: 'Load a db/geo/schema.rb file into the database',
-        dump: 'Create a db/geo/schema.rb file that is portable against any DB supported by AR.'
-      }.each do |task_name, task_desc|
-        desc task_desc
-        task task_name do
-          Rake::Task["db:schema:#{task_name}"].invoke
-        end
+    desc 'Creates the Geo tracking database from config/database_geo.yml for the current RAILS_ENV.'
+    task :create do
+      with_geo_db do
+        DatabaseTasks.create_current
       end
     end
 
-    namespace :migrate do
-      {
-        up: 'Runs the "up" for a given migration VERSION.',
-        down: 'Runs the "down" for a given migration VERSION.',
-        redo: 'Rollbacks the database one migration and re migrate up (options: STEP=x, VERSION=x).'
-      }.each do |task_name, task_desc|
-        desc task_desc
-        task task_name do
-          Rake::Task["db:migrate:#{task_name}"].invoke
-        end
+    desc 'Create the Geo tracking database, load the schema, and initialize with the seed data.'
+    task setup: ['geo:db:schema:load', 'geo:db:seed']
+
+    desc 'Migrate the Geo tracking database (options: VERSION=x, VERBOSE=false, SCOPE=blog).'
+    task migrate: [:environment] do
+      with_geo_db do
+        DatabaseTasks.migrate
       end
+      ns['_dump'].invoke
     end
 
-    namespace :test do
-      desc 'Check for pending migrations and load the test schema'
-      task :prepare do
-        Rake::Task['db:test:prepare'].invoke
+    desc 'Rolls the schema back to the previous version (specify steps w/ STEP=n).'
+    task rollback: [:environment] do
+      step = ENV['STEP'] ? ENV['STEP'].to_i : 1
+
+      with_geo_db do
+        ActiveRecord::Migrator.rollback(ActiveRecord::Migrator.migrations_paths, step)
       end
+      ns['_dump'].invoke
     end
 
-    # append and prepend proper tasks to all the tasks defined above
-    ns.tasks.each do |task|
-      task.enhance ['geo:config:check', 'geo:config:set'] do
-        Rake::Task['geo:config:restore'].invoke
+    desc 'Retrieves the current schema version number.'
+    task version: [:environment] do
+      with_geo_db do
+        puts "Current version: #{ActiveRecord::Migrator.current_version}"
+      end
+    end
+
+    desc 'Drops and recreates the database from db/geo/schema.rb for the current environment and loads the seeds.'
+    task reset: [:environment] do
+      ns['drop'].invoke
+      ns['create'].invoke
+      ns['setup'].invoke
+    end
+
+    desc 'Load the seed data from db/seeds.rb'
+    task seed: [:environment] do
+      ns['abort_if_pending_migrations'].invoke
 
-        # Reenable the tasks, otherwise the following tasks are run only once
-        # per invocation of `rake`!
-        Rake::Task['geo:config:check'].reenable
-        Rake::Task['geo:config:set'].reenable
-        Rake::Task['geo:config:restore'].reenable
+      with_geo_db do
+        DatabaseTasks.load_seed # Without setting DatabaseTasks.seed_loader it will load from db/seeds.rb
       end
     end
 
@@ -67,34 +66,155 @@ namespace :geo do
     task show_encryption_key: :environment do
       puts Rails.application.secrets.db_key_base
     end
-  end
 
-  namespace :config do
-    task :check do
-      unless File.exist?(Rails.root.join('config/database_geo.yml'))
-        abort('You should run these tasks only when GitLab Geo is enabled.')
+    # IMPORTANT: This task won't dump the schema if ActiveRecord::Base.dump_schema_after_migration is set to false
+    task :_dump do
+      with_geo_db do # TODO should this be in `with_geo_db` ?
+        if ActiveRecord::Base.dump_schema_after_migration
+          ns["schema:dump"].invoke
+        end
       end
+      # Allow this task to be called as many times as required. An example is the
+      # migrate:redo task, which calls other two internally that depend on this one.
+      ns['_dump'].reenable
     end
 
-    task :set do
-      # save current configuration
-      @previous_config = {
-        config: Rails.application.config.dup,
-        schema: ENV['SCHEMA']
-      }
-
-      # set config variables for geo database
-      ENV['SCHEMA'] = 'db/geo/schema.rb'
-      Rails.application.config.paths['db'] = ['db/geo']
-      Rails.application.config.paths['db/migrate'] = ['db/geo/migrate']
-      Rails.application.config.paths['db/seeds.rb'] = ['db/geo/seeds.rb']
-      Rails.application.config.paths['config/database'] = ['config/database_geo.yml']
+    # desc "Raises an error if there are pending migrations"
+    task abort_if_pending_migrations: [:environment] do
+      with_geo_db do
+        pending_migrations = ActiveRecord::Migrator.open(ActiveRecord::Migrator.migrations_paths).pending_migrations
+
+        if pending_migrations.any?
+          puts "You have #{pending_migrations.size} pending #{pending_migrations.size > 1 ? 'migrations:' : 'migration:'}"
+          pending_migrations.each do |pending_migration|
+            puts '  %4d %s' % [pending_migration.version, pending_migration.name]
+          end
+          abort %{Run `rake geo:db:migrate` to update your database then try again.}
+        end
+      end
+    end
+
+    namespace :schema do
+      desc 'Load a schema.rb file into the database'
+      task load: [:environment] do
+        with_geo_db do
+          ActiveRecord::Tasks::DatabaseTasks.load_schema_current(:ruby, ENV['SCHEMA'])
+        end
+      end
+
+      desc 'Create a db/geo/schema.rb file that is portable against any DB supported by AR'
+      task dump: [:environment] do
+        require 'active_record/schema_dumper'
+
+        with_geo_db do
+          filename = ENV['SCHEMA'] || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, 'schema.rb')
+          File.open(filename, "w:utf-8") do |file|
+            ActiveRecord::SchemaDumper.dump(Geo::BaseRegistry.connection, file)
+          end
+        end
+
+        ns['schema:dump'].reenable
+      end
+    end
+
+    namespace :migrate do
+      desc 'Runs the "up" for a given migration VERSION.'
+      task up: [:environment] do
+        version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
+        raise 'VERSION is required' unless version
+
+        with_geo_db do
+          ActiveRecord::Migrator.run(:up, ActiveRecord::Migrator.migrations_paths, version)
+        end
+        ns['_dump'].invoke
+      end
+
+      desc 'Runs the "down" for a given migration VERSION.'
+      task down: [:environment] do
+        version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
+        raise 'VERSION is required - To go down one migration, run db:rollback' unless version
+
+        with_geo_db do
+          ActiveRecord::Migrator.run(:down, ActiveRecord::Migrator.migrations_paths, version)
+        end
+        ns['_dump'].invoke
+      end
+
+      desc 'Rollbacks the database one migration and re migrate up (options: STEP=x, VERSION=x).'
+      task redo: [:environment] do
+        if ENV['VERSION']
+          ns['migrate:down'].invoke
+          ns['migrate:up'].invoke
+        else
+          ns['rollback'].invoke
+          ns['migrate'].invoke
+        end
+      end
+
+      desc 'Display status of migrations' # TODO test
+      task status: [:environment] do
+        with_geo_db do
+          unless ActiveRecord::SchemaMigration.table_exists?
+            abort 'Schema migrations table does not exist yet.'
+          end
+          db_list = ActiveRecord::SchemaMigration.normalized_versions
+
+          file_list =
+            ActiveRecord::Migrator.migrations_paths.flat_map do |path|
+            # match "20091231235959_some_name.rb" and "001_some_name.rb" pattern
+            Dir.foreach(path).grep(/^(\d{3,})_(.+)\.rb$/) do
+              version = ActiveRecord::SchemaMigration.normalize_migration_number($1)
+              status = db_list.delete(version) ? 'up' : 'down'
+              [status, version, $2.humanize]
+            end
+          end
+
+          db_list.map! do |version|
+            ['up', version, '********** NO FILE **********']
+          end
+          # output
+          puts "\ndatabase: #{ActiveRecord::Base.connection_config[:database]}\n\n"
+          puts "#{'Status'.center(8)}  #{'Migration ID'.ljust(14)}  Migration Name"
+          puts "-" * 50
+          (db_list + file_list).sort_by { |_, version, _| version }.each do |status, version, name|
+            puts "#{status.center(8)}  #{version.ljust(14)}  #{name}"
+          end
+          puts
+        end
+      end
     end
 
-    task :restore do
-      # restore config variables to previous values
-      ENV['SCHEMA'] = @previous_config[:schema]
-      Rails.application.config = @previous_config[:config]
+    namespace :test do
+      desc 'Check for pending migrations and load the test schema'
+      task prepare: [:environment] do
+        with_geo_db do
+          unless ActiveRecord::Base.configurations.blank?
+            ns['test:load'].invoke
+          end
+        end
+      end
+
+      # desc "Recreate the test database from the current schema"
+      task load: [:environment, 'geo:db:test:purge'] do
+        with_geo_db do
+          begin
+            should_reconnect = ActiveRecord::Base.connection_pool.active_connection?
+            ActiveRecord::Schema.verbose = false
+            ActiveRecord::Tasks::DatabaseTasks.load_schema_for ActiveRecord::Base.configurations['test'], :ruby, ENV['SCHEMA']
+          ensure
+            if should_reconnect
+              ActiveRecord::Base.establish_connection(ActiveRecord::Base.configurations[ActiveRecord::Tasks::DatabaseTasks.env])
+            end
+          end
+        end
+      end
+
+      # desc "Empty the test database"
+      task purge: [:environment] do
+        with_geo_db do
+          ActiveRecord::Tasks::DatabaseTasks.purge ActiveRecord::Base.configurations['test']
+        end
+      end
     end
   end
 
@@ -121,4 +241,50 @@ namespace :geo do
 
     puts "Error saving GeoNode:\n#{node.errors.full_messages.join("\n")}".color(:red) unless node.persisted?
   end
+
+  GEO_DATABASE_CONFIG = 'config/database_geo.yml'.freeze
+
+  def geo_settings
+    db_dir = 'db/geo'
+    {
+      database_config: YAML.load_file(GEO_DATABASE_CONFIG),
+      db_dir: db_dir,
+      migrations_paths: [Rails.root.join(db_dir, 'migrate')]
+    }
+  end
+
+  def abort_if_no_geo_config!
+    @geo_config_exists ||= File.exist?(Rails.root.join(GEO_DATABASE_CONFIG))
+
+    unless @geo_config_exists
+      abort("Failed to open #{GEO_DATABASE_CONFIG}. Consult the documentation on how to set up GitLab Geo.")
+    end
+  end
+
+  def with_geo_db
+    abort_if_no_geo_config!
+
+    original_settings = {
+      database_config: DatabaseTasks.database_configuration&.dup || YAML.load_file('config/database.yml'),
+      db_dir: DatabaseTasks.db_dir,
+      migrations_paths: DatabaseTasks.migrations_paths
+    }
+
+    set_db_env(geo_settings)
+
+    yield
+
+    set_db_env(original_settings)
+  end
+
+  def set_db_env(settings)
+    DatabaseTasks.database_configuration = settings[:database_config]
+    DatabaseTasks.db_dir = settings[:db_dir]
+    DatabaseTasks.migrations_paths = settings[:migrations_paths]
+
+    ActiveRecord::Base.configurations       = DatabaseTasks.database_configuration || {}
+    ActiveRecord::Migrator.migrations_paths = DatabaseTasks.migrations_paths
+
+    ActiveRecord::Base.establish_connection(ActiveRecord::Base.configurations[ActiveRecord::Tasks::DatabaseTasks.env])
+  end
 end