Splitting DB on "Rails" - https://speakerdeck.com/izumin5210/splitting-db-on-rails

active_record-relation.rb

require "join_detector_ext"
ActiveRecord::Relation.prepend JoinDetectorExt

find_relations.rb

require 'optparse'

class RelationFinder
  def self.create(app)
    app.eager_load!
    new(app)
  end

  def initialize(app)
    @app = app
    @ar_classes = ActiveRecord::Base.descendants.reject(&:abstract_class?).freeze
    @ar_by_table_name = @ar_classes.index_by(&:table_name).freeze
  end

  def associated_tables_to(ar_class, excluded_table_names = [])
    @ar_classes
      .flat_map(&:reflect_on_all_associations)
      .select { |r| Object.const_defined?(r.class_name) && (r.table_name == ar_class.table_name) }
      .map(&:active_record)
      .reject { |ar| excluded_table_names.include? ar.table_name }
      .uniq
  end

  def associated_tables_from(ar_class, excluded_table_names = [])
    ar_class
      .reflect_on_all_associations
      .select { |r| Object.const_defined?(r.class_name) }
      .map { |r| ar_class_of(r.table_name) }
      .compact
      .reject { |ar| excluded_table_names.include? ar.table_name }
  end

  def ar_class_of(table_name)
    @ar_by_table_name[table_name]
  end
end

opts = {
  excluded: [],
}

ARGV.shift if ARGV[0] == '--'

OptionParser.new do |opt|
  opt.on('-e', '--exclude one,two,three', Array, "Excluded table names (default: #{opts[:excluded]})") {|v|
    opts[:excluded] = v
  }

  opt.parse!(ARGV)
end

def main(args, opts)
  table_name = args[0]
  excluded_tables = opts[:excluded]

  finder = RelationFinder.create(Rails.application)
  ar_class = finder.ar_class_of(table_name)

  if ar_class.blank?
    warn("There does not found an ActiveRecord::Base subclass of #{table_name} table")
    exit 1
  end

  from_classes = finder.associated_tables_to(ar_class, excluded_tables)
  to_classes = finder.associated_tables_from(ar_class, excluded_tables)
  first = (from_classes + to_classes).uniq

  excluded_tables.push(ar_class.table_name)
  excluded_tables.push(*first.map(&:table_name))

  second = first.flat_map { |ar|
    from_classes = finder.associated_tables_to(ar, excluded_tables)
    to_classes = finder.associated_tables_from(ar, excluded_tables)
    to_classes + from_classes
  }.uniq

  excluded_tables.push(*second.map(&:table_name))

  third = second.flat_map { |ar|
    from_classes = finder.associated_tables_to(ar, excluded_tables)
    to_classes = finder.associated_tables_from(ar, excluded_tables)
    to_classes + from_classes
  }.uniq

  puts "==> first  #{?- * 16}"
  puts first.map(&:table_name).join("\n")
  puts "==> second #{?- * 16}"
  puts second.map(&:table_name).join("\n")
  puts "==> third #{?- * 16}"
  puts third.map(&:table_name).join("\n")
end

main(ARGV, opts)

join_detector_ext.rb

module JoinDetectorExt
  ENABLED = !Rails.env.production? || ENV["JOIN_DETECTION"] == '1'

  @@previous_detected_join_query = ''

  def eager_loading?
    loading = super
    if ENABLED && (loading || joins_values.any? || left_outer_joins_values.any?)
      caller_app = 
        caller.find do |c|
          !c.to_s.start_with?(Bundler.bundle_path.to_s) && c.to_s.start_with?(Rails.root.to_s)
        end
      info = {
        caller:                  caller_app || "",
        table:                   table.name,
        class:                   klass.name,
        preload:                 flatten_recursively(preload_values),
        joins:                   flatten_recursively(joins_values),
        left_outer_joins:        flatten_recursively(left_outer_joins_values),
        eager_load:              flatten_recursively(eager_load_values),
        includes:                flatten_recursively(includes_values),
        joined_includes:         flatten_recursively(joined_includes_values),
        references_eager_loaded: flatten_recursively(references_eager_loaded_tables),
        request_controller:      $request_controller,
        request_action:          $request_action,
      }
      json = info.to_json
      if @@previous_detected_join_query != json
        TD.event.post("join_queries", info)
        @@previous_detected_join_query = json
      end
    end
    loading
  end

  private 

  # ref: https://github.com/rails/rails/blob/813af4655f9bf3c712cf50205eebd337070cee52/activerecord/lib/active_record/relation.rb#L687-L702
  def references_eager_loaded_tables
    joined_tables = arel.join_sources.map do |join|
      if join.is_a?(Arel::Nodes::StringJoin)
        tables_in_string(join.left)
      else
        [join.left.table_name, join.left.table_alias]
      end
    end

    joined_tables += [table.name, table.table_alias]

    # always convert table names to downcase as in Oracle quoted table names are in uppercase
    joined_tables = joined_tables.flatten.compact.map(&:downcase).uniq

    references_values - joined_tables
  end

  def flatten_recursively(arr)
    arr.flat_map do |i|
      case i
      when Hash then hash_to_array(i)
      when Array then flatten_recursively(i)
      when Arel::Nodes::Node then i.to_sql
      when ActiveRecord::Associations::JoinDependency then join_dependency_to_table_names(i)
      else i.to_s
      end
    end
  end

  def hash_to_array(h)
    h.keys + (h.values.flat_map { |v|
      case v
      when Hash then hash_to_array(v)
      when Array then flatten_recursively(v)
      when Arel::Nodes::Node then v.to_sql
      when ActiveRecord::Associations::JoinDependency then join_dependency_to_table_names(v)
      else v.to_s
      end
    })
  end

  def join_dependency_to_table_names(jd)
    jd.aliases.instance_variable_get(:@tables).map { |a| a.table.table_name }
  end
end