Bigint Migration for 'events' Table (Step 3)

Author: philippthun

db/migrations/20250603103400_bigint_migration_events_step3a.rb

@@ -0,0 +1,21 @@
+require 'database/bigint_migration'
+
+Sequel.migration do
+  up do
+    if database_type == :postgres && !VCAP::BigintMigration.migration_completed?(self, :events) && !VCAP::BigintMigration.migration_skipped?(self, :events)
+      begin
+        VCAP::BigintMigration.add_check_constraint(self, :events)
+      rescue Sequel::CheckConstraintViolation
+        raise "Failed to add check constraint on 'events' table!\n" \
+              "There are rows where 'id_bigint' does not match 'id', thus step 3 of the bigint migration cannot be executed.\n" \
+              "Consider running rake task 'db:bigint_backfill[events]'."
+      end
+    end
+  end
+
+  down do
+    if database_type == :postgres && !VCAP::BigintMigration.migration_completed?(self, :events) && !VCAP::BigintMigration.migration_skipped?(self, :events)
+      VCAP::BigintMigration.drop_check_constraint(self, :events)
+    end
+  end
+end

db/migrations/20250603103500_bigint_migration_events_step3b.rb

@@ -0,0 +1,24 @@
+require 'database/bigint_migration'
+
+Sequel.migration do
+  up do
+    if database_type == :postgres && !VCAP::BigintMigration.migration_completed?(self, :events) && !VCAP::BigintMigration.migration_skipped?(self, :events)
+      transaction do
+        VCAP::BigintMigration.drop_check_constraint(self, :events)
+        VCAP::BigintMigration.drop_trigger_function(self, :events)
+        VCAP::BigintMigration.drop_pk_id_column(self, :events)
+        # ...
+      end
+    end
+  end
+
+  down do
+    if database_type == :postgres && !VCAP::BigintMigration.migration_completed?(self, :events) && !VCAP::BigintMigration.migration_skipped?(self, :events)
+      VCAP::BigintMigration.add_pk_id_column(self, :events)
+      VCAP::BigintMigration.create_trigger_function(self, :events)
+      VCAP::BigintMigration.backfill_id(self, :events)
+      VCAP::BigintMigration.add_check_constraint(self, :events)
+      # ...
+    end
+  end
+end

lib/database/bigint_migration.rb

@@ -62,6 +62,56 @@ def backfill(logger, db, table, batch_size: 10_000, iterations: -1)
       logger.info("finished bigint backfill on table '#{table}'")
     end
 
+    def migration_completed?(db, table)
+      column_type(db, table, :id) == 'bigint'
+    end
+
+    def migration_skipped?(db, table)
+      !column_exists?(db, table, :id_bigint)
+    end
+
+    def add_check_constraint(db, table)
+      return if check_constraint_exists?(db, table, :check_id_bigint_matches_id)
+
+      db.alter_table(table) do
+        add_constraint(:check_id_bigint_matches_id) do
+          Sequel.lit('id_bigint IS NOT NULL AND id_bigint = id')
+        end
+      end
+    end
+
+    def drop_check_constraint(db, table)
+      return unless check_constraint_exists?(db, table, :check_id_bigint_matches_id)
+
+      db.alter_table(table) do
+        drop_constraint(:check_id_bigint_matches_id)
+      end
+    end
+
+    def drop_pk_id_column(db, table)
+      db.drop_column(table, :id, if_exists: true)
+    end
+
+    def add_pk_id_column(db, table)
+      return if column_exists?(db, table, :id)
+
+      db.alter_table(table) do
+        add_primary_key(:id)
+      end
+    end
+
+    def backfill_id(db, table)
+      batch_size = 10_000
+      loop do
+        updated_rows = db.
+                       from(table, :batch).
+                       with(:batch, db[table].select(:id_bigint).where(id: nil).order(:id_bigint).limit(batch_size).for_update.skip_locked).
+                       where(Sequel.qualify(table, :id_bigint) => :batch__id_bigint).
+                       update(id: :batch__id_bigint)
+        break if updated_rows < batch_size
+      end
+    end
+
     private
 
     def column_type(db, table, column)
@@ -79,5 +129,17 @@ def trigger_name(table)
     def column_exists?(db, table, column)
       db[table].columns.include?(column)
     end
+
+    def check_constraint_exists?(db, table, constraint_name)
+      db.check_constraints(table).include?(constraint_name)
+    end
+
+    def backfill_batch(db, table, from_column, to_column, batch_size)
+      db.
+        from(table, :batch).
+        with(:batch, db[table].select(from_column).where(to_column => nil).order(from_column).limit(batch_size).for_update.skip_locked).
+        where(Sequel.qualify(table, from_column) => :"batch__#{from_column}").
+        update(to_column => :"batch__#{from_column}")
+    end
   end
 end

spec/migrations/20250603103400_bigint_migration_events_step3a_spec.rb

@@ -0,0 +1,17 @@
+require 'spec_helper'
+require 'migrations/helpers/bigint_migration_step3_shared_context'
+
+RSpec.describe 'bigint migration - events table - step3a', isolation: :truncation, type: :migration do
+  include_context 'bigint migration step3a' do
+    let(:migration_filename_step1) { '20250327142351_bigint_migration_events_step1.rb' }
+    let(:migration_filename_step3a) { '20250603103400_bigint_migration_events_step3a.rb' }
+    let(:table) { :events }
+    let(:insert) do
+      lambda do |db|
+        db[:events].insert(guid: SecureRandom.uuid, timestamp: Time.now.utc, type: 'type',
+                           actor: 'actor', actor_type: 'actor_type',
+                           actee: 'actee', actee_type: 'actee_type')
+      end
+    end
+  end
+end

spec/migrations/20250603103500_bigint_migration_events_step3b_spec.rb

@@ -0,0 +1,18 @@
+require 'spec_helper'
+require 'migrations/helpers/bigint_migration_step3_shared_context'
+
+RSpec.describe 'bigint migration - events table - step3b', isolation: :truncation, type: :migration do
+  include_context 'bigint migration step3b' do
+    let(:migration_filename_step1) { '20250327142351_bigint_migration_events_step1.rb' }
+    let(:migration_filename_step3a) { '20250603103400_bigint_migration_events_step3a.rb' }
+    let(:migration_filename_step3b) { '20250603103500_bigint_migration_events_step3b.rb' }
+    let(:table) { :events }
+    let(:insert) do
+      lambda do |db|
+        db[:events].insert(guid: SecureRandom.uuid, timestamp: Time.now.utc, type: 'type',
+                           actor: 'actor', actor_type: 'actor_type',
+                           actee: 'actee', actee_type: 'actee_type')
+      end
+    end
+  end
+end

spec/migrations/helpers/bigint_migration_step1_shared_context.rb

@@ -60,6 +60,10 @@
       context 'when the table is not empty' do
         let!(:old_id) { insert.call(db) }
 
+        after do
+          db[table].delete # Necessary to successfully run subsequent migrations in the after block of the migration shared context...
+        end
+
         it "does not change the id column's type" do
           expect(db).to have_table_with_column_and_type(table, :id, 'integer')
 
@@ -186,6 +190,10 @@
         run_migration
       end
 
+      after do
+        db[table].delete # Necessary to successfully run subsequent migrations in the after block of the migration shared context...
+      end
+
       it 'drops the id_bigint column' do
         expect(db).to have_table_with_column(table, :id_bigint)
 

spec/migrations/helpers/bigint_migration_step3_shared_context.rb

@@ -0,0 +1,207 @@
+require 'migrations/helpers/migration_shared_context'
+require 'database/bigint_migration'
+
+RSpec.shared_context 'bigint migration step3a' do
+  subject(:run_migration_step1) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index, allow_missing_migration_files: true) }
+
+  subject(:run_migration_step3a) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index_step3a, allow_missing_migration_files: true) }
+
+  let(:migration_filename) { migration_filename_step1 }
+  let(:current_migration_index_step3a) { migration_filename_step3a.match(/\A\d+/)[0].to_i }
+
+  include_context 'migration'
+
+  let(:skip_bigint_id_migration) { false }
+  let(:logger) { double(:logger, info: nil) }
+
+  before do
+    skip unless db.database_type == :postgres
+
+    allow_any_instance_of(VCAP::CloudController::Config).to receive(:get).with(:skip_bigint_id_migration).and_return(skip_bigint_id_migration)
+  end
+
+  describe 'up' do
+    context 'when migration step 1 was executed' do
+      context 'when the id_bigint column was added' do
+        before do
+          insert.call(db)
+          run_migration_step1
+        end
+
+        context 'when backfilling was completed' do
+          before do
+            VCAP::BigintMigration.backfill(logger, db, table)
+          end
+
+          it 'adds a check constraint' do
+            expect(db).not_to have_table_with_check_constraint(table)
+
+            run_migration_step3a
+
+            expect(db).to have_table_with_check_constraint(table)
+          end
+        end
+
+        context 'when backfilling was not completed' do
+          after do
+            db[table].delete # Necessary as the migration will be executed again in the after block of the migration shared context - and should not fail...
+          end
+
+          it 'fails ...' do
+            expect do
+              run_migration_step3a
+            end.to raise_error(/Failed to add check constraint on 'events' table!/)
+          end
+        end
+      end
+
+      context "when the migration was concluded (id column's type switched)" do
+        before do
+          db[table].delete
+          run_migration_step1
+        end
+
+        it 'does not add a check constraint' do
+          expect(db).not_to have_table_with_check_constraint(table)
+
+          run_migration_step3a
+
+          expect(db).not_to have_table_with_check_constraint(table)
+        end
+      end
+    end
+
+    context 'when migration step 1 was skipped' do
+      let(:skip_bigint_id_migration) { true }
+
+      before do
+        run_migration_step1
+      end
+
+      it 'does not add a check constraint' do
+        expect(db).not_to have_table_with_check_constraint(table)
+
+        run_migration_step3a
+
+        expect(db).not_to have_table_with_check_constraint(table)
+      end
+    end
+  end
+
+  describe 'down' do
+    subject(:run_rollback_step3a) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index_step3a - 1, allow_missing_migration_files: true) }
+
+    context 'when migration step 3a was executed' do
+      before do
+        insert.call(db)
+        run_migration_step1
+        VCAP::BigintMigration.backfill(logger, db, table)
+        run_migration_step3a
+      end
+
+      it 'drops the check constraint' do
+        expect(db).to have_table_with_check_constraint(table)
+
+        run_rollback_step3a
+
+        expect(db).not_to have_table_with_check_constraint(table)
+      end
+    end
+  end
+end
+
+RSpec.shared_context 'bigint migration step3b' do
+  subject(:run_migration_step1) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index, allow_missing_migration_files: true) }
+
+  subject(:run_migration_step3a) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index_step3a, allow_missing_migration_files: true) }
+
+  subject(:run_migration_step3b) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index_step3b, allow_missing_migration_files: true) }
+
+  let(:migration_filename) { migration_filename_step1 }
+  let(:current_migration_index_step3a) { migration_filename_step3a.match(/\A\d+/)[0].to_i }
+  let(:current_migration_index_step3b) { migration_filename_step3b.match(/\A\d+/)[0].to_i }
+
+  include_context 'migration'
+
+  let(:skip_bigint_id_migration) { false }
+  let(:logger) { double(:logger, info: nil) }
+
+  before do
+    skip unless db.database_type == :postgres
+
+    allow_any_instance_of(VCAP::CloudController::Config).to receive(:get).with(:skip_bigint_id_migration).and_return(skip_bigint_id_migration)
+  end
+
+  describe 'up' do
+    context 'when migration step 3a was executed' do
+      before do
+        insert.call(db)
+        run_migration_step1
+        VCAP::BigintMigration.backfill(logger, db, table)
+        run_migration_step3a
+      end
+
+      it 'drops the check constraint' do
+        expect(db).to have_table_with_check_constraint(table)
+
+        run_migration_step3b
+
+        expect(db).not_to have_table_with_check_constraint(table)
+      end
+
+      it 'drops the trigger function' do
+        expect(db).to have_trigger_function_for_table(table)
+
+        run_migration_step3b
+
+        expect(db).not_to have_trigger_function_for_table(table)
+      end
+
+      it 'drops the id primary key column' do
+        expect(db).to have_table_with_column(table, :id)
+
+        run_migration_step3b
+
+        expect(db).not_to have_table_with_column(table, :id)
+      end
+    end
+  end
+
+  describe 'down' do
+    subject(:run_rollback_step3b) { Sequel::Migrator.run(db, migrations_path, target: current_migration_index_step3b - 1, allow_missing_migration_files: true) }
+
+    context 'when migration step 3b was executed' do
+      before do
+        insert.call(db)
+        run_migration_step1
+        VCAP::BigintMigration.backfill(logger, db, table)
+        run_migration_step3a
+        run_migration_step3b
+      end
+
+      it 're-adds the id primary key column' do
+        expect(db).not_to have_table_with_column(table, :id)
+
+        run_rollback_step3b
+
+        expect(db).to have_table_with_column(table, :id)
+      end
+
+      it 're-creates the trigger function' do
+        expect(db).not_to have_trigger_function_for_table(table)
+
+        run_rollback_step3b
+
+        expect(db).to have_trigger_function_for_table(table)
+      end
+
+      it 're-adds the check constraint' do
+        expect(db).not_to have_table_with_check_constraint(table)
+
+        run_rollback_step3b
+
+        expect(db).to have_table_with_check_constraint(table)
+      end
+    end
+  end
+end

spec/migrations/helpers/matchers.rb

@@ -56,3 +56,9 @@
     db[table].where(column => nil).any?
   end
 end
+
+RSpec::Matchers.define :have_table_with_check_constraint do |table|
+  match do |db|
+    db.check_constraints(table).include?(:check_id_bigint_matches_id)
+  end
+end