From 1306a2a61e309040fa46d3b09af362a51491a96d Mon Sep 17 00:00:00 2001
From: hustewart <hustewart@gitlab.com>
Date: Mon, 21 Apr 2025 15:51:11 -0400
Subject: [PATCH] Add sharding key logic

---
 ee/app/models/merge_requests/approval_rule.rb | 12 ++++++
 .../merge_requests/approval_rule_spec.rb      | 37 +++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/ee/app/models/merge_requests/approval_rule.rb b/ee/app/models/merge_requests/approval_rule.rb
index 5d24cf7a9c5348..665a114424138e 100644
--- a/ee/app/models/merge_requests/approval_rule.rb
+++ b/ee/app/models/merge_requests/approval_rule.rb
@@ -43,6 +43,8 @@ class ApprovalRule < ApplicationRecord
       enum :origin, { group: 0, project: 1, merge_request: 2 }, prefix: :originates_from
     end
 
+    before_validation :set_sharding_key
+
     def approvers
       []
     end
@@ -57,6 +59,16 @@ def report_type
 
     private
 
+    # Currently this is narrowly scoped to issue
+    #   https://gitlab.com/gitlab-org/gitlab/-/issues/536752
+    # for the purposes of MR level approval rules, so we early return for now.
+    # As we implement more layers, this needs to be updated.
+    def set_sharding_key
+      return unless originates_from_merge_request?
+
+      self.project_id = merge_request.project_id
+    end
+
     def ensure_single_sharding_key
       return errors.add(:base, "Must have either `group_id` or `project_id`") if no_sharding_key?
 
diff --git a/ee/spec/models/merge_requests/approval_rule_spec.rb b/ee/spec/models/merge_requests/approval_rule_spec.rb
index 5fd721d76b74c0..c0dbfd80649f4f 100644
--- a/ee/spec/models/merge_requests/approval_rule_spec.rb
+++ b/ee/spec/models/merge_requests/approval_rule_spec.rb
@@ -9,6 +9,43 @@
 
   subject(:rule) { build(:merge_requests_approval_rule, attributes) }
 
+  describe 'setting sharding key' do
+    let_it_be(:project) { create(:project) }
+    let_it_be(:group) { create(:group) }
+    let_it_be(:merge_request) { create(:merge_request, source_project: project) }
+
+    context 'when originating from a merge request' do
+      let(:approval_rule) do
+        build(
+          :merge_requests_approval_rule,
+          origin: :merge_request,
+          merge_request: merge_request
+        )
+      end
+
+      it 'sets project_id to merge request project id before validating' do
+        expect(approval_rule.project_id).to be_nil
+        approval_rule.save!
+        expect(approval_rule.project_id).not_to be_nil
+        expect(approval_rule.project_id).to eq(merge_request.project_id)
+      end
+    end
+
+    context 'when not originating from a merge request' do
+      let(:approval_rule) do
+        build(:merge_requests_approval_rule, origin: :project)
+      end
+
+      it 'does not automatically set any sharding key' do
+        # This should fail validation since we're not setting any sharding key and the implementation to handle this
+        # case is out of scope for
+        #   https://gitlab.com/gitlab-org/gitlab/-/issues/536752
+        # Update as we add more layers of functionality.
+        expect { approval_rule.save! }.to raise_error(ActiveRecord::RecordInvalid)
+      end
+    end
+  end
+
   describe 'validations' do
     describe 'sharding key validation' do
       context 'with group_id' do
-- 
GitLab