diff --git a/ee/app/models/ee/user.rb b/ee/app/models/ee/user.rb index ebe12a71d1131ea4bee3c1be326437d2462c82bc..4123afaf3cfe6853bf33901c76408b5416319509 100644 --- a/ee/app/models/ee/user.rb +++ b/ee/app/models/ee/user.rb @@ -163,6 +163,13 @@ module User scope :excluding_enterprise_users_of_group, ->(group) { left_joins(:user_detail).where('user_details.enterprise_group_id != ? OR user_details.enterprise_group_id IS NULL', group.id) } + scope :security_policy_bots_for_projects, ->(projects) do + security_policy_bot + .joins(:members) + .where(members: { source: projects }) + .allow_cross_joins_across_databases(url: "https://gitlab.com/gitlab-org/gitlab/-/issues/422405") + end + accepts_nested_attributes_for :namespace accepts_nested_attributes_for :custom_attributes diff --git a/ee/app/workers/security/orchestration_policy_rule_schedule_namespace_worker.rb b/ee/app/workers/security/orchestration_policy_rule_schedule_namespace_worker.rb index f8f0567e5ef4b086460b345d0478ff1aced4c896..9cc8cd8dfaf4b97823e5c92f2d88bfbc61d7062f 100644 --- a/ee/app/workers/security/orchestration_policy_rule_schedule_namespace_worker.rb +++ b/ee/app/workers/security/orchestration_policy_rule_schedule_namespace_worker.rb @@ -2,6 +2,7 @@ module Security class OrchestrationPolicyRuleScheduleNamespaceWorker + BATCH_SIZE = 50 include ApplicationWorker feature_category :security_policy_management @@ -20,20 +21,38 @@ def perform(rule_schedule_id) schedule.schedule_next_run! - security_orchestration_policy_configuration.namespace.all_projects.not_aimed_for_deletion.find_in_batches.each do |projects| + projects_in_batches(security_orchestration_policy_configuration).each do |projects| + bots_by_project_id = security_policy_bot_ids_by_project_ids(projects) + projects.each do |project| - user = project.security_policy_bot - next prepare_security_policy_bot_user(project) unless user + user_id = bots_by_project_id[project.id] + next prepare_security_policy_bot_user(project) unless user_id - with_context(project: project, user: user) do - Security::ScanExecutionPolicies::RuleScheduleWorker.perform_async(project.id, user.id, schedule.id) + with_context(project: project) do + Security::ScanExecutionPolicies::RuleScheduleWorker.perform_async(project.id, user_id, schedule.id) end end end end + private + def prepare_security_policy_bot_user(project) Security::OrchestrationConfigurationCreateBotWorker.perform_async(project.id, nil) end + + def security_policy_bot_ids_by_project_ids(projects) + User.security_policy_bots_for_projects(projects).select(:id, :source_id).to_h do |bot| + [bot.source_id, bot.id] + end + end + + def projects_in_batches(configuration) + configuration + .namespace + .all_projects.not_aimed_for_deletion + .select(:id) + .find_in_batches(batch_size: BATCH_SIZE) # rubocop: disable CodeReuse/ActiveRecord -- A custom batch size is needed because querying too many bot users at once is too expensive + end end end diff --git a/ee/spec/models/ee/user_spec.rb b/ee/spec/models/ee/user_spec.rb index 82476b2746505008140409df637a427e941c48db..60d1a9c2896d089f4f6a6529c575e6faa6659402 100644 --- a/ee/spec/models/ee/user_spec.rb +++ b/ee/spec/models/ee/user_spec.rb @@ -348,6 +348,24 @@ end end end + + describe '.security_policy_bots_for_projects' do + let_it_be(:project_1) { create(:project) } + let_it_be(:project_2) { create(:project) } + let_it_be(:security_policy_bot_1) { create(:user, :security_policy_bot) } + let_it_be(:security_policy_bot_2) { create(:user, :security_policy_bot) } + + let(:projects) { [project_1, project_2] } + + before_all do + project_1.add_guest(security_policy_bot_1) + project_2.add_guest(security_policy_bot_2) + end + + subject { described_class.security_policy_bots_for_projects(projects) } + + it { is_expected.to contain_exactly(security_policy_bot_1, security_policy_bot_2) } + end end describe 'after_create' do diff --git a/ee/spec/workers/security/orchestration_policy_rule_schedule_namespace_worker_spec.rb b/ee/spec/workers/security/orchestration_policy_rule_schedule_namespace_worker_spec.rb index c761d610fc395f874c138e8f46f4c9b32fa33b9f..cf465aa2260f4170d1d926e9150bf8320a6b11dc 100644 --- a/ee/spec/workers/security/orchestration_policy_rule_schedule_namespace_worker_spec.rb +++ b/ee/spec/workers/security/orchestration_policy_rule_schedule_namespace_worker_spec.rb @@ -54,6 +54,16 @@ expect(schedule.reload.next_run_at).to be_future end + it 'does not trigger N+1 queries', :use_sql_query_cache do + control = ActiveRecord::QueryRecorder.new { worker.perform(schedule_id) } + + create(:project, namespace: namespace) + create(:project, namespace: namespace) + schedule.update_column(:next_run_at, 1.minute.ago) + + expect { worker.perform(schedule_id) }.to issue_same_number_of_queries_as(control) + end + context 'when there is a security_policy_bot in the project' do let_it_be(:security_policy_bot) { create(:user, :security_policy_bot) }