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

20 files changed, 226 insertions, 76 deletions

diff --git a/activerecord/lib/active_record/associations.rb b/activerecord/lib/active_record/associations.rb
index 4413a79a40..9253abd14d 100755
--- a/activerecord/lib/active_record/associations.rb
+++ b/activerecord/lib/active_record/associations.rb
@@ -628,67 +628,154 @@ module ActiveRecord
 
 
         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
diff --git a/activerecord/test/associations_go_eager_test.rb b/activerecord/test/associations_go_eager_test.rb
index a8e99c1bc6..36f79265a5 100644
--- a/activerecord/test/associations_go_eager_test.rb
+++ b/activerecord/test/associations_go_eager_test.rb
@@ -2,9 +2,10 @@ require 'abstract_unit'
 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 @@ class EagerAssociationTest < Test::Unit::TestCase
   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 @@ class EagerAssociationTest < Test::Unit::TestCase
     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
 
diff --git a/activerecord/test/fixtures/categories.yml b/activerecord/test/fixtures/categories.yml
new file mode 100644
index 0000000000..2ce34d6f9f
--- /dev/null
+++ b/activerecord/test/fixtures/categories.yml
@@ -0,0 +1,7 @@
+general:
+  id: 1
+  name: General
+
+technology:
+  id: 2
+  name: Technology
diff --git a/activerecord/test/fixtures/categories_posts.yml b/activerecord/test/fixtures/categories_posts.yml
new file mode 100644
index 0000000000..04eeb8f46c
--- /dev/null
+++ b/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
diff --git a/activerecord/test/fixtures/category.rb b/activerecord/test/fixtures/category.rb
new file mode 100644
index 0000000000..751413c0db
--- /dev/null
+++ b/activerecord/test/fixtures/category.rb
@@ -0,0 +1,3 @@
+class Category < ActiveRecord::Base
+  has_and_belongs_to_many :posts
+end
diff --git a/activerecord/test/fixtures/db_definitions/db2.drop.sql b/activerecord/test/fixtures/db_definitions/db2.drop.sql
index e718d4b321..844bf3e2f8 100644
--- a/activerecord/test/fixtures/db_definitions/db2.drop.sql
+++ b/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;
 
diff --git a/activerecord/test/fixtures/db_definitions/db2.sql b/activerecord/test/fixtures/db_definitions/db2.sql
index a24318bf4b..35efa3f6c7 100644
--- a/activerecord/test/fixtures/db_definitions/db2.sql
+++ b/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
+);
diff --git a/activerecord/test/fixtures/db_definitions/drop_oracle_tables.sql b/activerecord/test/fixtures/db_definitions/drop_oracle_tables.sql
deleted file mode 100644
index 2eb574df75..0000000000
--- a/activerecord/test/fixtures/db_definitions/drop_oracle_tables.sql
+++ /dev/null
@@ -1,37 +0,0 @@
-DROP TABLE accounts;
-DROP SEQUENCE accounts_id;
-DROP TABLE companies;
-DROP SEQUENCE companies_id;
-DROP TABLE topics;
-DROP SEQUENCE topics_id;
-DROP TABLE developers;
-DROP SEQUENCE developers_id;
-DROP TABLE projects;
-DROP SEQUENCE projects_id;
-DROP TABLE developers_projects;
-DROP SEQUENCE developers_projects_id;
-DROP TABLE customers;
-DROP SEQUENCE customers_id;
-DROP TABLE movies;
-DROP SEQUENCE movies_id;
-DROP TABLE subscribers;
-DROP SEQUENCE subscribers_id;
-DROP TABLE booleantests;
-DROP SEQUENCE booleantests_id;
-DROP TABLE auto_id_tests;
-DROP SEQUENCE auto_id_tests_id;
-DROP TABLE entrants;
-DROP SEQUENCE entrants_id;
-DROP TABLE colnametests;
-DROP SEQUENCE colnametests_id;
-DROP TABLE mixins;
-DROP SEQUENCE mixins_id;
-DROP TABLE people;
-DROP SEQUENCE people_id;
-DROP TABLE binaries;
-DROP SEQUENCE binaries_id;
-DROP TABLE computers;
-DROP SEQUENCE computers_id;
-DROP TABLE tasks;
-DROP SEQUENCE tasks_id;
-EXIT;
diff --git a/activerecord/test/fixtures/db_definitions/drop_oracle_tables2.sql b/activerecord/test/fixtures/db_definitions/drop_oracle_tables2.sql
deleted file mode 100644
index 934026dc89..0000000000
--- a/activerecord/test/fixtures/db_definitions/drop_oracle_tables2.sql
+++ /dev/null
@@ -1,3 +0,0 @@
-DROP TABLE courses;
-DROP SEQUENCE courses_id;
-EXIT;
diff --git a/activerecord/test/fixtures/db_definitions/mysql.drop.sql b/activerecord/test/fixtures/db_definitions/mysql.drop.sql
index d8eaf9b29c..12103a6d9e 100644
--- a/activerecord/test/fixtures/db_definitions/mysql.drop.sql
+++ b/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;
diff --git a/activerecord/test/fixtures/db_definitions/mysql.sql b/activerecord/test/fixtures/db_definitions/mysql.sql
index 17d443bffd..865e67a83f 100755
--- a/activerecord/test/fixtures/db_definitions/mysql.sql
+++ b/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
+);
diff --git a/activerecord/test/fixtures/db_definitions/oci.drop.sql b/activerecord/test/fixtures/db_definitions/oci.drop.sql
index c98a8790c9..532a2f4208 100644
--- a/activerecord/test/fixtures/db_definitions/oci.drop.sql
+++ b/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;
diff --git a/activerecord/test/fixtures/db_definitions/oci.sql b/activerecord/test/fixtures/db_definitions/oci.sql
index bb4726caaa..29e6c76e2a 100644
--- a/activerecord/test/fixtures/db_definitions/oci.sql
+++ b/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
+);
diff --git a/activerecord/test/fixtures/db_definitions/postgresql.drop.sql b/activerecord/test/fixtures/db_definitions/postgresql.drop.sql
index e718d4b321..844bf3e2f8 100644
--- a/activerecord/test/fixtures/db_definitions/postgresql.drop.sql
+++ b/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;
 
diff --git a/activerecord/test/fixtures/db_definitions/postgresql.sql b/activerecord/test/fixtures/db_definitions/postgresql.sql
index 01d6ee288c..ae41f80788 100644
--- a/activerecord/test/fixtures/db_definitions/postgresql.sql
+++ b/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
+);
diff --git a/activerecord/test/fixtures/db_definitions/sqlite.drop.sql b/activerecord/test/fixtures/db_definitions/sqlite.drop.sql
index d8eaf9b29c..12103a6d9e 100644
--- a/activerecord/test/fixtures/db_definitions/sqlite.drop.sql
+++ b/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;
diff --git a/activerecord/test/fixtures/db_definitions/sqlite.sql b/activerecord/test/fixtures/db_definitions/sqlite.sql
index 0d05ca6932..2a26ff8c37 100644
--- a/activerecord/test/fixtures/db_definitions/sqlite.sql
+++ b/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
+);
diff --git a/activerecord/test/fixtures/db_definitions/sqlserver.drop.sql b/activerecord/test/fixtures/db_definitions/sqlserver.drop.sql
index 2ec7dbf026..efa98fe500 100644
--- a/activerecord/test/fixtures/db_definitions/sqlserver.drop.sql
+++ b/activerecord/test/fixtures/db_definitions/sqlserver.drop.sql
@@ -19,3 +19,5 @@ DROP TABLE posts;
 DROP TABLE comments;
 DROP TABLE authors;
 DROP TABLE tasks;
+DROP TABLE categories;
+DROP TABLE categories_posts;
diff --git a/activerecord/test/fixtures/db_definitions/sqlserver.sql b/activerecord/test/fixtures/db_definitions/sqlserver.sql
index 8e0d78bb61..3d93dffe20 100644
--- a/activerecord/test/fixtures/db_definitions/sqlserver.sql
+++ b/activerecord/test/fixtures/db_definitions/sqlserver.sql
@@ -141,3 +141,13 @@ CREATE TABLE tasks (
   starting datetime default NULL,
   ending datetime default NULL
 );
+
+CREATE TABLE categories (
+  id int NOT NULL IDENTITY(1, 1) PRIMARY KEY,
+  name varchar(255)
+);
+
+CREATE TABLE categories_posts (
+  category_id int NOT NULL,
+  post_id int int NOT NULL
+);
diff --git a/activerecord/test/fixtures/post.rb b/activerecord/test/fixtures/post.rb
index 3c0ef927dc..ce9abcc693 100644
--- a/activerecord/test/fixtures/post.rb
+++ b/activerecord/test/fixtures/post.rb
@@ -1,6 +1,8 @@
 class Post < ActiveRecord::Base
   belongs_to :author
   has_many   :comments
+  has_and_belongs_to_many :categories
 end
 
-class SpecialPost < Post; end
\ No newline at end of file
+class SpecialPost < Post
+end
\ No newline at end of file