Added support for has_and_belongs_to_many associations in eager loading

#1064 [Dan Peterson]

git-svn-id: http://svn-commit.rubyonrails.org/rails/trunk@1132 5ecf4fe2-1ee6-0310-87b1-e25e094e27de

activerecord/lib/active_record/associations.rb

@@ -628,67 +628,154 @@ def add_multiple_associated_save_callbacks(association_name)
 
 
         def find_with_associations(options = {})
-          reflections = [ options[:include] ].flatten.collect { |association| reflect_on_association(association) }
-          rows = connection.select_all(construct_finder_sql_with_included_associations(options, reflections), "#{name} Load Including Associations")
-          records = rows.collect { |row| instantiate(extract_record(table_name, row)) }.uniq
+          reflections          = reflect_on_included_associations(options[:include])
+          schema_abbreviations = generate_schema_abbreviations(reflections)
+          primary_key_table    = generate_primary_key_table(reflections, schema_abbreviations)
 
-          reflections.each do |reflection| 
-            records.each do |record|
+          rows    = select_all_rows(options, schema_abbreviations, reflections)
+          records = extract_and_instantiate_records(schema_abbreviations, rows)
+
+          assign_associations_to_records(rows, records, reflections, schema_abbreviations, primary_key_table)
+
+          return records
+        end
+
+        def assign_associations_to_records(rows, records, reflections, schema_abbreviations, primary_key_table)
+          records.each do |record|
+            reflections.each do |reflection|
               case reflection.macro
-                when :has_many
-                  record.send(reflection.name).target = extract_association_for_record(record, rows, reflection)
+                when :has_many, :has_and_belongs_to_many
+                  record.send(reflection.name).target = 
+                    extract_association_for_record(record, schema_abbreviations, primary_key_table, rows, reflection)
                 when :has_one, :belongs_to
-                  record.send("set_#{reflection.name}_target", extract_association_for_record(record, rows, reflection).first)
+                  record.send(
+                    "set_#{reflection.name}_target", 
+                    extract_association_for_record(record, schema_abbreviations, primary_key_table, rows, reflection).first
+                  )
               end
             end
           end
+        end
+
+        def generate_schema_abbreviations(reflections)
+          schema = [ [ table_name, columns.collect { |c| c.name } ] ]
+          schema += reflections.collect { |r| [ r.klass.table_name, r.klass.columns.collect { |c| c.name } ] }
+
+          schema_abbreviations = {}
+          schema.each_with_index do |table_and_columns, i|
+            table, columns = table_and_columns
+            columns.each_with_index { |column, j| schema_abbreviations["t#{i}_r#{j}"] = [ table, column ] }
+          end
 
-          return records
+          return schema_abbreviations
         end
 
-        def construct_finder_sql_with_included_associations(options, reflections)
-          sql = "SELECT #{selected_columns(table_name, columns)}"
-          reflections.each { |reflection| sql << ", #{selected_columns(reflection.klass.table_name, reflection.klass.columns)}" }
-          sql << " FROM #{table_name} "
-
-          reflections.each do |reflection| 
-            case reflection.macro
-              when :has_many, :has_one
-                sql << " LEFT JOIN #{reflection.klass.table_name} ON " +
-                  "#{reflection.klass.table_name}.#{reflection.options[:foreign_key] || table_name.classify.foreign_key} = #{table_name}.#{primary_key} "
-              when :belongs_to
-                sql << " LEFT JOIN #{reflection.klass.table_name} ON " +
-                  "#{reflection.klass.table_name}.#{reflection.klass.primary_key} = #{table_name}.#{reflection.options[:foreign_key] || reflection.klass.table_name.classify.foreign_key} "
-            end
+        def generate_primary_key_table(reflections, schema_abbreviations)
+          primary_key_lookup_table = {}
+          primary_key_lookup_table[table_name] = 
+            schema_abbreviations.find { |cn, tc| tc == [ table_name, primary_key ] }.first
+
+          reflections.collect do |reflection| 
+            primary_key_lookup_table[reflection.klass.table_name] = schema_abbreviations.find { |cn, tc| 
+              tc == [ reflection.klass.table_name, reflection.klass.primary_key ]
+            }.first
           end
+
+          return primary_key_lookup_table
+        end
+
 
+        def construct_finder_sql_with_included_associations(options, schema_abbreviations, reflections)
+          habtm_associations = reflections.find_all { |r| r.macro == :has_and_belongs_to_many }
+
+          sql = "SELECT #{column_aliases(schema_abbreviations)} FROM #{table_name}"
+          add_habtm_join_tables!(habtm_associations, sql)
+          sql << " "
+
+          add_association_joins!(reflections, sql)
           sql << "#{options[:joins]} " if options[:joins]
+
+          add_habtm_conditions!(habtm_associations, options)
           add_conditions!(sql, options[:conditions])
+
           sql << "ORDER BY #{options[:order]} " if options[:order]
 
           return sanitize_sql(sql)
         end
 
-        def extract_association_for_record(record, rows, reflection)
-          association = rows.collect do |row| 
-            if row["#{table_name}__#{primary_key}"] == record.id.to_s && !row["#{reflection.klass.table_name}__#{reflection.klass.primary_key}"].nil?
-              reflection.klass.send(:instantiate, extract_record(reflection.klass.table_name, row))
+        def column_aliases(schema_abbreviations)
+          schema_abbreviations.collect { |cn, tc| "#{tc.join(".")} AS #{cn}" }.join(", ")
+        end
+
+        def add_habtm_join_tables!(habtm_associations, sql)
+          return if habtm_associations.empty?
+          sql << ", " + habtm_associations.collect { |a| [ a.klass.table_name, a.options[:join_table] ] }.join(", ")
+        end
+
+        def add_habtm_conditions!(habtm_associations, options)
+          return if habtm_associations.empty?
+          options[:conditions] = [
+            options[:conditions],
+            habtm_associations.collect { |r|
+              join_table = r.options[:join_table]
+              "#{join_table}.#{table_name.classify.foreign_key} = #{table_name}.#{primary_key} AND " +
+              "#{join_table}.#{r.klass.table_name.classify.foreign_key} = #{r.klass.table_name}.#{r.klass.primary_key}"
+            }
+          ].compact.join(" AND ")
+        end
+
+        def add_association_joins!(reflections, sql)
+          reflections.each { |reflection| sql << association_join(reflection) }
+        end
+
+        def association_join(reflection)
+          case reflection.macro
+            when :has_many, :has_one
+              " LEFT JOIN #{reflection.klass.table_name} ON " +
+              "#{reflection.klass.table_name}.#{reflection.options[:foreign_key] || table_name.classify.foreign_key} = " +
+              "#{table_name}.#{primary_key} "
+            when :belongs_to
+              " LEFT JOIN #{reflection.klass.table_name} ON " +
+              "#{reflection.klass.table_name}.#{reflection.klass.primary_key} = " +
+              "#{table_name}.#{reflection.options[:foreign_key] || reflection.klass.table_name.classify.foreign_key} "
+            else
+              ""
+          end          
+        end
+
+
+        def extract_and_instantiate_records(schema_abbreviations, rows)
+          rows.collect { |row| instantiate(extract_record(schema_abbreviations, table_name, row)) }.uniq
+        end
+
+        def extract_association_for_record(record, schema_abbreviations, primary_key_table, rows, reflection)
+          association = rows.collect do |row|
+            if row[primary_key_table[table_name]].to_s == record.id.to_s && !row[primary_key_table[reflection.klass.table_name]].nil?
+              reflection.klass.send(:instantiate, extract_record(schema_abbreviations, reflection.klass.table_name, row))
             end
           end
 
           return association.uniq.compact
         end
 
-        def extract_record(table_name, row)
+        def extract_record(schema_abbreviations, table_name, row)
           row.inject({}) do |record, pair|
-            prefix, column_name = pair.first.split("__")
+            prefix, column_name = schema_abbreviations[pair.first]
             record[column_name] = pair.last if prefix == table_name
             record
           end
         end
 
-        def selected_columns(table_name, columns)
-          columns.collect { |column| "#{table_name}.#{column.name} as #{table_name}__#{column.name}" }.join(", ")
+
+        def reflect_on_included_associations(associations)
+          [ associations ].flatten.collect { |association| reflect_on_association(association) }
+        end
+
+        def select_all_rows(options, schema_abbreviations, reflections)
+          connection.select_all(
+            construct_finder_sql_with_included_associations(options, schema_abbreviations, reflections), 
+            "#{name} Load Including Associations"
+          )
         end
     end
   end

activerecord/test/associations_go_eager_test.rb

@@ -2,9 +2,10 @@
 require 'fixtures/post'
 require 'fixtures/comment'
 require 'fixtures/author'
+require 'fixtures/category'
 
 class EagerAssociationTest < Test::Unit::TestCase
-  fixtures :posts, :comments, :authors
+  fixtures :posts, :comments, :authors, :categories, :categories_posts
 
   def test_loading_with_one_association
     posts = Post.find(:all, :include => :comments)
@@ -17,8 +18,9 @@ def test_loading_with_one_association
   end
 
   def test_loading_with_multiple_associations
-    posts = Post.find(:all, :include => [ :comments, :author ])
+    posts = Post.find(:all, :include => [ :comments, :author, :categories ])
     assert_equal 2, posts.first.comments.size
+    assert_equal 2, posts.first.categories.size
     assert_equal @greetings.body, posts.first.comments.first.body
   end
 
@@ -32,10 +34,17 @@ def test_eager_association_loading_with_belongs_to
     assert_equal @welcome.title, comments.first.post.title
     assert_equal @thinking.title, comments.last.post.title
   end
+
+  def test_eager_association_loading_with_habtm
+    posts = Post.find(:all, :include => :categories)
+    assert_equal 2, posts.first.categories.size
+    assert_equal 1, posts.last.categories.size
+    assert_equal @technology.name, posts.first.categories.last.name
+    assert_equal @general.name, posts.last.categories.first.name
+  end
 
   def test_eager_with_inheritance
     posts = SpecialPost.find(:all, :include => [ :comments ])
-  end
-
+  end  
 end
 

activerecord/test/fixtures/categories.yml

@@ -0,0 +1,7 @@
+general:
+  id: 1
+  name: General
+
+technology:
+  id: 2
+  name: Technology

activerecord/test/fixtures/categories_posts.yml

@@ -0,0 +1,11 @@
+general_welcome:
+  category_id: 1
+  post_id: 1
+
+technology_welcome:
+  category_id: 2
+  post_id: 1
+
+general_thinking:
+  category_id: 1
+  post_id: 2

activerecord/test/fixtures/category.rb

@@ -0,0 +1,3 @@
+class Category < ActiveRecord::Base
+  has_and_belongs_to_many :posts
+end

activerecord/test/fixtures/db_definitions/db2.drop.sql

@@ -19,4 +19,6 @@ DROP TABLE posts;
 DROP TABLE comments;
 DROP TABLE authors;
 DROP TABLE tasks;
+DROP TABLE categories;
+DROP TABLE categories_posts;
 

activerecord/test/fixtures/db_definitions/db2.sql

@@ -154,3 +154,13 @@ CREATE TABLE tasks (
   starting timestamp default NULL,
   ending timestamp default NULL
 );
+
+CREATE TABLE categories (
+  id int generated by default as identity (start with +10000),
+  name varchar(255) NOT NULL
+);
+
+CREATE TABLE categories_posts (
+  category_id int NOT NULL,
+  post_id int NOT NULL
+);

activerecord/test/fixtures/db_definitions/mysql.drop.sql

@@ -19,3 +19,5 @@ DROP TABLE tasks;
 DROP TABLE posts;
 DROP TABLE comments;
 DROP TABLE authors;
+DROP TABLE categories;
+DROP TABLE categories_posts;

activerecord/test/fixtures/db_definitions/mysql.sql

@@ -155,3 +155,13 @@ CREATE TABLE `tasks` (
   `ending` datetime NOT NULL default '0000-00-00 00:00:00',
   PRIMARY KEY  (`id`)
 );
+
+CREATE TABLE `categories` (
+  `id` int(11) NOT NULL auto_increment,
+  `name` VARCHAR(255) NOT NULL
+);
+
+CREATE TABLE `categories_posts` (
+  `category_id` int(11) NOT NULL,
+  `post_id` int(11) NOT NULL
+);

activerecord/test/fixtures/db_definitions/oci.drop.sql

@@ -18,4 +18,6 @@ drop table posts;
 drop table comments;
 drop table authors;
 drop table computers;
+drop table categories;
+drop table categories_posts;
 drop sequence rails_sequence;

activerecord/test/fixtures/db_definitions/oci.sql

@@ -192,3 +192,13 @@ create table tasks (
   starting date default null,
   ending date default null
 );
+
+create table categories (
+  id integer not null primary key,
+  name varchar(255) default null
+);
+
+create table categories_posts (
+  category_id integer not null references developers initially deferred disable,
+  post_id int integer not null references developers initially deferred disable
+);

activerecord/test/fixtures/db_definitions/postgresql.drop.sql

@@ -19,4 +19,6 @@ DROP TABLE posts;
 DROP TABLE comments;
 DROP TABLE authors;
 DROP TABLE tasks;
+DROP TABLE categories;
+DROP TABLE categories_posts;
 

activerecord/test/fixtures/db_definitions/postgresql.sql

@@ -173,3 +173,13 @@ CREATE TABLE taske (
   ending timestamp,
   PRIMARY KEY (id)
 );
+
+CREATE TABLE categories (
+  id serial,
+  name varchar(255)
+);
+
+CREATE TABLE categories_posts (
+  category_id integer NOT NULL,
+  post_id int integer NOT NULL
+);

activerecord/test/fixtures/db_definitions/sqlite.drop.sql

@@ -19,3 +19,5 @@ DROP TABLE tasks;
 DROP TABLE posts;
 DROP TABLE comments;
 DROP TABLE authors;
+DROP TABLE categories;
+DROP TABLE categories_posts;

activerecord/test/fixtures/db_definitions/sqlite.sql

@@ -142,3 +142,12 @@ CREATE TABLE 'tasks' (
   'ending' DATETIME DEFAULT NULL
 );
 
+CREATE TABLE 'categories' (
+  'id' INTEGER NOT NULL PRIMARY KEY,
+  'name' VARCHAR(255) NOT NULL
+);
+
+CREATE TABLE 'categories_posts' (
+  'category_id' INTEGER NOT NULL,
+  'post_id' INTEGER NOT NULL
+);