gitlab-foss/db/migrate/20250313171706_partition_vulnerability_archives.rb

# frozen_string_literal: true

class PartitionVulnerabilityArchives < Gitlab::Database::Migration[2.2]
  INDEX_NAME_ON_FK = 'index_vulnerability_archived_records_on_archive_id_and_date'
  PARTITIONED_TABLE_OPTIONS = {
    primary_key: %i[id date],
    options: 'PARTITION BY RANGE (date)'
  }.freeze

  milestone '17.11'

  def up
    remove_foreign_key :vulnerability_archived_records, column: :archive_id

    drop_table :vulnerability_archives # rubocop:disable Migration/DropTable -- We are temporarily dropping the table

    create_table :vulnerability_archives, **PARTITIONED_TABLE_OPTIONS do |t| # rubocop:disable Migration/EnsureFactoryForTable -- false positive
      t.timestamps_with_timezone null: false

      t.bigserial :id, null: false
      t.bigint :project_id, null: false
      t.integer :archived_records_count, null: false, default: 0
      t.date :date, null: false

      t.index %i[project_id date], unique: true

      t.check_constraint 'archived_records_count >= 0'
    end

    add_index :vulnerability_archived_records, %i[archive_id date], name: INDEX_NAME_ON_FK

    connection.execute(<<~SQL)
      ALTER TABLE vulnerability_archived_records
      ADD CONSTRAINT fk_rails_601e008d4b FOREIGN KEY (archive_id, date) REFERENCES vulnerability_archives(id, date) ON DELETE CASCADE;
    SQL
  end

  def down
    # Drop the constraint from partitioned table.
    connection.execute(<<~SQL)
      ALTER TABLE vulnerability_archived_records DROP CONSTRAINT fk_rails_601e008d4b;
    SQL

    # Drop the constraint from detached partitions.
    foreign_keys = connection.execute(<<~SQL)
      SELECT
        table_schema,
        table_name
      FROM
        information_schema.table_constraints
      WHERE
        table_name ILIKE 'vulnerability_archived_records%'
        AND constraint_name = 'fk_rails_601e008d4b';
    SQL

    foreign_keys.values.each do |schema, table_name| # rubocop:disable Style/HashEachMethods -- This is a Hash
      connection.execute(<<~SQL)
        ALTER TABLE connection.quote_table_name(#{schema}.#{table_name}) DROP CONSTRAINT fk_rails_601e008d4b;
      SQL
    end

    drop_table :vulnerability_archives

    create_table :vulnerability_archives do |t| # rubocop:disable Migration/EnsureFactoryForTable -- false positive
      t.timestamps_with_timezone null: false

      t.bigint :project_id, null: false
      t.integer :archived_records_count, null: false, default: 0
      t.date :date, null: false

      t.index %i[project_id date], unique: true

      t.check_constraint 'archived_records_count >= 0'
    end

    remove_index :vulnerability_archived_records, name: INDEX_NAME_ON_FK # rubocop:disable  Migration/RemoveIndex -- Table is empty

    add_foreign_key :vulnerability_archived_records, :vulnerability_archives, column: :archive_id, on_delete: :cascade # rubocop:disable Migration/AddConcurrentForeignKey -- Table is empty
  end
end