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.

Using a ruby module with methods used by the Rakefile.

10.0 backport of gitlab-org/gitlab-ee!3019.

Closes #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: 3052
+author:
+type: fixed

lib/gitlab/geo/database_tasks.rb

+module Gitlab
+  module Geo
+    module DatabaseTasks
+      extend self
+
+      DATABASE_CONFIG = 'config/database.yml'.freeze
+      GEO_DATABASE_CONFIG = 'config/database_geo.yml'.freeze
+      GEO_DB_DIR = 'db/geo'.freeze
+
+      def method_missing(method_name, *args, &block)
+        with_geo_db do
+          ActiveRecord::Tasks::DatabaseTasks.public_send(method_name, *args, &block) # rubocop:disable GitlabSecurity/PublicSend
+        end
+      end
+
+      def respond_to_missing?(method_name, include_private = false)
+        ActiveRecord::Tasks::DatabaseTasks.respond_to?(method_name) || super
+      end
+
+      def rollback
+        step = ENV['STEP'] ? ENV['STEP'].to_i : 1
+
+        with_geo_db do
+          ActiveRecord::Migrator.rollback(ActiveRecord::Migrator.migrations_paths, step)
+        end
+      end
+
+      def version
+        with_geo_db do
+          ActiveRecord::Migrator.current_version
+        end
+      end
+
+      def dump_schema_after_migration?
+        with_geo_db do
+          !!ActiveRecord::Base.dump_schema_after_migration
+        end
+      end
+
+      def pending_migrations
+        with_geo_db do
+          ActiveRecord::Migrator.open(ActiveRecord::Migrator.migrations_paths).pending_migrations
+        end
+      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
+
+      module Schema
+        extend self
+
+        def dump
+          require 'active_record/schema_dumper'
+
+          Gitlab::Geo::DatabaseTasks.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(ActiveRecord::Base.connection, file)
+            end
+          end
+        end
+      end
+
+      module Migrate
+        extend self
+
+        def up
+          version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
+          raise 'VERSION is required' unless version
+
+          Gitlab::Geo::DatabaseTasks.with_geo_db do
+            ActiveRecord::Migrator.run(:up, ActiveRecord::Migrator.migrations_paths, version)
+          end
+        end
+
+        def down
+          version = ENV['VERSION'] ? ENV['VERSION'].to_i : nil
+          raise 'VERSION is required - To go down one migration, run db:rollback' unless version
+
+          Gitlab::Geo::DatabaseTasks.with_geo_db do
+            ActiveRecord::Migrator.run(:down, ActiveRecord::Migrator.migrations_paths, version)
+          end
+        end
+
+        # rubocop: disable Rails/Output
+        def status
+          Gitlab::Geo::DatabaseTasks.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
+        # rubocop: enable Rails/Output
+      end
+
+      module Test
+        extend self
+
+        def load
+          Gitlab::Geo::DatabaseTasks.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
+
+        def purge
+          Gitlab::Geo::DatabaseTasks.with_geo_db do
+            ActiveRecord::Tasks::DatabaseTasks.purge ActiveRecord::Base.configurations['test']
+          end
+        end
+      end
+
+      def geo_settings
+        {
+          database_config: YAML.load_file(GEO_DATABASE_CONFIG),
+          db_dir: GEO_DB_DIR,
+          migrations_paths: [Rails.root.join(GEO_DB_DIR, 'migrate')],
+          seed_loader: SeedLoader.new
+        }
+      end
+
+      def with_geo_db
+        abort_if_no_geo_config!
+
+        original_settings = {
+          database_config: ActiveRecord::Tasks::DatabaseTasks.database_configuration&.dup || YAML.load_file(DATABASE_CONFIG),
+          db_dir: ActiveRecord::Tasks::DatabaseTasks.db_dir,
+          migrations_paths: ActiveRecord::Tasks::DatabaseTasks.migrations_paths,
+          seed_loader: ActiveRecord::Tasks::DatabaseTasks.seed_loader
+        }
+
+        set_db_env(geo_settings)
+
+        yield
+      ensure
+        set_db_env(original_settings)
+      end
+
+      def set_db_env(settings)
+        ActiveRecord::Tasks::DatabaseTasks.database_configuration = settings[:database_config]
+        ActiveRecord::Tasks::DatabaseTasks.db_dir = settings[:db_dir]
+        ActiveRecord::Tasks::DatabaseTasks.migrations_paths = settings[:migrations_paths]
+        ActiveRecord::Tasks::DatabaseTasks.seed_loader = settings[:seed_loader]
+
+        ActiveRecord::Base.configurations       = ActiveRecord::Tasks::DatabaseTasks.database_configuration || {}
+        ActiveRecord::Migrator.migrations_paths = ActiveRecord::Tasks::DatabaseTasks.migrations_paths
+
+        ActiveRecord::Base.establish_connection(ActiveRecord::Base.configurations[ActiveRecord::Tasks::DatabaseTasks.env])
+      end
+
+      class SeedLoader
+        def load_seed
+          load('db/geo/seeds.rb')
+        end
+      end
+    end
+  end
+end

lib/tasks/geo.rake

+require 'gitlab/geo'
+require 'gitlab/geo/database_tasks'
+
 task spec: ['geo:db:test:prepare']
 namespace :geo do
   namespace :db do |ns|
-    %i(drop create setup migrate rollback seed version reset).each do |task_name|
-      task task_name do
-        Rake::Task["db:#{task_name}"].invoke
-      end
+    desc 'Drops the Geo tracking database from config/database_geo.yml for the current RAILS_ENV.'
+    task :drop do
+      Gitlab::Geo::DatabaseTasks.drop_current
     end
-    namespace :schema do
-      %i(load dump).each do |task_name|
-        task task_name do
-          Rake::Task["db:schema:#{task_name}"].invoke
-        end
-      end
+    desc 'Creates the Geo tracking database from config/database_geo.yml for the current RAILS_ENV.'
+    task :create do
+      Gitlab::Geo::DatabaseTasks.create_current
     end
-    namespace :migrate do
-      %i(up down redo).each do |task_name|
-        task task_name do
-          Rake::Task["db:migrate:#{task_name}"].invoke
-        end
-      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
+      Gitlab::Geo::DatabaseTasks.migrate
+
+      ns['_dump'].invoke
     end
-    namespace :test do
-      task :prepare do
-        Rake::Task['db:test:prepare'].invoke
-      end
+    desc 'Rolls the schema back to the previous version (specify steps w/ STEP=n).'
+    task rollback: [:environment] do
+      Gitlab::Geo::DatabaseTasks.rollback
+
+      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
-        # 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
-      end
+    desc 'Retrieves the current schema version number.'
+    task version: [:environment] do
+      puts "Current version: #{Gitlab::Geo::DatabaseTasks.version}"
+    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/geo/seeds.rb'
+    task seed: [:environment] do
+      ns['abort_if_pending_migrations'].invoke
+
+      Gitlab::Geo::DatabaseTasks.load_seed
     end
     desc 'Display database encryption key'
     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
+      if Gitlab::Geo::DatabaseTasks.dump_schema_after_migration?
+        ns["schema:dump"].invoke
+      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
+
+    # desc "Raises an error if there are pending migrations"
+    task abort_if_pending_migrations: [:environment] do
+      pending_migrations = Gitlab::Geo::DatabaseTasks.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
+
+    namespace :schema do
+      desc 'Load a schema.rb file into the database'
+      task load: [:environment] do
+        Gitlab::Geo::DatabaseTasks.load_schema_current(:ruby, ENV['SCHEMA'])
+      end
+
+      desc 'Create a db/geo/schema.rb file that is portable against any DB supported by AR'
+      task dump: [:environment] do
+        Gitlab::Geo::DatabaseTasks::Schema.dump
+
+        ns['schema:dump'].reenable
       end
     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']
+    namespace :migrate do
+      desc 'Runs the "up" for a given migration VERSION.'
+      task up: [:environment] do
+        Gitlab::Geo::DatabaseTasks::Migrate.up
+
+        ns['_dump'].invoke
+      end
+
+      desc 'Runs the "down" for a given migration VERSION.'
+      task down: [:environment] do
+        Gitlab::Geo::DatabaseTasks::Migrate.down
+
+        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'
+      task status: [:environment] do
+        Gitlab::Geo::DatabaseTasks::Migrate.status
+      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
+        ns['test:load'].invoke
+      end
+
+      # desc "Recreate the test database from the current schema"
+      task load: [:environment, 'geo:db:test:purge'] do
+        Gitlab::Geo::DatabaseTasks::Test.load
+      end
+
+      # desc "Empty the test database"
+      task purge: [:environment] do
+        Gitlab::Geo::DatabaseTasks::Test.purge
+      end
     end
   end

spec/lib/gitlab/geo/database_tasks_spec.rb

+require 'spec_helper'
+
+describe Gitlab::Geo::DatabaseTasks do
+  let(:schema_file) { Rails.root.join('tmp', 'tests', 'geo_schema.rb').to_s }
+  subject { described_class }
+
+  before do
+    stub_env('SCHEMA', schema_file) # schema will be dumped to this file
+  end
+
+  after do
+    FileUtils.rm_rf(schema_file)
+  end
+
+  [:drop_current, :create_current, :migrate, :load_seed, :load_schema_current].each do |method_name|
+    it "defines the missing method #{method_name}" do
+      is_expected.to respond_to(method_name)
+    end
+
+    it "forwards method #{method_name} to ActiveRecord::Tasks::DatabaseTasks" do
+      expect(ActiveRecord::Tasks::DatabaseTasks).to receive(method_name)
+
+      subject.public_send(method_name)
+    end
+  end
+
+  describe '.rollback' do
+    context 'ENV["STEP"] not set' do
+      it 'calls ActiveRecord::Migrator.rollback with step 1' do
+        expect(ActiveRecord::Migrator).to receive(:rollback).with(anything, 1)
+
+        subject.rollback
+      end
+    end
+  end
+
+  describe '.version' do
+    it 'returns a Number' do
+      expect(subject.version).to be_an(Integer)
+    end
+  end
+
+  describe '.dump_schema_after_migration?' do
+    it 'returns a true value' do
+      expect(subject.dump_schema_after_migration?).to be_truthy
+    end
+  end
+
+  describe '.pending_migrations' do
+    it 'returns an array' do
+      expect(subject.pending_migrations).to be_an(Array)
+    end
+  end
+
+  describe described_class::Schema do
+    describe '.dump' do
+      it 'calls ActiveRecord::SchemaDumper.dump' do
+        expect(ActiveRecord::SchemaDumper).to receive(:dump)
+
+        subject.dump
+      end
+    end
+  end
+
+  describe described_class::Migrate do
+    describe '.up' do
+      it 'requires ENV["VERSION"] to be set' do
+        expect { subject.up }.to raise_error(String)
+      end
+
+      it 'calls ActiveRecord::Migrator.run' do
+        stub_env('VERSION', '19700101120000')
+        expect(ActiveRecord::Migrator).to receive(:run).with(:up, any_args)
+
+        subject.up
+      end
+    end
+
+    describe '.down' do
+      it 'requires ENV["VERSION"] to be set' do
+        expect { subject.down }.to raise_error(String)
+      end
+
+      it 'calls ActiveRecord::Migrator.run' do
+        stub_env('VERSION', '19700101120000')
+        expect(ActiveRecord::Migrator).to receive(:run).with(:down, any_args)
+
+        subject.down
+      end
+    end
+
+    describe '.status' do
+      it 'outputs "database: gitlabhq_geo_test"' do
+        expect(ActiveRecord::SchemaMigration).to receive(:normalized_versions).and_return([])
+
+        expect { subject.status }.to output(/database: gitlabhq_geo_test/).to_stdout
+      end
+    end
+  end
+
+  describe described_class::Test do
+    describe '.load' do
+      it 'calls ActiveRecord::Tasks::DatabaseTasks.load_schema_for' do
+        expect(ActiveRecord::Tasks::DatabaseTasks).to receive(:load_schema_for)
+
+        subject.load
+      end
+    end
+
+    describe '.purge' do
+      it 'calls ActiveRecord::Tasks::DatabaseTasks.load_schema_for' do
+        expect(ActiveRecord::Tasks::DatabaseTasks).to receive(:purge)
+
+        subject.purge
+      end
+    end
+  end
+end