88 lines
2.0 KiB
Ruby
88 lines
2.0 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module CrossDatabaseModification
|
|
extend ActiveSupport::Concern
|
|
|
|
class TransactionStackTrackRecord
|
|
def initialize(subject, gitlab_schema)
|
|
@subject = subject
|
|
@gitlab_schema = gitlab_schema
|
|
@subject.gitlab_transactions_stack.push(gitlab_schema)
|
|
end
|
|
|
|
def done!
|
|
unless @done
|
|
@done = true
|
|
@subject.gitlab_transactions_stack.pop
|
|
end
|
|
|
|
true
|
|
end
|
|
|
|
def trigger_transactional_callbacks?
|
|
false
|
|
end
|
|
|
|
def before_committed!
|
|
end
|
|
|
|
def rolledback!(force_restore_state: false, should_run_callbacks: true)
|
|
done!
|
|
end
|
|
|
|
def committed!(should_run_callbacks: true)
|
|
done!
|
|
end
|
|
end
|
|
|
|
included do
|
|
private_class_method :gitlab_schema
|
|
end
|
|
|
|
class_methods do
|
|
def gitlab_transactions_stack
|
|
Thread.current[:gitlab_transactions_stack] ||= []
|
|
end
|
|
|
|
def transaction(**options, &block)
|
|
if track_gitlab_schema_in_current_transaction?
|
|
super(**options) do
|
|
# Hook into current transaction to ensure that once
|
|
# the `COMMIT` is executed the `gitlab_transactions_stack`
|
|
# will be allowing to execute `after_commit_queue`
|
|
record = TransactionStackTrackRecord.new(self, gitlab_schema)
|
|
|
|
begin
|
|
connection.current_transaction.add_record(record)
|
|
|
|
yield
|
|
ensure
|
|
record.done!
|
|
end
|
|
end
|
|
else
|
|
super(**options, &block)
|
|
end
|
|
end
|
|
|
|
def track_gitlab_schema_in_current_transaction?
|
|
return false unless Feature::FlipperFeature.table_exists?
|
|
|
|
Feature.enabled?(:track_gitlab_schema_in_current_transaction, default_enabled: :yaml)
|
|
rescue ActiveRecord::NoDatabaseError, PG::ConnectionBad
|
|
false
|
|
end
|
|
|
|
def gitlab_schema
|
|
case self.name
|
|
when 'ActiveRecord::Base', 'ApplicationRecord'
|
|
:gitlab_main
|
|
when 'Ci::ApplicationRecord'
|
|
:gitlab_ci
|
|
else
|
|
Gitlab::Database::GitlabSchema.table_schema(table_name) if table_name
|
|
end
|
|
end
|
|
end
|
|
end
|