Update indexes and scopes for agent_id/cluster_id in Finding

db/post_migrate/20211217120000_modify_kubernetes_resource_location_index_to_vulnerability_occurrences.rb

+# frozen_string_literal: true
+
+class ModifyKubernetesResourceLocationIndexToVulnerabilityOccurrences < Gitlab::Database::Migration[1.0]
+  disable_ddl_transaction!
+
+  OLD_CLUSTER_ID_INDEX_NAME = 'index_vulnerability_occurrences_on_location_cluster_id'
+  OLD_AGENT_ID_INDEX_NAME = 'index_vulnerability_occurrences_on_location_agent_id'
+
+  NEW_CLUSTER_ID_INDEX_NAME = 'index_vulnerability_occurrences_on_location_k8s_cluster_id'
+  NEW_AGENT_ID_INDEX_NAME = 'index_vulnerability_occurrences_on_location_k8s_agent_id'
+
+  def up
+    add_concurrent_index :vulnerability_occurrences, "(location -> 'kubernetes_resource' -> 'cluster_id')",
+                         using: 'GIN',
+                         where: 'report_type = 7',
+                         name: NEW_CLUSTER_ID_INDEX_NAME
+
+    add_concurrent_index :vulnerability_occurrences, "(location -> 'kubernetes_resource' -> 'agent_id')",
+                         using: 'GIN',
+                         where: 'report_type = 7',
+                         name: NEW_AGENT_ID_INDEX_NAME
+
+    remove_concurrent_index_by_name :vulnerability_occurrences, OLD_CLUSTER_ID_INDEX_NAME
+    remove_concurrent_index_by_name :vulnerability_occurrences, OLD_AGENT_ID_INDEX_NAME
+  end
+
+  def down
+    add_concurrent_index :vulnerability_occurrences, "(location -> 'cluster_id')",
+                         using: 'GIN',
+                         where: 'report_type = 7',
+                         name: OLD_CLUSTER_ID_INDEX_NAME
+
+    add_concurrent_index :vulnerability_occurrences, "(location -> 'agent_id')",
+                         using: 'GIN',
+                         where: 'report_type = 7',
+                         name: OLD_AGENT_ID_INDEX_NAME
+
+    remove_concurrent_index_by_name :vulnerability_occurrences, NEW_CLUSTER_ID_INDEX_NAME
+    remove_concurrent_index_by_name :vulnerability_occurrences, NEW_AGENT_ID_INDEX_NAME
+  end
+end

db/schema_migrations/20211217120000

+d4360d6057602ec1f5e6e9d11c93cfbb16d878e9ecd4d5bfb1bed1c01e14c7a3
\ No newline at end of file

db/structure.sql

@@ -27873,11 +27873,11 @@ CREATE INDEX index_vulnerability_occurrences_deduplication ON vulnerability_occu
 
 CREATE INDEX index_vulnerability_occurrences_for_issue_links_migration ON vulnerability_occurrences USING btree (project_id, report_type, encode(project_fingerprint, 'hex'::text));
 
-CREATE INDEX index_vulnerability_occurrences_on_location_agent_id ON vulnerability_occurrences USING gin (((location -> 'agent_id'::text))) WHERE (report_type = 7);
+CREATE INDEX index_vulnerability_occurrences_on_location_image ON vulnerability_occurrences USING gin (((location -> 'image'::text))) WHERE (report_type = ANY (ARRAY[2, 7]));
 
-CREATE INDEX index_vulnerability_occurrences_on_location_cluster_id ON vulnerability_occurrences USING gin (((location -> 'cluster_id'::text))) WHERE (report_type = 7);
+CREATE INDEX index_vulnerability_occurrences_on_location_k8s_agent_id ON vulnerability_occurrences USING gin ((((location -> 'kubernetes_resource'::text) -> 'agent_id'::text))) WHERE (report_type = 7);
 
-CREATE INDEX index_vulnerability_occurrences_on_location_image ON vulnerability_occurrences USING gin (((location -> 'image'::text))) WHERE (report_type = ANY (ARRAY[2, 7]));
+CREATE INDEX index_vulnerability_occurrences_on_location_k8s_cluster_id ON vulnerability_occurrences USING gin ((((location -> 'kubernetes_resource'::text) -> 'cluster_id'::text))) WHERE (report_type = 7);
 
 CREATE INDEX index_vulnerability_occurrences_on_migrated_to_new_structure ON vulnerability_occurrences USING btree (migrated_to_new_structure, id);
 

ee/app/models/vulnerabilities/finding.rb

@@ -103,11 +103,11 @@ module Vulnerabilities
     end
     scope :by_location_cluster, -> (cluster_ids) do
       where(report_type: 'cluster_image_scanning')
-        .where("vulnerability_occurrences.location -> 'cluster_id' ?| array[:cluster_ids]", cluster_ids: cluster_ids)
+        .where("vulnerability_occurrences.location -> 'kubernetes_resource' -> 'cluster_id' ?| array[:cluster_ids]", cluster_ids: cluster_ids)
     end
     scope :by_location_cluster_agent, -> (agent_ids) do
       where(report_type: 'cluster_image_scanning')
-        .where("vulnerability_occurrences.location -> 'agent_id' ?| array[:agent_ids]", agent_ids: agent_ids)
+        .where("vulnerability_occurrences.location -> 'kubernetes_resource' -> 'agent_id' ?| array[:agent_ids]", agent_ids: agent_ids)
     end
 
     def self.counted_by_severity

ee/spec/factories/vulnerabilities/findings.rb

@@ -584,9 +584,11 @@ FactoryBot.define do
           },
           "operating_system": "alpine 3.7",
           "image": "alpine:3.7",
+          "kubernetes_resource": {
             "cluster_id": "1",
             "agent_id": "46357"
           }
+        }
         finding.raw_metadata = {
           "category": "cluster_image_scanning",
           "name": "CVE-2017-16997 in libc",
@@ -605,8 +607,10 @@ FactoryBot.define do
             },
             "operating_system": "alpine 3.7",
             "image": "alpine:3.7",
+            "kubernetes_resource": {
               "cluster_id": "1",
               "agent_id": "46357"
+            }
           },
           "identifiers": [{
             "type": "cve",

ee/spec/finders/security/vulnerabilities_finder_spec.rb

@@ -190,14 +190,14 @@ RSpec.describe Security::VulnerabilitiesFinder do
     let_it_be(:cluster_vulnerability) { create(:vulnerability, :cluster_image_scanning, project: project) }
     let_it_be(:finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: cluster_vulnerability) }
 
-    let(:filters) { { cluster_id: [finding.location['cluster_id']] } }
+    let(:filters) { { cluster_id: [finding.location['kubernetes_resource']['cluster_id']] } }
 
     it 'only returns vulnerabilities matching the given cluster_id' do
       is_expected.to contain_exactly(cluster_vulnerability)
     end
 
     context 'when different report_type is passed' do
-      let(:filters) { { report_type: %w[dast], cluster_id: [finding.location['cluster_id']] }}
+      let(:filters) { { report_type: %w[dast], cluster_id: [finding.location['kubernetes_resource']['cluster_id']] }}
 
       it 'returns empty list' do
         is_expected.to be_empty
@@ -209,14 +209,14 @@ RSpec.describe Security::VulnerabilitiesFinder do
     let_it_be(:cluster_vulnerability) { create(:vulnerability, :cluster_image_scanning, project: project) }
     let_it_be(:finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: cluster_vulnerability) }
 
-    let(:filters) { { cluster_agent_id: [finding.location['agent_id']] } }
+    let(:filters) { { cluster_agent_id: [finding.location['kubernetes_resource']['agent_id']] } }
 
     it 'only returns vulnerabilities matching the given agent_id' do
       is_expected.to contain_exactly(cluster_vulnerability)
     end
 
     context 'when different report_type is passed' do
-      let(:filters) { { report_type: %w[dast], cluster_agent_id: [finding.location['agent_id']] }}
+      let(:filters) { { report_type: %w[dast], cluster_agent_id: [finding.location['kubernetes_resource']['agent_id']] }}
 
       it 'returns empty list' do
         is_expected.to be_empty

ee/spec/graphql/resolvers/vulnerabilities_resolver_spec.rb

@@ -214,7 +214,7 @@ RSpec.describe Resolvers::VulnerabilitiesResolver do
     context 'when cluster_id is given' do
       let_it_be(:cluster_vulnerability) { create(:vulnerability, :cluster_image_scanning, project: project) }
       let_it_be(:cluster_finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: cluster_vulnerability) }
-      let_it_be(:cluster_gid) { ::Gitlab::GlobalId.as_global_id(cluster_finding.location['cluster_id'].to_i, model_name: 'Clusters::Cluster') }
+      let_it_be(:cluster_gid) { ::Gitlab::GlobalId.as_global_id(cluster_finding.location['kubernetes_resource']['cluster_id'].to_i, model_name: 'Clusters::Cluster') }
 
       let(:params) { { cluster_id: [cluster_gid] } }
 
@@ -234,7 +234,7 @@ RSpec.describe Resolvers::VulnerabilitiesResolver do
     context 'when cluster_agent_id is given' do
       let_it_be(:cluster_vulnerability) { create(:vulnerability, :cluster_image_scanning, project: project) }
       let_it_be(:cluster_finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: cluster_vulnerability) }
-      let_it_be(:cluster_gid) { ::Gitlab::GlobalId.as_global_id(cluster_finding.location['agent_id'].to_i, model_name: 'Clusters::Cluster') }
+      let_it_be(:cluster_gid) { ::Gitlab::GlobalId.as_global_id(cluster_finding.location['kubernetes_resource']['agent_id'].to_i, model_name: 'Clusters::Agent') }
 
       let(:params) { { cluster_agent_id: [cluster_gid] } }
 

ee/spec/models/ee/vulnerability_spec.rb

@@ -604,7 +604,7 @@ RSpec.describe Vulnerability do
   describe '.with_cluster_ids' do
     let_it_be(:vulnerability) { create(:vulnerability, project: project, report_type: 'cluster_image_scanning') }
     let_it_be(:finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: vulnerability) }
-    let_it_be(:cluster_ids) { [finding.location['cluster_id']] }
+    let_it_be(:cluster_ids) { [finding.location['kubernetes_resource']['cluster_id']] }
 
     before do
       finding_with_different_cluster_id = create(
@@ -612,7 +612,7 @@ RSpec.describe Vulnerability do
         :with_cluster_image_scanning_scanning_metadata,
         vulnerability: create(:vulnerability, report_type: 'cluster_image_scanning')
       )
-      finding_with_different_cluster_id.location['cluster_id'] = '2'
+      finding_with_different_cluster_id.location['kubernetes_resource']['cluster_id'] = '2'
       finding_with_different_cluster_id.save!
 
       finding_without_cluster_id = create(
@@ -620,7 +620,7 @@ RSpec.describe Vulnerability do
         :with_cluster_image_scanning_scanning_metadata,
         vulnerability: create(:vulnerability, report_type: 'cluster_image_scanning')
       )
-      finding_without_cluster_id.location['cluster_id'] = nil
+      finding_without_cluster_id.location['kubernetes_resource']['cluster_id'] = nil
       finding_without_cluster_id.save!
     end
 
@@ -634,7 +634,7 @@ RSpec.describe Vulnerability do
   describe '.with_cluster_agent_ids' do
     let_it_be(:vulnerability) { create(:vulnerability, project: project, report_type: 'cluster_image_scanning') }
     let_it_be(:finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: vulnerability) }
-    let_it_be(:cluster_agent_ids) { [finding.location['agent_id']] }
+    let_it_be(:cluster_agent_ids) { [finding.location['kubernetes_resource']['agent_id']] }
 
     before do
       finding_with_different_agent_id = create(
@@ -642,7 +642,7 @@ RSpec.describe Vulnerability do
         :with_cluster_image_scanning_scanning_metadata,
         vulnerability: create(:vulnerability, report_type: 'cluster_image_scanning')
       )
-      finding_with_different_agent_id.location['agent_id'] = '2'
+      finding_with_different_agent_id.location['kubernetes_resource']['agent_id'] = '2'
       finding_with_different_agent_id.save!
 
       finding_without_agent_id = create(
@@ -650,7 +650,7 @@ RSpec.describe Vulnerability do
         :with_cluster_image_scanning_scanning_metadata,
         vulnerability: create(:vulnerability, report_type: 'cluster_image_scanning')
       )
-      finding_without_agent_id.location['agent_id'] = nil
+      finding_without_agent_id.location['kubernetes_resource']['agent_id'] = nil
       finding_without_agent_id.save!
     end
 

ee/spec/models/vulnerabilities/finding_spec.rb

@@ -366,7 +366,7 @@ RSpec.describe Vulnerabilities::Finding do
     describe '.by_location_cluster' do
       let_it_be(:vulnerability) { create(:vulnerability, report_type: 'cluster_image_scanning') }
       let_it_be(:finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: vulnerability) }
-      let_it_be(:cluster_ids) { [finding.location['cluster_id']] }
+      let_it_be(:cluster_ids) { [finding.location['kubernetes_resource']['cluster_id']] }
 
       before do
         finding_with_different_cluster_id = create(
@@ -374,7 +374,7 @@ RSpec.describe Vulnerabilities::Finding do
           :with_cluster_image_scanning_scanning_metadata,
           vulnerability: create(:vulnerability, report_type: 'cluster_image_scanning')
         )
-        finding_with_different_cluster_id.location['cluster_id'] = '2'
+        finding_with_different_cluster_id.location['kubernetes_resource']['cluster_id'] = '2'
         finding_with_different_cluster_id.save!
 
         create(:vulnerabilities_finding, report_type: :dast)
@@ -390,7 +390,7 @@ RSpec.describe Vulnerabilities::Finding do
     describe '.by_location_cluster_agent' do
       let_it_be(:vulnerability) { create(:vulnerability, report_type: 'cluster_image_scanning') }
       let_it_be(:finding) { create(:vulnerabilities_finding, :with_cluster_image_scanning_scanning_metadata, vulnerability: vulnerability) }
-      let_it_be(:agent_ids) { [finding.location['agent_id']] }
+      let_it_be(:agent_ids) { [finding.location['kubernetes_resource']['agent_id']] }
 
       before do
         finding_with_different_agent_id = create(
@@ -398,7 +398,7 @@ RSpec.describe Vulnerabilities::Finding do
           :with_cluster_image_scanning_scanning_metadata,
           vulnerability: create(:vulnerability, report_type: 'cluster_image_scanning')
         )
-        finding_with_different_agent_id.location['agent_id'] = '2'
+        finding_with_different_agent_id.location['kubernetes_resource']['agent_id'] = '2'
         finding_with_different_agent_id.save!
 
         create(:vulnerabilities_finding, report_type: :dast)