Use connects_to to connect to the Geo tracking DB

Use connects_to to make Geo models use
the tracking database instead of connection
override in Geo::TrackingBase

Changelog: changed
EE: true

ee/app/models/geo/tracking_base.rb

@@ -10,13 +10,7 @@ module Geo
     SecondaryNotConfigured = Class.new(StandardError)
 
     if ::Gitlab::Geo.geo_database_configured?
-      # Mark current model as a `connection_class`
-      # This is the behavior when executing `connects_to`
-      # which indicates that given class is holding a connection object
-      # and should be treated accordingly.
-      self.connection_class = true
-
-      establish_connection Rails.configuration.geo_database
+      connects_to database: { writing: :geo, reading: :geo }
     end
 
     def self.connected?
@@ -37,5 +31,13 @@ module Geo
     rescue ActiveRecord::NoDatabaseError
       raise SecondaryNotConfigured, NOT_CONFIGURED_MSG
     end
+
+    class SchemaMigration < TrackingBase
+      class << self
+        def all_versions
+          order(:version).pluck(:version)
+        end
+      end
+    end
   end
 end

ee/lib/gitlab/geo.rb

@@ -75,7 +75,7 @@ module Gitlab
     end
 
     def self.geo_database_configured?
-      Rails.configuration.respond_to?(:geo_database)
+      ::Gitlab::Database.has_config?(:geo)
     end
 
     def self.primary_node_configured?

ee/lib/system_check/geo/geo_database_configured_check.rb

@@ -61,7 +61,27 @@ module SystemCheck
       end
 
       def tables_present?
-        Gitlab::Geo::DatabaseTasks.with_geo_db { !ActiveRecord::Base.connection.migration_context.needs_migration? }
+        !needs_migration?
+      end
+
+      def needs_migration?
+        (migrations.collect(&:version) - get_all_versions).size > 0 # rubocop:disable Style/ZeroLengthPredicate
+      end
+
+      def get_all_versions
+        if schema_migration.table_exists?
+          schema_migration.all_versions.map(&:to_i)
+        else
+          []
+        end
+      end
+
+      def migrations
+        ::Geo::TrackingBase.connection.migration_context.migrations
+      end
+
+      def schema_migration
+        ::Geo::TrackingBase::SchemaMigration
       end
 
       def geo_health_check

ee/spec/lib/system_check/geo/geo_database_configured_check_spec.rb

@@ -7,15 +7,14 @@ RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
 
   describe '#multi_check', :reestablished_active_record_base do
     it "checks database configuration" do
-      stub_configuration_check(false)
+      stub_database_state(subject, configured: false)
 
       expect(subject).to receive(:try_fixing_it).with(described_class::WRONG_CONFIGURATION_MESSAGE)
       expect(subject.multi_check).to be_falsey
     end
 
     it "checks database configuration" do
-      stub_configuration_check(true)
-      stub_connection_state(false)
+      stub_database_state(subject, active: false)
 
       expect(subject).to receive(:try_fixing_it).with(described_class::UNHEALTHY_CONNECTION_MESSAGE)
 
@@ -23,9 +22,7 @@ RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
     end
 
     it "checks table existence" do
-      stub_configuration_check(true)
-      stub_connection_state(true)
-      stub_tables_existence(false)
+      stub_database_state(subject, tables: false)
 
       expect(subject).to receive(:try_fixing_it).with(described_class::NO_TABLES_MESSAGE)
 
@@ -33,10 +30,7 @@ RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
     end
 
     it "checks if existing database is being reused" do
-      stub_configuration_check(true)
-      stub_connection_state(true)
-      stub_tables_existence(true)
-      stub_fresh_database(false)
+      stub_database_state(subject, fresh: false)
 
       expect(subject).to receive(:try_fixing_it).with(described_class::REUSING_EXISTING_DATABASE_MESSAGE)
 
@@ -44,10 +38,7 @@ RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
     end
 
     it "returns true when all checks passed" do
-      stub_configuration_check(true)
-      stub_connection_state(true)
-      stub_tables_existence(true)
-      stub_fresh_database(true)
+      stub_database_state(subject)
 
       expect(subject).not_to receive(:try_fixing_it)
 
@@ -55,23 +46,14 @@ RSpec.describe SystemCheck::Geo::GeoDatabaseConfiguredCheck, :silence_stdout do
     end
   end
 
-  def stub_configuration_check(state)
-    expect(Gitlab::Geo).to receive(:geo_database_configured?).and_return(state)
-  end
-
-  def stub_connection_state(state)
-    connection = double
-    expect(connection).to receive(:active?).and_return(state)
-    expect(::Geo::TrackingBase).to receive(:connection).and_return(connection)
-  end
+  def stub_database_state(subject, configured: true, active: true, tables: true, fresh: true)
+    allow(subject).to receive(:needs_migration?).and_return(!tables)
 
-  def stub_tables_existence(state)
-    expect_any_instance_of(ActiveRecord::MigrationContext).to receive(:needs_migration?).and_return(!state)
-  end
+    allow(::Gitlab::Geo).to receive(:geo_database_configured?).and_return(configured)
+    allow(::Geo::TrackingBase).to receive(:connection).and_return(double(active?: active))
 
-  def stub_fresh_database(state)
-    expect_next_instance_of(Gitlab::Geo::HealthCheck) do |geo_health_check|
-      expect(geo_health_check).to receive(:reusing_existing_tracking_database?).and_return(!state)
+    allow_next_instance_of(::Gitlab::Geo::HealthCheck) do |health_check|
+      allow(health_check).to receive(:reusing_existing_tracking_database?).and_return(!fresh)
     end
   end
 end

ee/spec/support/helpers/ee/migrations_helpers.rb

@@ -45,7 +45,7 @@ module EE
 
     def with_db_config(&block)
       if geo_migration?
-        ::Gitlab::Geo::DatabaseTasks.with_geo_db { yield }
+        ::Geo::TrackingBase.connected_to(database: :geo) { yield }
       else
         yield
       end