Merge branch 'ag-geo-add-vulnerability-export-replication' into 'master'

Geo: Add vulnerability export replication

See merge request gitlab-org/gitlab!36620

config/initializers_before_autoloader/000_inflections.rb

@@ -26,6 +26,7 @@ ActiveSupport::Inflector.inflections do |inflect|
     project_statistics
     system_note_metadata
     vulnerabilities_feedback
+    vulnerability_export_registry
     vulnerability_feedback
   )
   inflect.acronym 'EE'

db/migrate/20200720154007_create_vulnerabilities_export_verification_status.rb

+# frozen_string_literal: true
+
+class CreateVulnerabilitiesExportVerificationStatus < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  def up
+    unless table_exists?(:vulnerability_export_verification_status)
+      with_lock_retries do
+        create_table :vulnerability_export_verification_status, id: false do |t|
+          t.references :vulnerability_export,
+            primary_key: true,
+            null: false,
+            foreign_key: { on_delete: :cascade },
+            index:
+              { name: 'index_vulnerability_export_verification_status_on_export_id' }
+          t.datetime_with_timezone :verification_retry_at
+          t.datetime_with_timezone :verified_at
+          t.integer :verification_retry_count, limit: 2
+          t.binary :verification_checksum, using: 'verification_checksum::bytea'
+          t.text :verification_failure
+
+          t.index :verification_failure, where: "(verification_failure IS NOT NULL)", name: "vulnerability_exports_verification_failure_partial"
+          t.index :verification_checksum, where: "(verification_checksum IS NOT NULL)", name: "vulnerability_exports_verification_checksum_partial"
+        end
+      end
+    end
+
+    add_text_limit :vulnerability_export_verification_status, :verification_failure, 255
+  end
+
+  def down
+    with_lock_retries do
+      drop_table :vulnerability_export_verification_status
+    end
+  end
+end

db/schema_migrations/20200720154007

+8da79a65932d06a4a4cf7b601ecac742ae7b785a1c6799019b1ee758f844cacf
\ No newline at end of file

db/structure.sql

@@ -16167,6 +16167,25 @@ CREATE SEQUENCE public.vulnerabilities_id_seq
 
 ALTER SEQUENCE public.vulnerabilities_id_seq OWNED BY public.vulnerabilities.id;
 
+CREATE TABLE public.vulnerability_export_verification_status (
+    vulnerability_export_id bigint NOT NULL,
+    verification_retry_at timestamp with time zone,
+    verified_at timestamp with time zone,
+    verification_retry_count smallint,
+    verification_checksum bytea,
+    verification_failure text,
+    CONSTRAINT check_48fdf48546 CHECK ((char_length(verification_failure) <= 255))
+);
+
+CREATE SEQUENCE public.vulnerability_export_verification_s_vulnerability_export_id_seq
+    START WITH 1
+    INCREMENT BY 1
+    NO MINVALUE
+    NO MAXVALUE
+    CACHE 1;
+
+ALTER SEQUENCE public.vulnerability_export_verification_s_vulnerability_export_id_seq OWNED BY public.vulnerability_export_verification_status.vulnerability_export_id;
+
 CREATE TABLE public.vulnerability_exports (
     id bigint NOT NULL,
     created_at timestamp with time zone NOT NULL,
@@ -17235,6 +17254,8 @@ ALTER TABLE ONLY public.users_statistics ALTER COLUMN id SET DEFAULT nextval('pu
 
 ALTER TABLE ONLY public.vulnerabilities ALTER COLUMN id SET DEFAULT nextval('public.vulnerabilities_id_seq'::regclass);
 
+ALTER TABLE ONLY public.vulnerability_export_verification_status ALTER COLUMN vulnerability_export_id SET DEFAULT nextval('public.vulnerability_export_verification_s_vulnerability_export_id_seq'::regclass);
+
 ALTER TABLE ONLY public.vulnerability_exports ALTER COLUMN id SET DEFAULT nextval('public.vulnerability_exports_id_seq'::regclass);
 
 ALTER TABLE ONLY public.vulnerability_feedback ALTER COLUMN id SET DEFAULT nextval('public.vulnerability_feedback_id_seq'::regclass);
@@ -18517,6 +18538,9 @@ ALTER TABLE ONLY public.users_statistics
 ALTER TABLE ONLY public.vulnerabilities
     ADD CONSTRAINT vulnerabilities_pkey PRIMARY KEY (id);
 
+ALTER TABLE ONLY public.vulnerability_export_verification_status
+    ADD CONSTRAINT vulnerability_export_verification_status_pkey PRIMARY KEY (vulnerability_export_id);
+
 ALTER TABLE ONLY public.vulnerability_exports
     ADD CONSTRAINT vulnerability_exports_pkey PRIMARY KEY (id);
 
@@ -20761,6 +20785,8 @@ CREATE INDEX index_vulnerabilities_on_start_date_sourcing_milestone_id ON public
 
 CREATE INDEX index_vulnerabilities_on_updated_by_id ON public.vulnerabilities USING btree (updated_by_id);
 
+CREATE INDEX index_vulnerability_export_verification_status_on_export_id ON public.vulnerability_export_verification_status USING btree (vulnerability_export_id);
+
 CREATE INDEX index_vulnerability_exports_on_author_id ON public.vulnerability_exports USING btree (author_id);
 
 CREATE INDEX index_vulnerability_exports_on_file_store ON public.vulnerability_exports USING btree (file_store);
@@ -20905,6 +20931,10 @@ CREATE UNIQUE INDEX unique_merge_request_metrics_by_merge_request_id ON public.m
 
 CREATE UNIQUE INDEX users_security_dashboard_projects_unique_index ON public.users_security_dashboard_projects USING btree (project_id, user_id);
 
+CREATE INDEX vulnerability_exports_verification_checksum_partial ON public.vulnerability_export_verification_status USING btree (verification_checksum) WHERE (verification_checksum IS NOT NULL);
+
+CREATE INDEX vulnerability_exports_verification_failure_partial ON public.vulnerability_export_verification_status USING btree (verification_failure) WHERE (verification_failure IS NOT NULL);
+
 CREATE UNIQUE INDEX vulnerability_feedback_unique_idx ON public.vulnerability_feedback USING btree (project_id, category, feedback_type, project_fingerprint);
 
 CREATE UNIQUE INDEX vulnerability_occurrence_pipelines_on_unique_keys ON public.vulnerability_occurrence_pipelines USING btree (occurrence_id, pipeline_id);
@@ -21782,6 +21812,9 @@ ALTER TABLE ONLY public.approval_merge_request_rules
 ALTER TABLE ONLY public.namespace_statistics
     ADD CONSTRAINT fk_rails_0062050394 FOREIGN KEY (namespace_id) REFERENCES public.namespaces(id) ON DELETE CASCADE;
 
+ALTER TABLE ONLY public.vulnerability_export_verification_status
+    ADD CONSTRAINT fk_rails_00a22ee64f FOREIGN KEY (vulnerability_export_id) REFERENCES public.vulnerability_exports(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY public.clusters_applications_elastic_stacks
     ADD CONSTRAINT fk_rails_026f219f46 FOREIGN KEY (cluster_id) REFERENCES public.clusters(id) ON DELETE CASCADE;
 

ee/app/models/concerns/geo/blob_replicator_strategy.rb

@@ -13,6 +13,8 @@ module Geo
     end
 
     def handle_after_create_commit
+      return unless self.class.enabled?
+
       publish(:created, **created_params)
       schedule_checksum_calculation if needs_checksum?
     end
@@ -25,6 +27,8 @@ module Geo
     end
 
     def handle_after_destroy
+      return unless self.class.enabled?
+
       publish(:deleted, **deleted_params)
     end
 
@@ -114,7 +118,6 @@ module Geo
     end
 
     def needs_checksum?
-      return false unless self.class.enabled?
       return true unless model_record.respond_to?(:needs_checksum?)
 
       model_record.needs_checksum?

ee/app/models/geo/vulnerability_export_registry.rb

+# frozen_string_literal: true
+
+class Geo::VulnerabilityExportRegistry < Geo::BaseRegistry
+  include ::Geo::ReplicableRegistry
+
+  MODEL_CLASS = ::Vulnerabilities::Export
+  MODEL_FOREIGN_KEY = :vulnerability_export_id
+
+  belongs_to :vulnerability_export, class_name: 'Vulnerabilities::Export', foreign_key: :vulnerability_export_id
+end

ee/app/models/vulnerabilities/export.rb

@@ -2,12 +2,25 @@
 
 module Vulnerabilities
   class Export < ApplicationRecord
+    include ::Gitlab::Geo::ReplicableModel
+
     self.table_name = "vulnerability_exports"
 
+    with_replicator Geo::VulnerabilityExportReplicator
+
     belongs_to :project
     belongs_to :group
     belongs_to :author, optional: false, class_name: 'User'
 
+    has_one :vulnerability_export_verification_status, class_name: 'Vulnerabilities::ExportVerificationStatus', inverse_of: :vulnerability_export, foreign_key: :vulnerability_export_id
+
+    delegate :verification_retry_at, :verification_retry_at=,
+             :verified_at, :verified_at=,
+             :verification_checksum, :verification_checksum=,
+             :verification_failure, :verification_failure=,
+             :verification_retry_count, :verification_retry_count=,
+             to: :vulnerability_export_verification_status
+
     mount_uploader :file, AttachmentUploader
 
     after_save :update_file_store, if: :saved_change_to_file?
@@ -83,6 +96,18 @@ module Vulnerabilities
       self.update_column(:file_store, file.object_store)
     end
 
+    def vulnerability_export_verification_status
+      super.presence || build_vulnerability_export_verification_status
+    end
+
+    def local?
+      file_store == ObjectStorage::Store::LOCAL
+    end
+
+    def self.replicables_for_geo_node
+      self.all
+    end
+
     private
 
     def make_project_level_export(project)

ee/app/models/vulnerabilities/export_verification_status.rb

+# frozen_string_literal: true
+
+module Vulnerabilities
+  class ExportVerificationStatus < ApplicationRecord
+    self.primary_key = :vulnerability_export_id
+
+    self.table_name = 'vulnerability_export_verification_status'
+
+    belongs_to :vulnerability_export, class_name: 'Vulnerabilities::Export', inverse_of: :vulnerability_export_verification_status
+  end
+end

ee/app/replicators/geo/vulnerability_export_replicator.rb

+# frozen_string_literal: true
+
+module Geo
+  class VulnerabilityExportReplicator < Gitlab::Geo::Replicator
+    include ::Geo::BlobReplicatorStrategy
+
+    def self.model
+      ::Vulnerabilities::Export
+    end
+
+    def carrierwave_uploader
+      model_record.file
+    end
+
+    def self.replication_enabled_by_default?
+      false
+    end
+  end
+end

ee/app/workers/geo/secondary/registry_consistency_worker.rb

@@ -21,7 +21,8 @@ module Geo
         Geo::LfsObjectRegistry,
         Geo::PackageFileRegistry,
         Geo::ProjectRegistry,
-        Geo::UploadRegistry
+        Geo::UploadRegistry,
+        Geo::VulnerabilityExportRegistry
       ].freeze
 
       BATCH_SIZE = 10000

ee/changelogs/unreleased/ag-geo-add-vulnerability-export-replication.yml

+---
+title: 'Geo: Add migrations for registry and states tables for vulnerability export replication'
+merge_request: 36620
+author:
+type: added

ee/db/geo/migrate/20200710194046_create_vulnerability_export_registry.rb

+# frozen_string_literal: true
+
+class CreateVulnerabilityExportRegistry < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  def up
+    unless table_exists?(:vulnerability_export_registry)
+      create_table :vulnerability_export_registry, id: :bigserial, force: :cascade do |t|
+        t.datetime_with_timezone :retry_at
+        t.datetime_with_timezone :last_synced_at
+        t.datetime_with_timezone :created_at, null: false
+        t.bigint :vulnerability_export_id, null: false
+        t.integer :state, default: 0, null: false, limit: 2
+        t.integer :retry_count, default: 0, limit: 2
+        t.text :last_sync_failure
+
+        t.index :vulnerability_export_id, name: :index_vulnerability_export_registry_on_vulnerability_export_id, unique: true
+        t.index :retry_at
+        t.index :state
+      end
+    end
+
+    add_text_limit :vulnerability_export_registry, :last_sync_failure, 255
+  end
+
+  def down
+    drop_table :vulnerability_export_registry
+  end
+end

ee/db/geo/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 2020_07_07_011052) do
+ActiveRecord::Schema.define(version: 2020_07_10_194046) do
 
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
@@ -168,4 +168,17 @@ ActiveRecord::Schema.define(version: 2020_07_07_011052) do
     t.index ["wiki_verification_checksum_sha"], name: "idx_project_registry_on_wiki_checksum_sha_partial", where: "(wiki_verification_checksum_sha IS NULL)"
   end
 
+  create_table "vulnerability_export_registry", force: :cascade do |t|
+    t.datetime_with_timezone "retry_at"
+    t.datetime_with_timezone "last_synced_at"
+    t.datetime_with_timezone "created_at", null: false
+    t.bigint "vulnerability_export_id", null: false
+    t.integer "state", limit: 2, default: 0, null: false
+    t.integer "retry_count", limit: 2, default: 0
+    t.text "last_sync_failure"
+    t.index ["retry_at"], name: "index_vulnerability_export_registry_on_retry_at"
+    t.index ["state"], name: "index_vulnerability_export_registry_on_state"
+    t.index ["vulnerability_export_id"], name: "index_vulnerability_export_registry_on_vulnerability_export_id", unique: true
+  end
+
 end

ee/lib/gitlab/geo.rb

@@ -162,7 +162,8 @@ module Gitlab
     # solutions can be found at
     # https://gitlab.com/gitlab-org/gitlab/-/issues/227693
     def self.replicator_classes
-      classes = [::Geo::PackageFileReplicator]
+      classes = [::Geo::PackageFileReplicator,
+                 ::Geo::VulnerabilityExportReplicator]
 
       classes.select(&:enabled?)
     end

ee/lib/gitlab/geo/replicator.rb

@@ -213,8 +213,6 @@ module Gitlab
       # @param [Symbol] event_name
       # @param [Hash] event_data
       def publish(event_name, **event_data)
-        return unless self.class.enabled?
-
         raise ArgumentError, "Unsupported event: '#{event_name}'" unless self.class.event_supported?(event_name)
 
         create_event_with(

ee/spec/factories/geo/vulnerability_export_registry.rb

+# frozen_string_literal: true
+
+FactoryBot.define do
+  factory :geo_vulnerability_export_registry, class: 'Geo::VulnerabilityExportRegistry' do
+    association :vulnerability_export, factory: :vulnerability_export
+
+    state { Geo::VulnerabilityExportRegistry.state_value(:pending) }
+
+    trait :synced do
+      state { Geo::VulnerabilityExportRegistry.state_value(:synced) }
+      last_synced_at { 5.days.ago }
+    end
+
+    trait :failed do
+      state { Geo::VulnerabilityExportRegistry.state_value(:failed) }
+      last_synced_at { 1.day.ago }
+      retry_count { 2 }
+      last_sync_failure { 'Random error' }
+    end
+
+    trait :started do
+      state { Geo::VulnerabilityExportRegistry.state_value(:started) }
+      last_synced_at { 1.day.ago }
+      retry_count { 0 }
+    end
+  end
+end

ee/spec/lib/gitlab/geo/geo_node_status_check_spec.rb

@@ -13,6 +13,8 @@ RSpec.describe Gitlab::Geo::GeoNodeStatusCheck do
   describe '#replication_verification_complete?' do
     before do
       allow(Gitlab.config.geo.registry_replication).to receive(:enabled).and_return(true)
+
+      stub_feature_flags(geo_vulnerability_export_replication: false)
     end
 
     it 'prints messages for all verification checks' do

ee/spec/lib/gitlab/geo/replicator_spec.rb

@@ -106,14 +106,6 @@ RSpec.describe Gitlab::Geo::Replicator do
     describe '#publish' do
       subject { Geo::DummyReplicator.new }
 
-      context 'when replication is disabled' do
-        before do
-          stub_feature_flags(geo_dummy_replication: false)
-        end
-
-        it_behaves_like 'does not replicate'
-      end
-
       context 'when publishing a supported events with required params' do
         it 'creates event with associated event log record' do
           stub_current_geo_node(primary_node)

ee/spec/lib/gitlab/geo_spec.rb

@@ -346,11 +346,11 @@ RSpec.describe Gitlab::Geo, :geo, :request_store do
 
     context 'when replication is disabled' do
       before do
-        stub_feature_flags(geo_package_file_replication: false)
+        stub_feature_flags(geo_vulnerability_export_replication: false)
       end
 
-      it 'does not return the replicator class' do
-        expect(described_class.replicator_classes).not_to include(Geo::PackageFileReplicator)
+      it 'does not return the corresponding replicator class' do
+        expect(described_class.replicator_classes).not_to include(Geo::VulnerabilityExportReplicator)
       end
     end
   end

ee/spec/models/geo/vulnerability_export_registry_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Geo::VulnerabilityExportRegistry, :geo, type: :model do
+  let_it_be(:registry) { create(:geo_vulnerability_export_registry) }
+
+  specify 'factory is valid' do
+    expect(registry).to be_valid
+  end
+
+  include_examples 'a Geo framework registry'
+end

ee/spec/replicators/geo/vulnerability_export_replicator_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Geo::VulnerabilityExportReplicator do
+  let(:model_record) { build(:vulnerability_export, :with_csv_file) }
+
+  it_behaves_like 'a blob replicator'
+end

ee/spec/services/geo/registry_consistency_service_spec.rb

@@ -14,6 +14,7 @@ RSpec.describe Geo::RegistryConsistencyService, :geo, :use_clean_rails_memory_st
   def model_class_factory_name(registry_class)
     return :project_with_design if registry_class == ::Geo::DesignRegistry
     return :package_file_with_file if registry_class == ::Geo::PackageFileRegistry
+    return :vulnerability_export if registry_class == ::Geo::VulnerabilityExportRegistry
 
     registry_class::MODEL_CLASS.underscore.tr('/', '_').to_sym
   end

ee/spec/support/shared_examples/models/concerns/blob_replicator_strategy_shared_examples.rb

@@ -14,14 +14,6 @@ RSpec.shared_examples 'a blob replicator' do
 
   subject(:replicator) { model_record.replicator }
 
-  shared_examples 'does not schedule the checksum calculation' do
-    it do
-      expect(Geo::BlobVerificationPrimaryWorker).not_to receive(:perform_async)
-
-      replicator.handle_after_create_commit
-    end
-  end
-
   before do
     stub_current_geo_node(primary)
   end
@@ -50,7 +42,17 @@ RSpec.shared_examples 'a blob replicator' do
         stub_feature_flags("geo_#{replicator.replicable_name}_replication": false)
       end
 
-      it_behaves_like 'does not schedule the checksum calculation'
+      it 'does not schedule the checksum calculation' do
+        expect(Geo::BlobVerificationPrimaryWorker).not_to receive(:perform_async)
+
+        replicator.handle_after_create_commit
+      end
+
+      it 'does not publish' do
+        expect(replicator).not_to receive(:publish)
+
+        replicator.handle_after_create_commit
+      end
     end
   end
 
@@ -63,6 +65,18 @@ RSpec.shared_examples 'a blob replicator' do
       expect(::Geo::Event.last.attributes).to include(
         "replicable_name" => replicator.replicable_name, "event_name" => "deleted", "payload" => { "model_record_id" => replicator.model_record.id, "blob_path" => replicator.blob_path })
     end
+
+    context 'when replication feature flag is disabled' do
+      before do
+        stub_feature_flags("geo_#{replicator.replicable_name}_replication": false)
+      end
+
+      it 'does not publish' do
+        expect(replicator).not_to receive(:publish)
+
+        replicator.handle_after_create_commit
+      end
+    end
   end
 
   describe '#calculate_checksum!' do

ee/spec/tasks/geo_rake_spec.rb

@@ -341,6 +341,7 @@ RSpec.describe 'geo rake tasks', :geo do
       before do
         stub_licensed_features(geo: true)
         stub_current_geo_node(current_node)
+        stub_feature_flags(geo_vulnerability_export_replication: false)
 
         allow(GeoNodeStatus).to receive(:current_node_status).and_return(geo_node_status)
         allow(Gitlab.config.geo.registry_replication).to receive(:enabled).and_return(true)

ee/spec/tasks/gitlab/geo_rake_spec.rb

@@ -8,6 +8,7 @@ RSpec.describe 'gitlab:geo rake tasks', :geo do
   before do
     Rake.application.rake_require 'tasks/gitlab/geo'
     stub_licensed_features(geo: true)
+    stub_feature_flags(geo_vulnerability_export_replication: false)
   end
 
   describe 'gitlab:geo:check_replication_verification_status' do

ee/spec/workers/geo/secondary/registry_consistency_worker_spec.rb

@@ -83,6 +83,7 @@ RSpec.describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
       create(:design, project: project)
       upload = create(:upload)
       package_file = create(:conan_package_file, :conan_package)
+      vulnerability_export = create(:vulnerability_export, :with_csv_file)
 
       expect(Geo::LfsObjectRegistry.where(lfs_object_id: lfs_object.id).count).to eq(0)
       expect(Geo::JobArtifactRegistry.where(artifact_id: job_artifact.id).count).to eq(0)
@@ -90,6 +91,7 @@ RSpec.describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
       expect(Geo::DesignRegistry.where(project_id: project.id).count).to eq(0)
       expect(Geo::UploadRegistry.where(file_id: upload.id).count).to eq(0)
       expect(Geo::PackageFileRegistry.where(package_file_id: package_file.id).count).to eq(0)
+      expect(Geo::VulnerabilityExportRegistry.where(vulnerability_export_id: vulnerability_export.id).count).to eq(0)
 
       subject.perform
 
@@ -99,6 +101,7 @@ RSpec.describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
       expect(Geo::DesignRegistry.where(project_id: project.id).count).to eq(1)
       expect(Geo::UploadRegistry.where(file_id: upload.id).count).to eq(1)
       expect(Geo::PackageFileRegistry.where(package_file_id: package_file.id).count).to eq(1)
+      expect(Geo::VulnerabilityExportRegistry.where(vulnerability_export_id: vulnerability_export.id).count).to eq(1)
     end
 
     context 'when geo_design_registry_ssot_sync is disabled' do
@@ -116,6 +119,7 @@ RSpec.describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
         allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::PackageFileRegistry, batch_size: batch_size).and_call_original
         allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::UploadRegistry, batch_size: batch_size).and_call_original
         allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::ProjectRegistry, batch_size: batch_size).and_call_original
+        allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::VulnerabilityExportRegistry, batch_size: batch_size).and_call_original
 
         expect(Geo::RegistryConsistencyService).not_to receive(:new).with(Geo::DesignRegistry, batch_size: batch_size)