59 lines
2.6 KiB
Ruby
59 lines
2.6 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module Gitlab
|
|
module Database
|
|
module Migrations
|
|
module ReestablishedConnectionStack
|
|
# This is workaround for `db:migrate` that switches `ActiveRecord::Base.connection`
|
|
# depending on execution. This is subject to be removed once proper fix is implemented:
|
|
# https://gitlab.com/gitlab-org/gitlab/-/issues/362341
|
|
#
|
|
# In some cases when we run application code we need to restore application connection stack:
|
|
# - ApplicationRecord (in fact ActiveRecord::Base): points to main
|
|
# - Ci::ApplicationRecord: points to ci
|
|
#
|
|
# rubocop:disable Database/MultipleDatabases
|
|
def with_restored_connection_stack(&block)
|
|
original_handler = ActiveRecord::Base.connection_handler
|
|
|
|
original_db_config = ActiveRecord::Base.connection_db_config
|
|
if ActiveRecord::Base.configurations.primary?(original_db_config.name)
|
|
return yield(ActiveRecord::Base.connection)
|
|
end
|
|
|
|
# If the `ActiveRecord::Base` connection is different than `:main`
|
|
# re-establish and configure `SharedModel` context accordingly
|
|
# to previously established `ActiveRecord::Base` to allow the application
|
|
# code to use `ApplicationRecord` and `Ci::ApplicationRecord` usual way.
|
|
# We swap a connection handler as migration context does hold an actual
|
|
# connection which we cannot close.
|
|
base_model = Gitlab::Database.database_base_models.fetch(original_db_config.name.to_sym)
|
|
|
|
# copy connections over to new connection handler
|
|
db_configs = original_handler.connection_pool_names.map do |connection_pool_name|
|
|
[connection_pool_name.constantize, connection_pool_name.constantize.connection_db_config]
|
|
end
|
|
|
|
new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
|
|
ActiveRecord::Base.connection_handler = new_handler
|
|
|
|
db_configs.each do |klass, db_config|
|
|
new_handler.establish_connection(db_config, owner_name: klass)
|
|
end
|
|
|
|
# re-establish ActiveRecord::Base to main
|
|
ActiveRecord::Base.establish_connection :main # rubocop:disable Database/EstablishConnection
|
|
|
|
Gitlab::Database::SharedModel.using_connection(base_model.connection) do
|
|
yield(base_model.connection)
|
|
end
|
|
ensure
|
|
ActiveRecord::Base.connection_handler = original_handler
|
|
new_handler&.clear_all_connections!
|
|
end
|
|
# rubocop:enable Database/MultipleDatabases
|
|
end
|
|
end
|
|
end
|
|
end
|