Rails ActiveRecord UNION scope extension - adds chainable union functionality to ActiveRecord

union_scope.rb

# MIT License, Copyright (c) 2019 Brian Bugh

# Do a UNION on two or more scopes, combining the results. Unlike Rails arel
# unions/`or`, this will preserve the ActiveRecord_Relation and allows for
# `joins` and `references`. It can also be chained with a scope without issue.
#
# NOTE: the scope `select` for each union must match; you can't select a smaller
#       column subset and UNION it (this is a limitation of SQL).
#
# ==== Examples
#
#   User.where(name: "John").union(User.where(id: 5))
#   User.union(User.where(name: "John"), User.where(id: 5))
#   User.union(User.where(id: 5), team1.users, team2.users)
#
# Adapted, fixed, and improved from:
#  https://gist.github.com/lsiden/260167a4d3574a580d97
#  https://gist.github.com/tlowrimore/5162327
#
module UnionScope
  extend ActiveSupport::Concern

  module ClassMethods
    def union(*scopes)
      return all if scopes.length.zero?

      unless scopes.all? { |s| s.is_a? ActiveRecord::Relation }
        raise ArgumentError, "Scopes must be ActiveRecord::Relation objects."
      end

      id_column = "#{table_name}.id"
      scope_default = "(#{default_scoped.select(id_column).to_sql})"
      subquery = [self, *scopes]
                 .map { |s| "(#{s.select(id_column).to_sql})" }
                 .reject { |q| q == "()" || q == scope_default } # Remove empty or `all`
                 .join(" UNION ")
      return default_scoped.where("#{id_column} IN (#{subquery})") if subquery.present?

      all
    end
  end
end

union_scope_spec.rb

# MIT License, Copyright (c) 2019 Brian Bugh

require 'rails_helper'

describe UnionScope do
  let(:users) { create_list(:user, 2) }

  shared_examples 'shared examples' do |base_scope|
    it 'returns merged scopes' do
      result = base_scope.union(User.where(id: users[0].id), User.where(id: users[1].id))
      expect(result).to match_array users
    end

    context 'with joins/references' do
      let(:users) { create_list(:user, 2) { |u| u.update(team: create(:team)) } }

      it 'returns merged scopes' do
        result = base_scope.union(User.joins(:team).where(team: users.first.team), User.joins(:team).where(team: users.second.team))
        expect(result).to match_array(User.all)
      end
    end

    context 'with none in union' do
      it 'removes none and returns results' do
        result = base_scope.union(User.none, User.where(id: users[0].id))
        expect(result).to match_array([users[0]])
      end
    end

    # if the model uses default_scope then the union should respect that
    describe 'with default_scope' do
      before do
        allow(User).to receive(:default_scoped).and_return(User.where(email: "superman@kansas.gov"))
        users.first.update(email: "superman@kansas.gov")
        users.second.update(email: "president.lex.luthor@whitehouse.gov")
      end

      it 'respects default scope' do
        result = base_scope.union(User.where(id: users[0].id), User.where(id: users[1].id))
        expect(result).to eq [users.first]
      end
    end
  end

  describe 'union on default scope' do
    include_examples 'shared examples', User
  end

  describe 'union on all scope' do
    include_examples 'shared examples', User.all
  end

  describe 'union on none scope' do
    include_examples 'shared examples', User.none
  end

  describe 'chaining examples' do
    context 'default scope' do
      it 'returns merged scopes' do
        result = User.where(id: users[0].id).union(User.where(id: users[1].id))
        expect(result).to match_array users
      end
    end

    context 'all scope' do
      it 'returns merged scopes' do
        result = User.all.where(id: users[0].id).union(User.where(id: users[1].id))
        expect(result).to match_array users
      end
    end

    context 'none scope' do
      it 'returns merged scopes without base scope' do
        result = User.none.where(id: users[0].id).union(User.where(id: users[1].id))
        expect(result).to match_array [users.second]
      end
    end
  end

  describe 'passing a class instead of relation' do
    it 'raises an ArgumentError' do
      expect do
        User.where(id: users[0].id).union(User.find(users[1].id))
      end.to raise_error ArgumentError
    end
  end
end