Populate timelogs.project_id

We have added this redundant column for performance
optimizations in a previous merge request. Now we need
to populate it

Changelog: added

changelogs/unreleased/populate-timelogs-project_id.yml

+---
+title: Populate timelogs.project_id
+merge_request: 60439
+author: Lee Tickett @leetickett
+type: added

db/post_migrate/20210427212034_schedule_update_timelogs_project_id.rb

+# frozen_string_literal: true
+
+class ScheduleUpdateTimelogsProjectId < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+  BATCH_SIZE = 50_000
+  DELAY_INTERVAL = 2.minutes
+  MIGRATION = 'UpdateTimelogsProjectId'
+
+  disable_ddl_transaction!
+
+  class Timelog < ActiveRecord::Base
+    include EachBatch
+
+    self.table_name = 'timelogs'
+    self.inheritance_column = :_type_disabled
+  end
+
+  def up
+    queue_background_migration_jobs_by_range_at_intervals(
+      Timelog.all,
+      MIGRATION,
+      DELAY_INTERVAL,
+      batch_size: BATCH_SIZE
+    )
+  end
+
+  def down
+    # no-op
+  end
+end

db/schema_migrations/20210427212034

+2ffe65c4abcb8f638198943e1b74de710387438fb7c93addb05ccb3e86729934
\ No newline at end of file

lib/gitlab/background_migration/update_timelogs_project_id.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # Class to populate project_id for timelogs
+    class UpdateTimelogsProjectId
+      BATCH_SIZE = 1000
+
+      def perform(start_id, stop_id)
+        (start_id..stop_id).step(BATCH_SIZE).each do |offset|
+          update_issue_timelogs(offset, offset + BATCH_SIZE)
+          update_merge_request_timelogs(offset, offset + BATCH_SIZE)
+        end
+      end
+
+      def update_issue_timelogs(batch_start, batch_stop)
+        execute(<<~SQL)
+          UPDATE timelogs
+          SET project_id = issues.project_id
+          FROM issues
+          WHERE issues.id = timelogs.issue_id
+          AND timelogs.id BETWEEN #{batch_start} AND #{batch_stop}
+          AND timelogs.project_id IS NULL;
+        SQL
+      end
+
+      def update_merge_request_timelogs(batch_start, batch_stop)
+        execute(<<~SQL)
+          UPDATE timelogs
+          SET project_id = merge_requests.target_project_id
+          FROM merge_requests
+          WHERE merge_requests.id = timelogs.merge_request_id
+          AND timelogs.id BETWEEN #{batch_start} AND #{batch_stop}
+          AND timelogs.project_id IS NULL;
+        SQL
+      end
+
+      def execute(sql)
+        @connection ||= ::ActiveRecord::Base.connection
+        @connection.execute(sql)
+      end
+    end
+  end
+end

spec/lib/gitlab/background_migration/update_timelogs_project_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::UpdateTimelogsProjectId, schema: 20210427212034 do
+  let!(:namespace) { table(:namespaces).create!(name: 'namespace', path: 'namespace') }
+  let!(:project1) { table(:projects).create!(namespace_id: namespace.id) }
+  let!(:project2) { table(:projects).create!(namespace_id: namespace.id) }
+  let!(:issue1) { table(:issues).create!(project_id: project1.id) }
+  let!(:issue2) { table(:issues).create!(project_id: project2.id) }
+  let!(:merge_request1) { table(:merge_requests).create!(target_project_id: project1.id, source_branch: 'master', target_branch: 'feature') }
+  let!(:merge_request2) { table(:merge_requests).create!(target_project_id: project2.id, source_branch: 'master', target_branch: 'feature') }
+  let!(:timelog1) { table(:timelogs).create!(issue_id: issue1.id, time_spent: 60) }
+  let!(:timelog2) { table(:timelogs).create!(issue_id: issue1.id, time_spent: 60) }
+  let!(:timelog3) { table(:timelogs).create!(issue_id: issue2.id, time_spent: 60) }
+  let!(:timelog4) { table(:timelogs).create!(merge_request_id: merge_request1.id, time_spent: 600) }
+  let!(:timelog5) { table(:timelogs).create!(merge_request_id: merge_request1.id, time_spent: 600) }
+  let!(:timelog6) { table(:timelogs).create!(merge_request_id: merge_request2.id, time_spent: 600) }
+  let!(:timelog7) { table(:timelogs).create!(issue_id: issue2.id, time_spent: 60, project_id: project1.id) }
+  let!(:timelog8) { table(:timelogs).create!(merge_request_id: merge_request2.id, time_spent: 600, project_id: project1.id) }
+
+  describe '#perform' do
+    context 'when timelogs belong to issues' do
+      it 'sets correct project_id' do
+        subject.perform(timelog1.id, timelog3.id)
+
+        expect(timelog1.reload.project_id).to eq(issue1.project_id)
+        expect(timelog2.reload.project_id).to eq(issue1.project_id)
+        expect(timelog3.reload.project_id).to eq(issue2.project_id)
+      end
+    end
+
+    context 'when timelogs belong to merge requests' do
+      it 'sets correct project ids' do
+        subject.perform(timelog4.id, timelog6.id)
+
+        expect(timelog4.reload.project_id).to eq(merge_request1.target_project_id)
+        expect(timelog5.reload.project_id).to eq(merge_request1.target_project_id)
+        expect(timelog6.reload.project_id).to eq(merge_request2.target_project_id)
+      end
+    end
+
+    context 'when timelogs already belong to projects' do
+      it 'does not update the project id' do
+        subject.perform(timelog7.id, timelog8.id)
+
+        expect(timelog7.reload.project_id).to eq(project1.id)
+        expect(timelog8.reload.project_id).to eq(project1.id)
+      end
+    end
+  end
+end

spec/migrations/schedule_update_timelogs_project_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require Rails.root.join('db', 'post_migrate', '20210427212034_schedule_update_timelogs_project_id.rb')
+
+RSpec.describe ScheduleUpdateTimelogsProjectId do
+  let!(:namespace) { table(:namespaces).create!(name: 'namespace', path: 'namespace') }
+  let!(:project) { table(:projects).create!(namespace_id: namespace.id) }
+  let!(:issue) { table(:issues).create!(project_id: project.id) }
+  let!(:merge_request) { table(:merge_requests).create!(target_project_id: project.id, source_branch: 'master', target_branch: 'feature') }
+  let!(:timelog1) { table(:timelogs).create!(issue_id: issue.id, time_spent: 60) }
+  let!(:timelog2) { table(:timelogs).create!(merge_request_id: merge_request.id, time_spent: 600) }
+  let!(:timelog3) { table(:timelogs).create!(merge_request_id: merge_request.id, time_spent: 60) }
+  let!(:timelog4) { table(:timelogs).create!(issue_id: issue.id, time_spent: 600) }
+
+  it 'correctly schedules background migrations' do
+    stub_const("#{described_class}::BATCH_SIZE", 2)
+
+    Sidekiq::Testing.fake! do
+      freeze_time do
+        migrate!
+
+        expect(described_class::MIGRATION)
+          .to be_scheduled_delayed_migration(2.minutes, timelog1.id, timelog2.id)
+
+        expect(described_class::MIGRATION)
+          .to be_scheduled_delayed_migration(4.minutes, timelog3.id, timelog4.id)
+
+        expect(BackgroundMigrationWorker.jobs.size).to eq(2)
+      end
+    end
+  end
+end