Merge branch 'sh-fix-orphaned-tokens-migration' into 'master'

Fix OrphanedInviteTokensCleanup migration

See merge request gitlab-org/gitlab!68784

db/post_migrate/20210809123658_orphaned_invite_tokens_cleanup.rb

@@ -6,15 +6,12 @@ class OrphanedInviteTokensCleanup < ActiveRecord::Migration[6.1]
   disable_ddl_transaction!
 
   TMP_INDEX_NAME = 'tmp_idx_orphaned_invite_tokens'
-  QUERY_CONDITION = "invite_token IS NOT NULL and invite_accepted_at IS NOT NULL and invite_accepted_at < created_at"
 
   def up
-    membership = define_batchable_model('members')
+    add_concurrent_index('members', :id, where: query_condition, name: TMP_INDEX_NAME)
 
-    add_concurrent_index('members', :id, where: QUERY_CONDITION, name: TMP_INDEX_NAME)
-
-    membership.where(QUERY_CONDITION).pluck(:id).each_slice(10) do |group|
-      membership.where(id: group).where(QUERY_CONDITION).update_all(invite_token: nil)
+    membership.where(query_condition).pluck(:id).each_slice(10) do |group|
+      membership.where(id: group).where(query_condition).update_all(invite_token: nil)
     end
 
     remove_concurrent_index_by_name('members', TMP_INDEX_NAME)
@@ -25,4 +22,30 @@ class OrphanedInviteTokensCleanup < ActiveRecord::Migration[6.1]
 
     # This migration is irreversible
   end
+
+  private
+
+  def membership
+    @membership ||= define_batchable_model('members')
+  end
+
+  # We need to ensure we're comparing timestamp with time zones across
+  # the board since that is an immutable comparison. Some database
+  # schemas have a mix of timestamp without time zones and and timestamp
+  # with time zones: https://gitlab.com/groups/gitlab-org/-/epics/2473
+  def query_condition
+    "invite_token IS NOT NULL and invite_accepted_at IS NOT NULL and #{timestamptz("invite_accepted_at")} < #{timestamptz("created_at")}"
+  end
+
+  def timestamptz(name)
+    if column_type(name) == "timestamp without time zone"
+      "TIMEZONE('UTC', #{name})"
+    else
+      name
+    end
+  end
+
+  def column_type(name)
+    membership.columns_hash[name].sql_type
+  end
 end

spec/migrations/orphaned_invite_tokens_cleanup_spec.rb

@@ -16,7 +16,7 @@ RSpec.describe OrphanedInviteTokensCleanup, :migration do
     table(:members).create!(defaults.merge(extra_attributes))
   end
 
-  describe '#up', :aggregate_failures do
+  shared_examples 'removes orphaned invite tokens' do
     it 'removes invite tokens for accepted records with invite_accepted_at < created_at' do
       record1 = create_member(invite_token: 'foo', invite_accepted_at: 1.day.ago, created_at: 1.hour.ago)
       record2 = create_member(invite_token: 'foo2', invite_accepted_at: nil, created_at: 1.hour.ago)
@@ -29,4 +29,22 @@ RSpec.describe OrphanedInviteTokensCleanup, :migration do
       expect(table(:members).find(record3.id).invite_token).to eq 'foo3'
     end
   end
+
+  describe '#up', :aggregate_failures do
+    it_behaves_like 'removes orphaned invite tokens'
+  end
+
+  context 'when there is a mix of timestamptz and timestamp types' do
+    around do |example|
+      ActiveRecord::Base.connection.execute "ALTER TABLE members alter created_at type timestamp with time zone"
+
+      example.run
+
+      ActiveRecord::Base.connection.execute "ALTER TABLE members alter created_at type timestamp without time zone"
+    end
+
+    describe '#up', :aggregate_failures do
+      it_behaves_like 'removes orphaned invite tokens'
+    end
+  end
 end