Merge branch 'bwill/refactor-sbom-post-ingestion' into 'master'

Refactor SBoM ingestion hash lookup

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/148326



Merged-by: Bojan Marjanovic <bmarjanovic@gitlab.com>
Approved-by: Oscar Tovar <otovar@gitlab.com>
Approved-by: Bojan Marjanovic <bmarjanovic@gitlab.com>
Reviewed-by: Brian Williams <bwilliams@gitlab.com>
Reviewed-by: Oscar Tovar <otovar@gitlab.com>
Co-authored-by: Brian Williams <bwilliams@gitlab.com>

ee/app/services/sbom/ingestion/tasks/base.rb

@@ -4,6 +4,9 @@ module Sbom
   module Ingestion
     module Tasks
       class Base
+        include Gitlab::Utils::StrongMemoize
+        include Gitlab::Ingestion::BulkInsertableTask
+
         def self.execute(pipeline, occurrence_maps)
           new(pipeline, occurrence_maps).execute
         end
@@ -13,15 +16,53 @@ def initialize(pipeline, occurrence_maps)
           @occurrence_maps = occurrence_maps
         end
 
-        def execute
-          raise NoMethodError, "Implement the `execute` template method!"
-        end
-
         private
 
         attr_reader :pipeline, :occurrence_maps
 
         delegate :project, to: :pipeline, private: true
+
+        def insertable_maps
+          occurrence_maps
+        end
+
+        def each_pair
+          validate_unique_by!
+
+          return_data.each do |row|
+            occurrence_maps_for_row(row).each { |map| yield map, row }
+          end
+        end
+
+        def occurrence_maps_for_row(row)
+          indexed_occurrence_maps[grouping_key_for_row(row)]
+        end
+
+        def indexed_occurrence_maps
+          insertable_maps.group_by { |map| grouping_key_for_map(map) }
+        end
+        strong_memoize_attr :indexed_occurrence_maps
+
+        def grouping_key_for_map(occurrence_map)
+          occurrence_map.to_h.values_at(*unique_by)
+        end
+
+        def unique_attr_indices
+          unique_by.map { |attr| uses.find_index(attr) }
+        end
+        strong_memoize_attr :unique_attr_indices
+
+        def grouping_key_for_row(row)
+          unique_attr_indices.map { |index| row[index] }
+        end
+
+        def validate_unique_by!
+          raise ArgumentError, '#each_pair can only be used with unique_by attributes' if unique_by.blank?
+
+          return if unique_by.all? { |attr| uses.include?(attr) }
+
+          raise ArgumentError, 'All unique_by attributes must be included in returned columns'
+        end
       end
     end
   end

ee/app/services/sbom/ingestion/tasks/ingest_component_versions.rb

@@ -4,8 +4,6 @@ module Sbom
   module Ingestion
     module Tasks
       class IngestComponentVersions < Base
-        include Gitlab::Ingestion::BulkInsertableTask
-
         COMPONENT_VERSION_ATTRIBUTES = %i[component_id version].freeze
 
         self.model = Sbom::ComponentVersion
@@ -15,33 +13,19 @@ class IngestComponentVersions < Base
         private
 
         def after_ingest
-          return_data.each do |component_version_id, component_id, version|
-            maps_with(component_id, version)&.each do |occurrence_map|
-              occurrence_map.component_version_id = component_version_id
-            end
+          each_pair do |occurrence_map, row|
+            occurrence_map.component_version_id = row.first
           end
         end
 
         def attributes
-          valid_occurrence_maps.map do |occurrence_map|
+          insertable_maps.map do |occurrence_map|
             occurrence_map.to_h.slice(*COMPONENT_VERSION_ATTRIBUTES)
           end
         end
 
-        def valid_occurrence_maps
-          @valid_occurrence_maps ||= occurrence_maps.filter(&:version_present?)
-        end
-
-        def maps_with(component_id, version)
-          grouped_maps[[component_id, version]]
-        end
-
-        def grouped_maps
-          @grouped_maps ||= valid_occurrence_maps.group_by do |occurrence_map|
-            report_component = occurrence_map.report_component
-
-            [occurrence_map.component_id, report_component.version]
-          end
+        def insertable_maps
+          super.filter(&:version_present?)
         end
       end
     end

ee/app/services/sbom/ingestion/tasks/ingest_components.rb

@@ -4,8 +4,6 @@ module Sbom
   module Ingestion
     module Tasks
       class IngestComponents < Base
-        include Gitlab::Ingestion::BulkInsertableTask
-
         COMPONENT_ATTRIBUTES = %i[name purl_type component_type].freeze
 
         self.model = Sbom::Component
@@ -15,28 +13,16 @@ class IngestComponents < Base
         private
 
         def after_ingest
-          return_data.each do |id, name, purl_type, component_type|
-            maps_with(name, purl_type, component_type)&.each do |occurrence_map|
-              occurrence_map.component_id = id
-            end
+          each_pair do |occurrence_map, row|
+            occurrence_map.component_id = row.first
           end
         end
 
         def attributes
-          occurrence_maps.map do |occurrence_map|
+          insertable_maps.map do |occurrence_map|
             occurrence_map.to_h.slice(*COMPONENT_ATTRIBUTES)
           end
         end
-
-        def maps_with(name, purl_type, component_type)
-          grouped_maps[[name, purl_type, component_type]]
-        end
-
-        def grouped_maps
-          @grouped_maps ||= occurrence_maps.group_by do |occurrence_map|
-            occurrence_map.to_h.values_at(:name, :purl_type, :component_type)
-          end
-        end
       end
     end
   end

ee/app/services/sbom/ingestion/tasks/ingest_occurrences.rb

@@ -4,18 +4,17 @@ module Sbom
   module Ingestion
     module Tasks
       class IngestOccurrences < Base
-        include Gitlab::Ingestion::BulkInsertableTask
         include Gitlab::Utils::StrongMemoize
 
         self.model = Sbom::Occurrence
-        self.unique_by = :uuid
-        self.uses = :id
+        self.unique_by = %i[uuid].freeze
+        self.uses = %i[id uuid].freeze
 
         private
 
         def after_ingest
-          return_data.each_with_index do |occurrence_id, index|
-            occurrence_maps[index].occurrence_id = occurrence_id
+          each_pair do |occurrence_map, row|
+            occurrence_map.occurrence_id = row.first
           end
         end
 
@@ -58,6 +57,10 @@ def uuid(occurrence_map)
           ::Sbom::OccurrenceUUID.generate(**uuid_attributes)
         end
 
+        def grouping_key_for_map(map)
+          [uuid(map)]
+        end
+
         def licenses
           Licenses.new(project, occurrence_maps)
         end

ee/app/services/sbom/ingestion/tasks/ingest_occurrences_vulnerabilities.rb

@@ -4,15 +4,13 @@ module Sbom
   module Ingestion
     module Tasks
       class IngestOccurrencesVulnerabilities < Base
-        include Gitlab::Ingestion::BulkInsertableTask
-
         self.model = Sbom::OccurrencesVulnerability
-        self.unique_by = %i[sbom_occurrence_id vulnerability_id]
+        self.unique_by = %i[sbom_occurrence_id vulnerability_id].freeze
 
         private
 
         def attributes
-          occurrence_maps.flat_map do |occurrence_map|
+          insertable_maps.flat_map do |occurrence_map|
             occurrence_map.vulnerability_ids.map do |vulnerability_id|
               {
                 sbom_occurrence_id: occurrence_map.occurrence_id,

ee/app/services/sbom/ingestion/tasks/ingest_source_packages.rb

@@ -4,8 +4,6 @@ module Sbom
   module Ingestion
     module Tasks
       class IngestSourcePackages < Base
-        include Gitlab::Ingestion::BulkInsertableTask
-
         SOURCE_PACKAGE_ATTRIBUTES = %i[name purl_type].freeze
 
         self.model = Sbom::SourcePackage
@@ -15,15 +13,13 @@ class IngestSourcePackages < Base
         private
 
         def after_ingest
-          return_data.each do |source_package_id, source_package_name, purl_type|
-            maps_with(source_package_name, purl_type)&.each do |occurrence_map|
-              occurrence_map.source_package_id = source_package_id
-            end
+          each_pair do |occurrence_map, row|
+            occurrence_map.source_package_id = row.first
           end
         end
 
         def attributes
-          valid_occurrence_maps.map do |occurrence_map|
+          insertable_maps.map do |occurrence_map|
             {
               name: occurrence_map.source_package_name,
               purl_type: occurrence_map.purl_type
@@ -31,20 +27,12 @@ def attributes
           end
         end
 
-        def valid_occurrence_maps
-          @valid_occurrence_maps ||= occurrence_maps.filter(&:source_package_name)
-        end
-
-        def maps_with(source_package_name, purl_type)
-          grouped_maps[[source_package_name, purl_type]]
+        def insertable_maps
+          super.filter(&:source_package_name)
         end
 
-        def grouped_maps
-          @grouped_maps ||= valid_occurrence_maps.group_by do |occurrence_map|
-            report_component = occurrence_map.report_component
-
-            [report_component.source_package_name, report_component.purl_type]
-          end
+        def grouping_key_for_map(map)
+          [map.source_package_name, map.purl_type]
         end
       end
     end

ee/app/services/sbom/ingestion/tasks/ingest_sources.rb

@@ -8,10 +8,8 @@ class IngestSources < Base
         # so all occurrence maps in the batch will use the same one.
         # This is likely to change in the future, so the interface
         # allows for multiple sources.
-        include Gitlab::Ingestion::BulkInsertableTask
-
         self.model = Sbom::Source
-        self.uses = :id
+        self.uses = %i[id].freeze
         self.unique_by = %i[source_type source].freeze
 
         private

ee/spec/services/sbom/ingestion/tasks/base_spec.rb

@@ -5,9 +5,51 @@
 RSpec.describe Sbom::Ingestion::Tasks::Base, feature_category: :dependency_management do
   let(:pipeline) { instance_double('Ci::Pipeline') }
   let(:occurrence_maps) { [instance_double('Sbom::Ingestion::OccurrenceMap')] }
-  let(:implementation) { Class.new(described_class) }
 
-  it 'raises error when execute is not implemented' do
-    expect { implementation.execute(pipeline, occurrence_maps) }.to raise_error(NoMethodError)
+  describe '#each_pair' do
+    context 'when implementation does not have unique_by columns in uses' do
+      let(:implementation) do
+        Class.new(described_class) do
+          self.model = Sbom::ComponentVersion
+          self.unique_by = %i[component_id version].freeze
+          self.uses = %i[id].freeze
+
+          def execute
+            each_pair do |map, row|
+              map.id = row.first
+            end
+          end
+        end
+      end
+
+      it 'raises an ArgumentError' do
+        expect { implementation.execute(pipeline, occurrence_maps) }.to raise_error(
+          ArgumentError,
+          'All unique_by attributes must be included in returned columns'
+        )
+      end
+    end
+
+    context 'when implementation does not have unique_by' do
+      let(:implementation) do
+        Class.new(described_class) do
+          self.model = Sbom::ComponentVersion
+          self.uses = %i[id].freeze
+
+          def execute
+            each_pair do |map, row|
+              map.id = row.first
+            end
+          end
+        end
+      end
+
+      it 'raises an ArgumentError' do
+        expect { implementation.execute(pipeline, occurrence_maps) }.to raise_error(
+          ArgumentError,
+          '#each_pair can only be used with unique_by attributes'
+        )
+      end
+    end
   end
 end