Update vulnerability statistics

Update the vulnerability statistics based on the stat_diff of given
vulnerability object in Statistics::UpdateService service class by
using PostgreSQL's UPSERT functionality.

ee/app/models/vulnerabilities/stat_diff.rb

@@ -10,7 +10,7 @@ module Vulnerabilities
       update_required? ? changes.keys : []
     end
 
-    def change_values
+    def changed_values
       update_required? ? changes.values : []
     end
 

ee/app/models/vulnerabilities/statistic.rb

@@ -16,10 +16,31 @@ module Vulnerabilities
     validates :unknown, numericality: { greater_than_or_equal_to: 0 }
     validates :info, numericality: { greater_than_or_equal_to: 0 }
 
+    before_save :assign_letter_grade
+
     class << self
-      def update_stats_with(vulnerability)
-        true
+      # Takes an object which responds to `#[]` method call
+      # like an instance of ActiveRecord::Base or a Hash and
+      # returns the letter grade value for given object.
+      def letter_grade_for(object)
+        if object['critical'].to_i > 0
+          letter_grades[:f]
+        elsif object['high'].to_i > 0 || object['unknown'].to_i > 0
+          letter_grades[:d]
+        elsif object['medium'].to_i > 0
+          letter_grades[:c]
+        elsif object['low'].to_i > 0
+          letter_grades[:b]
+        else
+          letter_grades[:a]
+        end
       end
     end
+
+    private
+
+    def assign_letter_grade
+      self.letter_grade = self.class.letter_grade_for(self)
+    end
   end
 end

ee/app/models/vulnerability.rb

@@ -153,13 +153,13 @@ class Vulnerability < ApplicationRecord
   alias_method :after_note_created,   :after_note_changed
   alias_method :after_note_destroyed, :after_note_changed
 
+  def stat_diff
+    Vulnerabilities::StatDiff.new(self)
+  end
+
   private
 
   def user_notes_count_service
     @user_notes_count_service ||= Vulnerabilities::UserNotesCountService.new(self) # rubocop: disable CodeReuse/ServiceClass
   end
-
-  def stat_diff
-    Vulnerabilities::StatDiff.new(self)
-  end
 end

ee/app/services/vulnerabilities/statistics/update_service.rb

@@ -3,6 +3,29 @@
 module Vulnerabilities
   module Statistics
     class UpdateService
+      LETTER_GRADE_SQL = <<~SQL
+        CASE
+        WHEN TARGET.critical + EXCLUDED.critical > 0 THEN
+          #{Statistic.letter_grades[:f]}
+        WHEN TARGET.high + TARGET.unknown + EXCLUDED.high + EXCLUDED.unknown > 0 THEN
+          #{Statistic.letter_grades[:d]}
+        WHEN TARGET.medium + EXCLUDED.medium > 0 THEN
+          #{Statistic.letter_grades[:c]}
+        WHEN TARGET.low + EXCLUDED.low > 0 THEN
+          #{Statistic.letter_grades[:b]}
+        ELSE
+          #{Statistic.letter_grades[:a]}
+        END
+      SQL
+
+      UPSERT_SQL = <<~SQL
+        INSERT INTO #{Statistic.table_name} AS target (project_id, %{insert_attributes}, letter_grade, created_at, updated_at)
+          VALUES (%{project_id}, %{insert_values}, %{letter_grade}, now(), now())
+        ON CONFLICT (project_id)
+          DO UPDATE SET
+            %{update_values}, letter_grade = (#{LETTER_GRADE_SQL}), updated_at = now()
+      SQL
+
       def self.update_for(vulnerability)
         new(vulnerability).execute
       end
@@ -12,16 +35,53 @@ module Vulnerabilities
       end
 
       def execute
+        return unless stat_diff.update_required?
 
+        connection.execute(upsert_sql)
       end
 
       private
 
       attr_accessor :vulnerability
 
+      delegate :connection, to: Statistic, private: true
+      delegate :quote, :quote_column_name, to: :connection, private: true
+
       def stat_diff
         @stat_diff ||= vulnerability.stat_diff
       end
+
+      def upsert_sql
+        format(
+          UPSERT_SQL,
+          project_id: stat_diff.project_id,
+          insert_attributes: insert_attributes,
+          insert_values: insert_values,
+          letter_grade: letter_grade,
+          update_values: update_values
+        )
+      end
+
+      def insert_attributes
+        stat_diff.changed_attributes.map { |attribute| quote_column_name(attribute) }.join(', ')
+      end
+
+      def insert_values
+        stat_diff.changed_values.map { |value| quote(value) }.join(', ')
+      end
+
+      def letter_grade
+        quote(Statistic.letter_grade_for(stat_diff.changes))
+      end
+
+      def update_values
+        stat_diff.changes.map do |attribute, value|
+          column_name = quote_column_name(attribute)
+          quoted_value = quote(value)
+
+          "#{column_name} = GREATEST(TARGET.#{column_name} + #{quoted_value}, 0)"
+        end.join(', ')
+      end
     end
   end
 end

ee/spec/models/vulnerabilities/stat_diff_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Vulnerabilities::StatDiff do
+  using RSpec::Parameterized::TableSyntax
+
   let!(:vulnerability) { create(:vulnerability, :detected, severity: :high, title: 'Title') }
   let(:stat_diff) { described_class.new(vulnerability) }
   let(:update_vulnerability) {}
@@ -25,7 +27,25 @@ RSpec.describe Vulnerabilities::StatDiff do
 
       context 'when the severity is not changed' do
         context 'when the state is changed' do
-          shared_examples 'state changes' do |from:, to:, is_update_required:|
+          where(:from, :to, :is_update_required) do
+            'confirmed' | 'detected'  | false
+            'confirmed' | 'resolved'  | true
+            'confirmed' | 'dismissed' | true
+
+            'detected'  | 'confirmed' | false
+            'detected'  | 'resolved'  | true
+            'detected'  | 'dismissed' | true
+
+            'resolved'  | 'dismissed' | false
+            'resolved'  | 'confirmed' | true
+            'resolved'  | 'detected'  | true
+
+            'dismissed' | 'resolved'  | false
+            'dismissed' | 'confirmed' | true
+            'dismissed' | 'detected'  | true
+          end
+
+          with_them do
             let(:update_vulnerability) { vulnerability.update_attribute(:state, to) }
 
             before do
@@ -34,22 +54,6 @@ RSpec.describe Vulnerabilities::StatDiff do
 
             it { is_expected.to eq(is_update_required) }
           end
-
-          it_behaves_like 'state changes', from: 'confirmed', to: 'detected', is_update_required: false
-          it_behaves_like 'state changes', from: 'confirmed', to: 'resolved', is_update_required: true
-          it_behaves_like 'state changes', from: 'confirmed', to: 'dismissed', is_update_required: true
-
-          it_behaves_like 'state changes', from: 'detected', to: 'confirmed', is_update_required: false
-          it_behaves_like 'state changes', from: 'detected', to: 'resolved', is_update_required: true
-          it_behaves_like 'state changes', from: 'detected', to: 'dismissed', is_update_required: true
-
-          it_behaves_like 'state changes', from: 'resolved', to: 'dismissed', is_update_required: false
-          it_behaves_like 'state changes', from: 'resolved', to: 'confirmed', is_update_required: true
-          it_behaves_like 'state changes', from: 'resolved', to: 'detected', is_update_required: true
-
-          it_behaves_like 'state changes', from: 'dismissed', to: 'resolved', is_update_required: false
-          it_behaves_like 'state changes', from: 'dismissed', to: 'confirmed', is_update_required: true
-          it_behaves_like 'state changes', from: 'dismissed', to: 'detected', is_update_required: true
         end
 
         context 'when the state is not changed' do
@@ -81,7 +85,25 @@ RSpec.describe Vulnerabilities::StatDiff do
         end
 
         context 'when the state is changed' do
-          shared_examples 'state & severity change together' do |from:, to:, expected_changes:|
+          where(:from, :to, :expected_changes) do
+            'confirmed' | 'detected'  | { 'total' => 0, 'high' => -1, 'critical' => 1 }
+            'confirmed' | 'resolved'  | { 'total' => -1, 'high' => -1 }
+            'confirmed' | 'dismissed' | { 'total' => -1, 'high' => -1 }
+
+            'detected'  | 'confirmed' | { 'total' => 0, 'high' => -1, 'critical' => 1 }
+            'detected'  | 'resolved'  | { 'total' => -1, 'high' => -1 }
+            'detected'  | 'dismissed' | { 'total' => -1, 'high' => -1 }
+
+            'resolved'  | 'dismissed' | { 'total' => 0 }
+            'resolved'  | 'confirmed' | { 'total' => 1, 'critical' => 1 }
+            'resolved'  | 'detected'  | { 'total' => 1, 'critical' => 1 }
+
+            'dismissed' | 'resolved'  | { 'total' => 0 }
+            'dismissed' | 'confirmed' | { 'total' => 1, 'critical' => 1 }
+            'dismissed' | 'detected'  | { 'total' => 1, 'critical' => 1 }
+          end
+
+          with_them do
             let(:update_vulnerability) { vulnerability.update!(state: to, severity: :critical) }
 
             before do
@@ -90,28 +112,30 @@ RSpec.describe Vulnerabilities::StatDiff do
 
             it { is_expected.to eq(expected_changes) }
           end
-
-          it_behaves_like 'state & severity change together', from: 'confirmed', to: 'detected', expected_changes: { 'total' => 0, 'high' => -1, 'critical' => 1 }
-          it_behaves_like 'state & severity change together', from: 'confirmed', to: 'resolved', expected_changes: { 'total' => -1, 'high' => -1 }
-          it_behaves_like 'state & severity change together', from: 'confirmed', to: 'dismissed', expected_changes: { 'total' => -1, 'high' => -1 }
-
-          it_behaves_like 'state & severity change together', from: 'detected', to: 'confirmed', expected_changes: { 'total' => 0, 'high' => -1, 'critical' => 1 }
-          it_behaves_like 'state & severity change together', from: 'detected', to: 'resolved', expected_changes: { 'total' => -1, 'high' => -1 }
-          it_behaves_like 'state & severity change together', from: 'detected', to: 'dismissed', expected_changes: { 'total' => -1, 'high' => -1 }
-
-          it_behaves_like 'state & severity change together', from: 'resolved', to: 'dismissed', expected_changes: { 'total' => 0 }
-          it_behaves_like 'state & severity change together', from: 'resolved', to: 'confirmed', expected_changes: { 'total' => 1, 'critical' => 1 }
-          it_behaves_like 'state & severity change together', from: 'resolved', to: 'detected', expected_changes: { 'total' => 1, 'critical' => 1 }
-
-          it_behaves_like 'state & severity change together', from: 'dismissed', to: 'resolved', expected_changes: { 'total' => 0 }
-          it_behaves_like 'state & severity change together', from: 'dismissed', to: 'confirmed', expected_changes: { 'total' => 1, 'critical' => 1 }
-          it_behaves_like 'state & severity change together', from: 'dismissed', to: 'detected', expected_changes: { 'total' => 1, 'critical' => 1 }
         end
       end
 
       context 'when the severity is not changed' do
         context 'when the state is changed' do
-          shared_examples 'state changes' do |from:, to:, expected_changes:|
+          where(:from, :to, :expected_changes) do
+            'confirmed' | 'detected'  | { 'total' => 0 }
+            'confirmed' | 'resolved'  | { 'total' => -1, 'high' => -1 }
+            'confirmed' | 'dismissed' | { 'total' => -1, 'high' => -1 }
+
+            'detected'  | 'confirmed' | { 'total' => 0 }
+            'detected'  | 'resolved'  | { 'total' => -1, 'high' => -1 }
+            'detected'  | 'dismissed' | { 'total' => -1, 'high' => -1 }
+
+            'resolved'  | 'dismissed' | { 'total' => 0 }
+            'resolved'  | 'confirmed' | { 'total' => 1, 'high' => 1 }
+            'resolved'  | 'detected'  | { 'total' => 1, 'high' => 1 }
+
+            'dismissed' | 'resolved'  | { 'total' => 0 }
+            'dismissed' | 'confirmed' | { 'total' => 1, 'high' => 1 }
+            'dismissed' | 'detected'  | { 'total' => 1, 'high' => 1 }
+          end
+
+          with_them do
             let(:update_vulnerability) { vulnerability.update_attribute(:state, to) }
 
             before do
@@ -120,22 +144,6 @@ RSpec.describe Vulnerabilities::StatDiff do
 
             it { is_expected.to eq(expected_changes) }
           end
-
-          it_behaves_like 'state changes', from: 'confirmed', to: 'detected', expected_changes: { 'total' => 0 }
-          it_behaves_like 'state changes', from: 'confirmed', to: 'resolved', expected_changes: { 'total' => -1, 'high' => -1 }
-          it_behaves_like 'state changes', from: 'confirmed', to: 'dismissed', expected_changes: { 'total' => -1, 'high' => -1 }
-
-          it_behaves_like 'state changes', from: 'detected', to: 'confirmed', expected_changes: { 'total' => 0 }
-          it_behaves_like 'state changes', from: 'detected', to: 'resolved', expected_changes: { 'total' => -1, 'high' => -1 }
-          it_behaves_like 'state changes', from: 'detected', to: 'dismissed', expected_changes: { 'total' => -1, 'high' => -1 }
-
-          it_behaves_like 'state changes', from: 'resolved', to: 'dismissed', expected_changes: { 'total' => 0 }
-          it_behaves_like 'state changes', from: 'resolved', to: 'confirmed', expected_changes: { 'total' => 1, 'high' => 1 }
-          it_behaves_like 'state changes', from: 'resolved', to: 'detected', expected_changes: { 'total' => 1, 'high' => 1 }
-
-          it_behaves_like 'state changes', from: 'dismissed', to: 'resolved', expected_changes: { 'total' => 0 }
-          it_behaves_like 'state changes', from: 'dismissed', to: 'confirmed', expected_changes: { 'total' => 1, 'high' => 1 }
-          it_behaves_like 'state changes', from: 'dismissed', to: 'detected', expected_changes: { 'total' => 1, 'high' => 1 }
         end
 
         context 'when the state is not changed' do
@@ -172,8 +180,8 @@ RSpec.describe Vulnerabilities::StatDiff do
     end
   end
 
-  describe '#change_values' do
-    subject { stat_diff.change_values }
+  describe '#changed_values' do
+    subject { stat_diff.changed_values }
 
     context 'when there are changes' do
       let(:expected_values) { [-1, -1] }

ee/spec/models/vulnerabilities/statistic_spec.rb

@@ -18,7 +18,31 @@ RSpec.describe Vulnerabilities::Statistic do
     it { is_expected.to define_enum_for(:letter_grade).with_values(%i(a b c d f)) }
   end
 
-  describe '.update_stats_with' do
-    pending('This functionality is WIP')
+  describe '.before_save' do
+    describe '#assign_letter_grade' do
+      let(:statistic) { build(:vulnerability_statistic, letter_grade: nil, critical: 5) }
+
+      subject(:save_statistic) { statistic.save! }
+
+      it 'assigns the letter_grade' do
+        expect { save_statistic }.to change { statistic.letter_grade }.from(nil).to('f')
+      end
+    end
+  end
+
+  describe '.letter_grade_for' do
+    subject { described_class.letter_grade_for(object) }
+
+    context 'when the given object is an instance of Vulnerabilities::Statistic' do
+      let(:object) { build(:vulnerability_statistic, critical: 1) }
+
+      it { is_expected.to eq(4) }
+    end
+
+    context 'when the given object is a Hash' do
+      let(:object) { { 'high' => 1 } }
+
+      it { is_expected.to eq(3) }
+    end
   end
 end

ee/spec/models/vulnerability_spec.rb

@@ -424,4 +424,12 @@ RSpec.describe Vulnerability do
 
     it { is_expected.to eq(after_note_changed_method) }
   end
+
+  describe '#stat_diff' do
+    let(:vulnerability) { build(:vulnerability) }
+
+    subject { vulnerability.stat_diff }
+
+    it { is_expected.to be_an_instance_of(Vulnerabilities::StatDiff) }
+  end
 end

ee/spec/services/vulnerabilities/statistics/update_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Vulnerabilities::Statistics::UpdateService do
+  describe '.update_for' do
+    let(:mock_service_object) { instance_double(described_class, execute: true) }
+    let(:vulnerability) { instance_double(Vulnerability) }
+
+    subject(:update_stats) { described_class.update_for(vulnerability) }
+
+    before do
+      allow(described_class).to receive(:new).with(vulnerability).and_return(mock_service_object)
+    end
+
+    it 'instantiates an instance of service class and calls execute on it' do
+      update_stats
+
+      expect(mock_service_object).to have_received(:execute)
+    end
+  end
+
+  describe '#execute' do
+    let_it_be(:project) { create(:project) }
+    let_it_be(:statistic) { create(:vulnerability_statistic, project: project) }
+
+    let(:vulnerability) { create(:vulnerability, severity: :high, project: project) }
+    let(:stat_diff) { Vulnerabilities::StatDiff.new(vulnerability) }
+
+    subject(:update_stats) { described_class.new(vulnerability).execute }
+
+    context 'when the diff is empty' do
+      it 'does not change existing statistic entity' do
+        expect { update_stats }.not_to change { statistic.reload }
+      end
+    end
+
+    context 'when the diff is not empty' do
+      before do
+        vulnerability.update_attribute(:severity, :critical)
+      end
+
+      context 'when there is already a record in the database' do
+        it 'changes the existing statistic entity' do
+          expect { update_stats }.to change { statistic.reload.critical }.by(1)
+                                 .and not_change { statistic.reload.high }
+        end
+      end
+
+      context 'when there is no existing record in the database' do
+        before do
+          statistic.destroy!
+        end
+
+        it 'creates a new record in the database' do
+          expect { update_stats }.to change { Vulnerabilities::Statistic.count }.by(1)
+        end
+      end
+    end
+  end
+end