Add weights to burnup chart data

Add weights to graph data

Use union of events to reduce queries

Extract scopes and use Gitlab::SQL::Union

Refactor and extract logic out of resource events

Apply suggestion to app/models/concerns/issue_resource_event.rb
Apply suggestion to ee/app/services/burnup_chart_service.rb
Apply suggestion to ee/app/services/burnup_chart_service.rb
Change scope name

Make attribute readers private

Use field names instead of index access

app/models/concerns/issue_resource_event.rb

+# frozen_string_literal: true
+
+module IssueResourceEvent
+  extend ActiveSupport::Concern
+
+  included do
+    belongs_to :issue
+
+    scope :by_issue, ->(issue) { where(issue_id: issue.id) }
+
+    scope :by_issue_ids_and_created_at_earlier_or_equal_to, ->(issue_ids, time) { where(issue_id: issue_ids).where('created_at <= ?', time) }
+  end
+end

app/models/concerns/merge_request_resource_event.rb

+# frozen_string_literal: true
+
+module MergeRequestResourceEvent
+  extend ActiveSupport::Concern
+
+  included do
+    belongs_to :merge_request
+
+    scope :by_merge_request, ->(merge_request) { where(merge_request_id: merge_request.id) }
+  end
+end

app/models/resource_label_event.rb

@@ -2,16 +2,14 @@
 
 class ResourceLabelEvent < ResourceEvent
   include CacheMarkdownField
+  include IssueResourceEvent
+  include MergeRequestResourceEvent
 
   cache_markdown_field :reference
 
-  belongs_to :issue
-  belongs_to :merge_request
   belongs_to :label
 
   scope :inc_relations, -> { includes(:label, :user) }
-  scope :by_issue, ->(issue) { where(issue_id: issue.id) }
-  scope :by_merge_request, ->(merge_request) { where(merge_request_id: merge_request.id) }
 
   validates :label, presence: { unless: :importing? }, on: :create
   validate :exactly_one_issuable

app/models/resource_milestone_event.rb

@@ -2,14 +2,11 @@
 
 class ResourceMilestoneEvent < ResourceEvent
   include IgnorableColumns
+  include IssueResourceEvent
+  include MergeRequestResourceEvent
 
-  belongs_to :issue
-  belongs_to :merge_request
   belongs_to :milestone
 
-  scope :by_issue, ->(issue) { where(issue_id: issue.id) }
-  scope :by_merge_request, ->(merge_request) { where(merge_request_id: merge_request.id) }
-
   validate :exactly_one_issuable
 
   enum action: {

app/models/resource_weight_event.rb

@@ -3,7 +3,5 @@
 class ResourceWeightEvent < ResourceEvent
   validates :issue, presence: true
 
-  belongs_to :issue
-
-  scope :by_issue, ->(issue) { where(issue_id: issue.id) }
+  include IssueResourceEvent
 end

ee/app/services/burnup_chart_service.rb

@@ -3,7 +3,12 @@
 class BurnupChartService
   include Gitlab::Utils::StrongMemoize
 
-  attr_reader :milestone, :start_date, :due_date, :end_date, :user
+  EVENT_TYPE = 'event_type'.freeze
+  CREATED_AT = 'created_at'.freeze
+  MILESTONE_ID = 'value'.freeze
+  WEIGHT = 'value'.freeze
+  ACTION = 'action'.freeze
+  ISSUE_ID = 'issue_id'.freeze
 
   def initialize(milestone:, user:)
     @user = user
@@ -15,10 +20,12 @@ class BurnupChartService
   def execute
     return [] unless milestone.burnup_charts_available?
 
+    return [] unless can_read_milestone?
+
     events = []
     assigned_milestones = {}
 
-    resource_milestone_events.each do |event|
+    resource_events.each do |event|
       handle_added_milestone(event, assigned_milestones)
 
       events << create_burnup_graph_event_by(event, assigned_milestones)
@@ -31,27 +38,76 @@ class BurnupChartService
 
   private
 
+  attr_reader :milestone, :start_date, :due_date, :end_date, :user
+
+  def can_read_milestone?
+    Ability.allowed?(user, :read_milestone, milestone_parent)
+  end
+
+  def milestone_parent
+    if milestone.group_milestone?
+      milestone.group
+    else
+      milestone.project
+    end
+  end
+
   def handle_added_milestone(event, assigned_milestones)
-    if event.add?
-      assigned_milestones[event.issue_id] = event.milestone_id
+    if add_milestone?(event)
+      assigned_milestones[event[ISSUE_ID]] = event[MILESTONE_ID]
     end
   end
 
   def handle_removed_milestone(event, assigned_milestones)
-    if event.remove?
-      assigned_milestones[event.issue_id] = nil
+    if remove_milestone?(event)
+      assigned_milestones[event[ISSUE_ID]] = nil
     end
   end
 
+  def resource_events
+    strong_memoize(:resource_events) do
+      union = Gitlab::SQL::Union.new(all_events).to_sql # rubocop: disable Gitlab/Union
+      ActiveRecord::Base.connection.execute("(#{union}) ORDER BY created_at")
+    end
+  end
+
+  def all_events
+    [milestone_events, weight_events]
+  end
+
+  def weight_events
+    ResourceWeightEvent.by_issue_ids_and_created_at_earlier_or_equal_to(relevant_issue_ids, end_time)
+      .select('\'weight\' as event_type, created_at, weight as value, null as action, issue_id')
+  end
+
+  def milestone_events
+    ResourceMilestoneEvent.by_issue_ids_and_created_at_earlier_or_equal_to(relevant_issue_ids, end_time)
+      .select('\'milestone\' as event_type, created_at, milestone_id as value, action, issue_id')
+  end
+
   def create_burnup_graph_event_by(event, assigned_milestones)
     {
-      created_at: event.created_at,
-      action: event.action,
+      event_type: event[EVENT_TYPE],
+      created_at: event[CREATED_AT],
+      action: action_of(event),
       milestone_id: milestone_id_of(event, assigned_milestones),
-      issue_id: event.issue_id
+      issue_id: event[ISSUE_ID],
+      weight: weight_of(event)
     }
   end
 
+  def action_of(event)
+    return unless milestone_event?(event)
+
+    ResourceMilestoneEvent.actions.key(event[ACTION])
+  end
+
+  def weight_of(event)
+    return unless event[EVENT_TYPE] == 'weight'
+
+    event[WEIGHT]
+  end
+
   def assign_dates_by_milestone
     @start_date = milestone.start_date
     @due_date = milestone.due_date
@@ -63,45 +119,43 @@ class BurnupChartService
   end
 
   def milestone_id_of(event, assigned_milestones)
-    if event.remove? && event.milestone_id.nil?
-      return assigned_milestones[event.issue_id]
+    if remove_milestone?(event) && event[MILESTONE_ID].nil?
+      return assigned_milestones[event[ISSUE_ID]]
     end
 
-    event.milestone_id
+    event[MILESTONE_ID]
   end
 
-  # This is just temporarily disabled until https://gitlab.com/gitlab-org/gitlab/-/merge_requests/29578
-  # is merged.
-  # rubocop: disable CodeReuse/ActiveRecord
-  def resource_milestone_events
-    # Here we use the relevant issues to get *all* milestone events for
-    # them.
-    strong_memoize(:resource_milestone_events) do
-      ResourceMilestoneEvent
-        .where(issue_id: relevant_issue_ids)
-        .where('created_at <= ?', end_time)
-        .order(:created_at)
-    end
+  def add_milestone?(event)
+    return unless milestone_event?(event)
+
+    event[ACTION] == ResourceMilestoneEvent.actions[:add]
+  end
+
+  def remove_milestone?(event)
+    return unless milestone_event?(event)
+
+    event[ACTION] == ResourceMilestoneEvent.actions[:remove]
   end
 
+  def milestone_event?(event)
+    event[EVENT_TYPE] == 'milestone'
+  end
+
+  # rubocop: disable CodeReuse/ActiveRecord
   def relevant_issue_ids
     # We are using all resource milestone events where the
     # milestone in question was added to identify the relevant
     # issues.
     strong_memoize(:relevant_issue_ids) do
-      ids = ResourceMilestoneEvent
-              .select(:issue_id)
-              .where(milestone_id: milestone.id)
-              .where(action: :add)
-              .distinct
-
-      # We need to perform an additional check whether all these
-      # issues are visible to the given user
-      IssuesFinder.new(user)
-        .execute.select(:id).where(id: ids)
+      ResourceMilestoneEvent
+            .select(:issue_id)
+            .where(milestone_id: milestone.id)
+            .where(action: :add)
+            .distinct
     end
+    # rubocop: enable CodeReuse/ActiveRecord
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
   def end_time
     @end_time ||= @end_date.end_of_day

ee/spec/fixtures/api/schemas/burnup_event.json

@@ -8,6 +8,12 @@
     "issue_id"
   ],
   "properties": {
+    "event_type": {
+      "type": { "enum": [ "milestone", "weight" ] }
+    },
+    "weight": {
+      "type": ["integer", "null"]
+    },
     "milestone_id": {
       "type": ["integer", "null"]
     },
@@ -15,7 +21,7 @@
       "type": "integer"
     },
     "action": {
-      "type": "string"
+      "type": ["string", "null"]
     },
     "created_at": {
       "type": "date"

ee/spec/services/burnup_chart_service_spec.rb

@@ -19,8 +19,14 @@ describe BurnupChartService do
   let_it_be(:issue3) { create(:issue, project: project, milestone: milestone1) }
   let_it_be(:other_issue) { create(:issue, project: project) }
 
+  let_it_be(:weight_event1) { create(:resource_weight_event, issue: issue1, weight: 9, created_at: start_date + 2.seconds) }
+  let_it_be(:weight_event2) { create(:resource_weight_event, issue: issue2, weight: 3, created_at: start_date + 3.seconds) }
+  let_it_be(:weight_event3) { create(:resource_weight_event, issue: issue3, weight: 1, created_at: start_date + 2.days) }
+  let_it_be(:weight_event4) { create(:resource_weight_event, issue: issue3, weight: 2, created_at: start_date + 2.days + 23.hours + 59.minutes + 59.seconds) }
+  let_it_be(:weight_event5) { create(:resource_weight_event, issue: issue3, weight: 7, created_at: start_date + 4.days + 1.second) }
+
   let_it_be(:event1) { create(:resource_milestone_event, issue: issue1, action: :add, milestone: milestone1, created_at: start_date + 1.second) }
-  let_it_be(:event2) { create(:resource_milestone_event, issue: issue2, action: :add, milestone: milestone1, created_at: start_date + 2.seconds) }
+  let_it_be(:event2) { create(:resource_milestone_event, issue: issue2, action: :add, milestone: milestone1, created_at: start_date + 4.seconds) }
   let_it_be(:event3) { create(:resource_milestone_event, issue: issue3, action: :add, milestone: milestone1, created_at: start_date + 1.day) }
   let_it_be(:event4) { create(:resource_milestone_event, issue: issue3, action: :remove, milestone: nil, created_at: start_date + 2.days + 1.second) }
   let_it_be(:event5) { create(:resource_milestone_event, issue: issue3, action: :add, milestone: milestone2, created_at: start_date + 3.days) }
@@ -43,17 +49,20 @@ describe BurnupChartService do
 
       data = described_class.new(milestone: milestone1, user: user).execute
 
-      expected_events = [
-        { action: 'add', issue_id: issue1.id, milestone_id: milestone2.id, created_at: start_date.beginning_of_day - 1.second },
-        { action: 'add', issue_id: issue1.id, milestone_id: milestone1.id, created_at: start_date + 1.second },
-        { action: 'add', issue_id: issue2.id, milestone_id: milestone1.id, created_at: start_date + 2.seconds },
-        { action: 'add', issue_id: issue3.id, milestone_id: milestone1.id, created_at: start_date + 1.day },
-        { action: 'remove', issue_id: issue3.id, milestone_id: milestone1.id, created_at: start_date + 2.days + 1.second },
-        { action: 'add', issue_id: issue3.id, milestone_id: milestone2.id, created_at: start_date + 3.days },
-        { action: 'remove', issue_id: issue3.id, milestone_id: milestone2.id, created_at: start_date + 4.days }
-      ]
-
-      expect(data).to eq(expected_events)
+      expect(data.size).to eq(12)
+
+      expect(data[0]).to include(event_type: 'milestone', action: 'add', issue_id: issue1.id, milestone_id: milestone2.id, created_at: start_date.beginning_of_day - 1.second)
+      expect(data[1]).to include(event_type: 'milestone', action: 'add', issue_id: issue1.id, milestone_id: milestone1.id, created_at: start_date + 1.second)
+      expect(data[2]).to include(event_type: 'weight', issue_id: issue1.id, weight: 9, created_at: start_date + 2.seconds)
+      expect(data[3]).to include(event_type: 'weight', issue_id: issue2.id, weight: 3, created_at: start_date + 3.seconds)
+      expect(data[4]).to include(event_type: 'milestone', action: 'add', issue_id: issue2.id, milestone_id: milestone1.id, created_at: start_date + 4.seconds)
+      expect(data[5]).to include(event_type: 'milestone', action: 'add', issue_id: issue3.id, milestone_id: milestone1.id, created_at: start_date + 1.day)
+      expect(data[6]).to include(event_type: 'weight', issue_id: issue3.id, weight: 1, created_at: start_date + 2.days)
+      expect(data[7]).to include(event_type: 'milestone', action: 'remove', issue_id: issue3.id, milestone_id: milestone1.id, created_at: start_date + 2.days + 1.second)
+      expect(data[8]).to include(event_type: 'weight', issue_id: issue3.id, weight: 2, created_at: start_date + 2.days + 23.hours + 59.minutes + 59.seconds)
+      expect(data[9]).to include(event_type: 'milestone', action: 'add', issue_id: issue3.id, milestone_id: milestone2.id, created_at: start_date + 3.days)
+      expect(data[10]).to include(event_type: 'milestone', action: 'remove', issue_id: issue3.id, milestone_id: milestone2.id, created_at: start_date + 4.days)
+      expect(data[11]).to include(event_type: 'weight', issue_id: issue3.id, weight: 7, created_at: start_date + 4.days + 1.second)
     end
 
     it 'excludes issues which should not be visible to the user ' do

ee/spec/support/shared_examples/requests/api/milestone_burnup_shared_examples.rb

@@ -17,13 +17,11 @@ RSpec.shared_examples 'group and project milestone burnups' do |route_definition
       expect(json_response).to be_an Array
       expect(json_response).to match_schema('burnup_events', dir: 'ee')
 
-      expected_events = [
-        { 'issue_id' => issue1.id, 'milestone_id' => milestone.id, 'action' => 'add', 'created_at' => event1.created_at.iso8601(3) },
-        { 'issue_id' => issue2.id, 'milestone_id' => milestone.id, 'action' => 'add', 'created_at' => event2.created_at.iso8601(3) },
-        { 'issue_id' => issue1.id, 'milestone_id' => milestone.id, 'action' => 'remove', 'created_at' => event3.created_at.iso8601(3) }
-      ]
+      expect(json_response.size).to eq(3)
 
-      expect(json_response).to eq(expected_events)
+      expect(json_response.first).to include('issue_id' => issue1.id, 'milestone_id' => milestone.id, 'action' => 'add', 'created_at' => event_time - 1.hour)
+      expect(json_response.second).to include('issue_id' => issue2.id, 'milestone_id' => milestone.id, 'action' => 'add', 'created_at' => event_time + 1.hour)
+      expect(json_response.third).to include('issue_id' => issue1.id, 'milestone_id' => milestone.id, 'action' => 'remove', 'created_at' => event_time + 2.hours)
     end
 
     it 'returns 404 when user is not authorized to read milestone' do

spec/models/resource_label_event_spec.rb

@@ -15,9 +15,6 @@ RSpec.describe ResourceLabelEvent, type: :model do
   it_behaves_like 'a resource event for merge requests'
 
   describe 'associations' do
-    it { is_expected.to belong_to(:user) }
-    it { is_expected.to belong_to(:issue) }
-    it { is_expected.to belong_to(:merge_request) }
     it { is_expected.to belong_to(:label) }
   end
 

spec/support/shared_examples/resource_events.rb

@@ -83,6 +83,24 @@ shared_examples 'a resource event for issues' do
       expect(events).to be_empty
     end
   end
+
+  describe '.by_issue_ids_and_created_at_earlier_or_equal_to' do
+    let_it_be(:event1) { create(described_class.name.underscore.to_sym, issue: issue1, created_at: '2020-03-10') }
+    let_it_be(:event2) { create(described_class.name.underscore.to_sym, issue: issue2, created_at: '2020-03-10') }
+    let_it_be(:event3) { create(described_class.name.underscore.to_sym, issue: issue1, created_at: '2020-03-12') }
+
+    it 'returns the expected records for an issue with events' do
+      events = described_class.by_issue_ids_and_created_at_earlier_or_equal_to([issue1.id, issue2.id], '2020-03-11 23:59:59')
+
+      expect(events).to contain_exactly(event1, event2)
+    end
+
+    it 'returns the expected records for an issue with no events' do
+      events = described_class.by_issue_ids_and_created_at_earlier_or_equal_to(issue3, '2020-03-12')
+
+      expect(events).to be_empty
+    end
+  end
 end
 
 shared_examples 'a resource event for merge requests' do