BackgroundMigration: relation defined with `includes` raises exception in batch iterator

[davejachimiak]

Given a Background migration with a relation that uses includes, e.g.

module OnlineMigrations::BackgroundMigrations
  class BMigration < OnlineMigrations::BackgroundMigration
    def relation
      Model.includes(association: :other_association)
    end

    # ...
  end
end

, running it will result in an exception:

ActiveModel::MissingAttributeError: missing attribute 'association_id' for Model (ActiveModel::MissingAttributeError)

Likely culprit: in BatchIterator, the base_relation — which only selects the primary key — is used to query against an arel clause from the relation. For some reason, ActiveRecord expects the foreign key to exist on the base_relation when passing an arel clause to where that's derived from a relation that uses includes.

[davejachimiak]

# frozen_string_literal: true

require "bundler/inline"

gemfile(true) do
  source "https://rubygems.org"

  gem "rails"
  # If you want to test against edge Rails replace the previous line with this:
  # gem "rails", github: "rails/rails", branch: "main"

  gem "pg"

  gem 'online_migrations', '~> 0.20'
end

require "active_record"
require 'rails/generators'
require "minitest/autorun"
require "logger"
require "online_migrations"
require "rake"

system("createdb om-bug-demo-test")
Rails::Generators.invoke("online_migrations:install")
ActiveRecord::Base.establish_connection(adapter: "postgresql", database: "om-bug-demo-test")
ActiveRecord::Base.logger = Logger.new(STDOUT)

ActiveRecord::Schema.define do
  create_table :foos, force: true do |t|
    t.belongs_to :user
    t.integer :bar_id
  end

  create_table :users, force: true do |t|
    t.belongs_to :bar
  end

  create_table :bars, force: true do |t|
  end
end

Dir["./db/migrate/*.rb"].each {|file| puts file; require_relative file }

ActiveRecord::Migration.run InstallOnlineMigrations

class Foo < ActiveRecord::Base
  belongs_to :user
end

class User < ActiveRecord::Base
  belongs_to :bar
end

class Bar < ActiveRecord::Base
end

module OnlineMigrations::BackgroundMigrations
  class BackfillFoosBars < OnlineMigrations::BackgroundMigration
    def relation
      Foo.includes(user: :bar).where(bar_id: nil)
    end

    def process_batch(relation)
      relation.each do |foo|
        next if foo.bar_id.present?

        foo.update(bar_id: 5000)
      end
    end
  end
end

class BackfillMigration < ActiveRecord::Migration::Current # or use a specific version via `Migration[number]`
  def up
    enqueue_background_data_migration('BackfillFoosBars')
  end

  def down
  end
end

class BugTest < ActiveSupport::TestCase
  def test_migration_up
    bar = Bar.create
    user = User.create(bar: bar)
    Foo.create(user: user)

    BackfillMigration.migrate(:up)

    bar_id = Foo.first.bar_id

    assert_equal bar_id, 5000
  end
end