From 9c39306860249aa9499c612d5f4f7f95dfa49a17 Mon Sep 17 00:00:00 2001
From: John Ferlito <johnf@inodes.org>
Date: Thu, 24 Mar 2011 18:37:02 +1100
Subject: [PATCH] Cache the connection we use to clean the DB with (Take 2)

The connection caching was a little bit too liberal and was caching
across databases (I think). No we just grab it when we enter clean.

Added an iterator so there isn't too much code duplication between
delete and truncate.
---
 lib/database_cleaner/active_record/deletion.rb |  6 ++----
 .../active_record/truncation.rb                | 18 ++++++++++--------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/lib/database_cleaner/active_record/deletion.rb b/lib/database_cleaner/active_record/deletion.rb
index feec7894..835f5056 100644
--- a/lib/database_cleaner/active_record/deletion.rb
+++ b/lib/database_cleaner/active_record/deletion.rb
@@ -53,10 +53,8 @@ module DatabaseCleaner::ActiveRecord
   class Deletion < Truncation
 
     def clean
-      connection.disable_referential_integrity do
-        tables_to_truncate.each do |table_name|
-          connection.delete_table table_name
-        end
+      each_table do |connection, table_name|
+        connection.delete_table table_name
       end
     end
 
diff --git a/lib/database_cleaner/active_record/truncation.rb b/lib/database_cleaner/active_record/truncation.rb
index 4cda30ee..a92eab44 100755
--- a/lib/database_cleaner/active_record/truncation.rb
+++ b/lib/database_cleaner/active_record/truncation.rb
@@ -103,24 +103,26 @@ class Truncation
     include ::DatabaseCleaner::Generic::Truncation
 
     def clean
+      each_table do |connection, table_name|
+        connection.truncate_table table_name
+      end
+    end
+
+    def each_table
+      connection = connection_klass.connection
       connection.disable_referential_integrity do
-        tables_to_truncate.each do |table_name|
-          connection.truncate_table table_name
+        tables_to_truncate(connection).each do |table_name|
+          yield connection, table_name
         end
       end
     end
 
     private
 
-    def tables_to_truncate
+    def tables_to_truncate(connection)
        (@only || connection.tables) - @tables_to_exclude - connection.views
     end
 
-    def connection
-       #::ActiveRecord::Base.connection
-       @connection ||= connection_klass.connection
-    end
-
     # overwritten
     def migration_storage_name
       'schema_migrations'