1 files changed, 44 insertions, 0 deletions
diff --git a/lib/gitlab/patch/legacy_database_config.rb b/lib/gitlab/patch/legacy_database_config.rb
index a7d4fdf7490..6040f737c75 100644
--- a/lib/gitlab/patch/legacy_database_config.rb
+++ b/lib/gitlab/patch/legacy_database_config.rb
@@ -35,6 +35,40 @@ module Gitlab
         attr_reader :uses_legacy_database_config
       end
 
+      def load_database_yaml
+        return super unless Gitlab.ee?
+
+        super.deep_merge(load_geo_database_yaml)
+      end
+
+      # This method is taken from Rails to load a database YAML file without
+      # evaluating ERB. This allows us to create the rake tasks for the Geo
+      # tracking database without filling in the configuration values or
+      # loading the environment. To be removed when we start configure Geo
+      # tracking database in database.yml instead of custom database_geo.yml
+      #
+      # https://github.com/rails/rails/blob/v6.1.4/railties/lib/rails/application/configuration.rb#L255
+      def load_geo_database_yaml
+        path = Rails.root.join("config/database_geo.yml")
+        return {} unless File.exist?(path)
+
+        require "rails/application/dummy_erb_compiler"
+
+        yaml = DummyERB.new(Pathname.new(path).read).result
+        config = YAML.load(yaml) || {} # rubocop:disable Security/YAMLLoad
+
+        config.to_h do |env, configs|
+          # This check is taken from Rails where the transformation
+          # of a flat database.yml is done into `primary:`
+          # https://github.com/rails/rails/blob/v6.1.4/activerecord/lib/active_record/database_configurations.rb#L169
+          if configs.is_a?(Hash) && !configs.all? { |_, v| v.is_a?(Hash) }
+            configs = { "geo" => configs }
+          end
+
+          [env, configs]
+        end
+      end
+
       def database_configuration
         @uses_legacy_database_config = false # rubocop:disable Gitlab/ModuleWithInstanceVariables
 
@@ -48,6 +82,16 @@ module Gitlab
             @uses_legacy_database_config = true # rubocop:disable Gitlab/ModuleWithInstanceVariables
           end
 
+          if Gitlab.ee? && File.exist?(Rails.root.join("config/database_geo.yml"))
+            migrations_paths = ["ee/db/geo/migrate"]
+            migrations_paths << "ee/db/geo/post_migrate" unless ENV['SKIP_POST_DEPLOYMENT_MIGRATIONS']
+
+            configs["geo"] =
+              Rails.application.config_for(:database_geo)
+                .merge(migrations_paths: migrations_paths, schema_migrations_path: "ee/db/geo/schema_migrations")
+                .stringify_keys
+          end
+
           [env, configs]
         end
       end