From fa145e30591667582b14f1d1b32776aeaa89a3ac Mon Sep 17 00:00:00 2001
From: Mehmet Emin INAC <minac@gitlab.com>
Date: Fri, 28 Feb 2025 12:51:55 +0000
Subject: [PATCH] Introduce `Vulnerabilities::ArchiveExport` model
Changelog: other
---
config/gitlab_loose_foreign_keys.yml | 7 +
config/initializers/postgres_partitioning.rb | 3 +-
db/docs/vulnerability_archive_exports.yml | 12 ++
...35_create_vulnerability_archive_exports.rb | 39 ++++++
db/schema_migrations/20250214121035 | 1 +
db/structure.sql | 39 ++++++
.../models/vulnerabilities/archive_export.rb | 96 +++++++++++++
.../vulnerabilities/archive_exports.rb | 20 +++
.../vulnerabilities/archive_export.csv | 1 +
.../vulnerabilities/archive_export_spec.rb | 128 ++++++++++++++++++
10 files changed, 345 insertions(+), 1 deletion(-)
create mode 100644 db/docs/vulnerability_archive_exports.yml
create mode 100644 db/migrate/20250214121035_create_vulnerability_archive_exports.rb
create mode 100644 db/schema_migrations/20250214121035
create mode 100644 ee/app/models/vulnerabilities/archive_export.rb
create mode 100644 ee/spec/factories/vulnerabilities/archive_exports.rb
create mode 100644 ee/spec/fixtures/vulnerabilities/archive_export.csv
create mode 100644 ee/spec/models/vulnerabilities/archive_export_spec.rb
diff --git a/config/gitlab_loose_foreign_keys.yml b/config/gitlab_loose_foreign_keys.yml
index 228a372725182..54cc3159d1109 100644
--- a/config/gitlab_loose_foreign_keys.yml
+++ b/config/gitlab_loose_foreign_keys.yml
@@ -620,6 +620,13 @@ vulnerabilities:
- table: projects
column: project_id
on_delete: async_delete
+vulnerability_archive_exports:
+ - table: projects
+ column: project_id
+ on_delete: async_delete
+ - table: users
+ column: author_id
+ on_delete: async_delete
vulnerability_archived_records:
- table: projects
column: project_id
diff --git a/config/initializers/postgres_partitioning.rb b/config/initializers/postgres_partitioning.rb
index 744167b4e8a00..9ec06212e82bf 100644
--- a/config/initializers/postgres_partitioning.rb
+++ b/config/initializers/postgres_partitioning.rb
@@ -49,7 +49,8 @@
Ci::FinishedBuildChSyncEvent,
Search::Zoekt::Task,
Ai::CodeSuggestionEvent,
- Ai::DuoChatEvent
+ Ai::DuoChatEvent,
+ Vulnerabilities::ArchiveExport
])
else
Gitlab::Database::Partitioning.register_tables(
diff --git a/db/docs/vulnerability_archive_exports.yml b/db/docs/vulnerability_archive_exports.yml
new file mode 100644
index 0000000000000..0a0735603a3a9
--- /dev/null
+++ b/db/docs/vulnerability_archive_exports.yml
@@ -0,0 +1,12 @@
+---
+table_name: vulnerability_archive_exports
+classes:
+- Vulnerabilities::ArchiveExport
+feature_categories:
+- vulnerability_management
+description: Stores the archive export information.
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/182290
+milestone: '17.10'
+gitlab_schema: gitlab_sec
+sharding_key:
+ project_id: projects
diff --git a/db/migrate/20250214121035_create_vulnerability_archive_exports.rb b/db/migrate/20250214121035_create_vulnerability_archive_exports.rb
new file mode 100644
index 0000000000000..f945aaadf1529
--- /dev/null
+++ b/db/migrate/20250214121035_create_vulnerability_archive_exports.rb
@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+class CreateVulnerabilityArchiveExports < Gitlab::Database::Migration[2.2]
+ CREATE_TABLE_OPTIONS = {
+ primary_key: %i[id partition_number],
+ options: 'PARTITION BY LIST (partition_number)'
+ }.freeze
+
+ milestone '17.10'
+
+ def up
+ create_table :vulnerability_archive_exports, **CREATE_TABLE_OPTIONS do |t| # rubocop:disable Migration/EnsureFactoryForTable -- false positive
+ t.bigserial :id, null: false
+ t.bigint :partition_number, null: false, default: 1
+ t.timestamps_with_timezone null: false
+ t.datetime_with_timezone :started_at
+ t.datetime_with_timezone :finished_at
+ t.bigint :project_id, null: false, index: true
+ t.bigint :author_id, null: false, index: true
+ t.daterange :date_range, null: false
+ t.integer :file_store, limit: 2
+ t.integer :format, limit: 2
+ t.text :file, limit: 255
+ t.text :status, limit: 8
+
+ t.index :status
+ end
+
+ connection.execute(<<~SQL)
+ CREATE TABLE IF NOT EXISTS gitlab_partitions_dynamic.vulnerability_archive_exports_1
+ PARTITION OF vulnerability_archive_exports
+ FOR VALUES IN (1);
+ SQL
+ end
+
+ def down
+ drop_table :vulnerability_archive_exports
+ end
+end
diff --git a/db/schema_migrations/20250214121035 b/db/schema_migrations/20250214121035
new file mode 100644
index 0000000000000..98f0abf79ed2b
--- /dev/null
+++ b/db/schema_migrations/20250214121035
@@ -0,0 +1 @@
+26bbbc8ad73e099557a4c2fe31b61b4674857e5045900e5308d6fb832d7cd1ba
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index 62156c0ff99fe..943169e456d8c 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -4453,6 +4453,25 @@ PARTITION BY RANGE (created_at);
COMMENT ON TABLE verification_codes IS 'JiHu-specific table';
+CREATE TABLE vulnerability_archive_exports (
+ id bigint NOT NULL,
+ partition_number bigint DEFAULT 1 NOT NULL,
+ created_at timestamp with time zone NOT NULL,
+ updated_at timestamp with time zone NOT NULL,
+ started_at timestamp with time zone,
+ finished_at timestamp with time zone,
+ project_id bigint NOT NULL,
+ author_id bigint NOT NULL,
+ date_range daterange NOT NULL,
+ file_store smallint,
+ format smallint,
+ file text,
+ status text,
+ CONSTRAINT check_3423276100 CHECK ((char_length(file) <= 255)),
+ CONSTRAINT check_aada0b0f45 CHECK ((char_length(status) <= 8))
+)
+PARTITION BY LIST (partition_number);
+
CREATE TABLE web_hook_logs (
id bigint NOT NULL,
web_hook_id bigint NOT NULL,
@@ -23204,6 +23223,15 @@ CREATE SEQUENCE vulnerabilities_id_seq
ALTER SEQUENCE vulnerabilities_id_seq OWNED BY vulnerabilities.id;
+CREATE SEQUENCE vulnerability_archive_exports_id_seq
+ START WITH 1
+ INCREMENT BY 1
+ NO MINVALUE
+ NO MAXVALUE
+ CACHE 1;
+
+ALTER SEQUENCE vulnerability_archive_exports_id_seq OWNED BY vulnerability_archive_exports.id;
+
CREATE TABLE vulnerability_archived_records (
id bigint NOT NULL,
created_at timestamp with time zone NOT NULL,
@@ -26259,6 +26287,8 @@ ALTER TABLE ONLY vs_code_settings ALTER COLUMN id SET DEFAULT nextval('vs_code_s
ALTER TABLE ONLY vulnerabilities ALTER COLUMN id SET DEFAULT nextval('vulnerabilities_id_seq'::regclass);
+ALTER TABLE ONLY vulnerability_archive_exports ALTER COLUMN id SET DEFAULT nextval('vulnerability_archive_exports_id_seq'::regclass);
+
ALTER TABLE ONLY vulnerability_archived_records ALTER COLUMN id SET DEFAULT nextval('vulnerability_archived_records_id_seq'::regclass);
ALTER TABLE ONLY vulnerability_archives ALTER COLUMN id SET DEFAULT nextval('vulnerability_archives_id_seq'::regclass);
@@ -29276,6 +29306,9 @@ ALTER TABLE ONLY vs_code_settings
ALTER TABLE ONLY vulnerabilities
ADD CONSTRAINT vulnerabilities_pkey PRIMARY KEY (id);
+ALTER TABLE ONLY vulnerability_archive_exports
+ ADD CONSTRAINT vulnerability_archive_exports_pkey PRIMARY KEY (id, partition_number);
+
ALTER TABLE ONLY vulnerability_archived_records
ADD CONSTRAINT vulnerability_archived_records_pkey PRIMARY KEY (id);
@@ -35594,6 +35627,12 @@ CREATE INDEX index_vulnerabilities_project_id_and_id_on_default_branch ON vulner
CREATE INDEX index_vulnerabilities_project_id_state_severity_default_branch ON vulnerabilities USING btree (project_id, state, severity, present_on_default_branch);
+CREATE INDEX index_vulnerability_archive_exports_on_author_id ON ONLY vulnerability_archive_exports USING btree (author_id);
+
+CREATE INDEX index_vulnerability_archive_exports_on_project_id ON ONLY vulnerability_archive_exports USING btree (project_id);
+
+CREATE INDEX index_vulnerability_archive_exports_on_status ON ONLY vulnerability_archive_exports USING btree (status);
+
CREATE INDEX index_vulnerability_archived_records_on_archive_id ON vulnerability_archived_records USING btree (archive_id);
CREATE INDEX index_vulnerability_archived_records_on_project_id ON vulnerability_archived_records USING btree (project_id);
diff --git a/ee/app/models/vulnerabilities/archive_export.rb b/ee/app/models/vulnerabilities/archive_export.rb
new file mode 100644
index 0000000000000..748fc56e59f2e
--- /dev/null
+++ b/ee/app/models/vulnerabilities/archive_export.rb
@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module Vulnerabilities
+ class ArchiveExport < Gitlab::Database::SecApplicationRecord
+ include FileStoreMounter
+ include PartitionedTable
+
+ RETENTION_PERIOD = 1.month
+
+ mount_file_store_uploader AttachmentUploader
+
+ self.table_name = 'vulnerability_archive_exports'
+ self.primary_key = :id
+
+ attr_readonly :partition_number
+
+ partitioned_by :partition_number,
+ strategy: :sliding_list,
+ next_partition_if: ->(partition) { requires_new_partition?(partition.value) },
+ detach_partition_if: ->(partition) { detach_partition?(partition.value) }
+
+ belongs_to :project, optional: false
+ belongs_to :author, class_name: 'User', optional: false
+
+ enum format: { csv: 0 }
+
+ validates :date_range, presence: true
+ validates :status, presence: true
+ validates :format, presence: true
+ validates :file, presence: true, if: :finished?
+
+ state_machine :status, initial: :created do
+ state :created
+ state :running
+ state :finished
+ state :failed
+ state :purged
+
+ event :start do
+ transition created: :running
+ end
+
+ event :finish do
+ transition running: :finished
+ end
+
+ event :failed do
+ transition [:created, :running] => :failed
+ end
+
+ event :reset_state do
+ transition running: :created
+ end
+
+ before_transition created: :running do |export|
+ export.started_at = Time.current
+ end
+
+ before_transition running: :finished do |export|
+ export.finished_at = Time.current
+ end
+
+ before_transition running: :created do |export|
+ export.started_at = nil
+ end
+ end
+
+ class << self
+ def requires_new_partition?(partition_number)
+ first_record = first_record_in(partition_number)
+
+ return unless first_record
+
+ first_record.created_at < RETENTION_PERIOD.ago
+ end
+
+ def detach_partition?(partition_number)
+ where(partition_number: partition_number).where.not(status: :purged).none?
+ end
+
+ private
+
+ def first_record_in(partition_number)
+ where(partition_number: partition_number).first
+ end
+ end
+
+ def uploads_sharding_key
+ { project_id: project_id }
+ end
+
+ def retrieve_upload(_identifier, paths)
+ Upload.find_by(model: self, path: paths)
+ end
+ end
+end
diff --git a/ee/spec/factories/vulnerabilities/archive_exports.rb b/ee/spec/factories/vulnerabilities/archive_exports.rb
new file mode 100644
index 0000000000000..2139057aba580
--- /dev/null
+++ b/ee/spec/factories/vulnerabilities/archive_exports.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+FactoryBot.define do
+ factory :vulnerability_archive_export, class: 'Vulnerabilities::ArchiveExport' do
+ project
+ author factory: :user
+
+ date_range { (5.days.ago..Time.zone.today) }
+ format { :csv }
+
+ trait :with_csv_file do
+ file { fixture_file_upload('ee/spec/fixtures/vulnerabilities/archive_export.csv') }
+ end
+
+ trait :running do
+ status { :running }
+ started_at { 1.minute.ago }
+ end
+ end
+end
diff --git a/ee/spec/fixtures/vulnerabilities/archive_export.csv b/ee/spec/fixtures/vulnerabilities/archive_export.csv
new file mode 100644
index 0000000000000..f38bbbd153c8c
--- /dev/null
+++ b/ee/spec/fixtures/vulnerabilities/archive_export.csv
@@ -0,0 +1 @@
+Scanner Type,Scanner Name,Status,Vulnerability,Details,Severity,CVE,CWE,Other Identifiers,Detected At,Location,Activity,Comments,CVSS Vectors,Dismissal Reason
diff --git a/ee/spec/models/vulnerabilities/archive_export_spec.rb b/ee/spec/models/vulnerabilities/archive_export_spec.rb
new file mode 100644
index 0000000000000..69f3c14693e1c
--- /dev/null
+++ b/ee/spec/models/vulnerabilities/archive_export_spec.rb
@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Vulnerabilities::ArchiveExport, feature_category: :vulnerability_management do
+ describe 'associations' do
+ it { is_expected.to belong_to(:project).required }
+ it { is_expected.to belong_to(:author).class_name('User').required }
+ end
+
+ describe 'validations' do
+ it { is_expected.to validate_presence_of(:date_range) }
+ it { is_expected.to validate_presence_of(:status) }
+ it { is_expected.to validate_presence_of(:format) }
+ it { is_expected.not_to validate_presence_of(:file) }
+
+ context 'when the status is `finished`' do
+ subject { build(:vulnerability_archive_export, status: :finished) }
+
+ it { is_expected.to validate_presence_of(:file) }
+ end
+ end
+
+ describe 'state machine' do
+ describe '#start' do
+ let(:export) { create(:vulnerability_archive_export) }
+
+ subject(:start) { export.start! }
+
+ it 'changes the status of the export' do
+ expect { start }.to change { export.reload.status }
+ end
+
+ it 'sets the `started_at` attribute', :freeze_time do
+ expect { start }.to change { export.reload.started_at }.from(nil).to(Time.current)
+ end
+ end
+
+ describe '#finish' do
+ let(:export) { create(:vulnerability_archive_export, :with_csv_file, :running) }
+
+ subject(:finish) { export.finish! }
+
+ it 'changes the status of the export' do
+ expect { finish }.to change { export.reload.status }.to('finished')
+ end
+
+ it 'sets the `finished_at` attribute', :freeze_time do
+ expect { finish }.to change { export.reload.finished_at }.from(nil).to(Time.current)
+ end
+ end
+
+ describe '#failed' do
+ let(:export) { create(:vulnerability_archive_export, :running) }
+
+ subject(:failed) { export.failed! }
+
+ it 'changes the status of the export' do
+ expect { failed }.to change { export.reload.status }.to('failed')
+ end
+ end
+
+ describe '#reset_state' do
+ let(:export) { create(:vulnerability_archive_export, :running) }
+
+ subject(:reset_state) { export.reset_state! }
+
+ it 'changes the status of the export' do
+ expect { reset_state }.to change { export.reload.status }.to('created')
+ end
+
+ it 'resets the `started_at` attribute' do
+ expect { reset_state }.to change { export.reload.started_at }.to(nil)
+ end
+ end
+ end
+
+ describe 'partition helpers' do
+ let(:partitioning_strategy) { described_class.partitioning_strategy }
+ let(:active_partition) { partitioning_strategy.active_partition }
+
+ describe '.partitioning_strategy#detach_partition_if' do
+ subject { partitioning_strategy.detach_partition_if.call(active_partition) }
+
+ before do
+ create(:vulnerability_archive_export, status: status_of_existing_record)
+ end
+
+ context 'when there is a non-purged record' do
+ let(:status_of_existing_record) { :created }
+
+ it { is_expected.to be_falsey }
+ end
+
+ context 'when there is no non-purged record' do
+ let(:status_of_existing_record) { :purged }
+
+ it { is_expected.to be_truthy }
+ end
+ end
+
+ describe '.partitioning_strategy#next_partition_if' do
+ subject { partitioning_strategy.next_partition_if.call(active_partition) }
+
+ context 'when there is no record in the partition' do
+ it { is_expected.to be_falsey }
+ end
+
+ context 'when there is at least one record in the partition', :freeze_time do
+ before do
+ create(:vulnerability_archive_export, created_at: record_created_at)
+ end
+
+ context 'when the first record in the partition is not older than 1 month' do
+ let(:record_created_at) { 1.week.ago }
+
+ it { is_expected.to be_falsey }
+ end
+
+ context 'when the first record in the partition is older than 1 month' do
+ let(:record_created_at) { 2.months.ago }
+
+ it { is_expected.to be_truthy }
+ end
+ end
+ end
+ end
+end
--
GitLab