Refactor db:structure:load task.

9 files changed, 129 insertions, 19 deletions

diff --git a/activerecord/lib/active_record/railties/databases.rake b/activerecord/lib/active_record/railties/databases.rake
index b8205a99a8..8199b5c2e0 100644
--- a/activerecord/lib/active_record/railties/databases.rake
+++ b/activerecord/lib/active_record/railties/databases.rake
@@ -303,18 +303,8 @@ db_namespace = namespace :db do
       abcs = ActiveRecord::Base.configurations
       filename = ENV['DB_STRUCTURE'] || File.join(Rails.root, "db", "structure.sql")
       case abcs[env]['adapter']
-      when /mysql/
-        ActiveRecord::Base.establish_connection(abcs[env])
-        ActiveRecord::Base.connection.execute('SET foreign_key_checks = 0')
-        IO.read(filename).split("\n\n").each do |table|
-          ActiveRecord::Base.connection.execute(table)
-        end
-      when /postgresql/
-        set_psql_env(abcs[env])
-        `psql -f "#{filename}" #{abcs[env]['database']}`
-      when /sqlite/
-        dbfile = abcs[env]['database']
-        `sqlite3 #{dbfile} < "#{filename}"`
+      when /mysql/, /postgresql/, /sqlite/
+        ActiveRecord::Tasks::DatabaseTasks.structure_load(abcs[Rails.env], filename)
       when 'sqlserver'
         `sqlcmd -S #{abcs[env]['host']} -d #{abcs[env]['database']} -U #{abcs[env]['username']} -P #{abcs[env]['password']} -i #{filename}`
       when 'oci', 'oracle'
@@ -462,10 +452,3 @@ end
 def firebird_db_string(config)
   FireRuby::Database.db_string_for(config.symbolize_keys)
 end
-
-def set_psql_env(config)
-  ENV['PGHOST']     = config['host']          if config['host']
-  ENV['PGPORT']     = config['port'].to_s     if config['port']
-  ENV['PGPASSWORD'] = config['password'].to_s if config['password']
-  ENV['PGUSER']     = config['username'].to_s if config['username']
-end
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index d53fe307c2..999b2ebc85 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -66,6 +66,12 @@ module ActiveRecord
         class_for_adapter(configuration['adapter']).new(*arguments).structure_dump(filename)
       end
 
+      def structure_load(*arguments)
+        configuration = arguments.first
+        filename = arguments.delete_at 1
+        class_for_adapter(configuration['adapter']).new(*arguments).structure_load(filename)
+      end
+
       private
 
       def class_for_adapter(adapter)
diff --git a/activerecord/lib/active_record/tasks/mysql_database_tasks.rb b/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
index 052ab60e81..b39cd2282f 100644
--- a/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
@@ -49,6 +49,14 @@ module ActiveRecord
         File.open(filename, "w:utf-8") { |f| f << ActiveRecord::Base.connection.structure_dump }
       end
 
+      def structure_load(filename)
+        establish_connection(configuration)
+        connection.execute('SET foreign_key_checks = 0')
+        IO.read(filename).split("\n\n").each do |table|
+          connection.execute(table)
+        end
+      end
+
       private
 
       def configuration
diff --git a/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb b/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
index dabc69b9cd..a210392e53 100644
--- a/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
@@ -48,6 +48,11 @@ module ActiveRecord
         File.open(filename, "a") { |f| f << "SET search_path TO #{ActiveRecord::Base.connection.schema_search_path};\n\n" }
       end
 
+      def structure_load(filename)
+        set_psql_env
+        Kernel.system("psql -f #{filename} #{configuration['database']}")
+      end
+
       private
 
       def configuration
diff --git a/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb b/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
index 68f0c304a8..da01058a82 100644
--- a/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
@@ -36,6 +36,11 @@ module ActiveRecord
         `sqlite3 #{dbfile} .schema > #{filename}`
       end
 
+      def structure_load(filename)
+        dbfile = configuration['database']
+        `sqlite3 #{dbfile} < "#{filename}"`
+      end
+
       private
 
       def configuration
diff --git a/activerecord/test/cases/database_tasks_test.rb b/activerecord/test/cases/database_tasks_test.rb
index 41267c2df3..943c6e39b2 100644
--- a/activerecord/test/cases/database_tasks_test.rb
+++ b/activerecord/test/cases/database_tasks_test.rb
@@ -362,4 +362,38 @@ module ActiveRecord
       ActiveRecord::Tasks::DatabaseTasks.structure_dump({'adapter' => 'sqlite3'}, "awesome-file.sql")
     end
   end
+
+  class DatabaseTasksStructureLoadTest < ActiveRecord::TestCase
+    def setup
+      @mysql_tasks, @postgresql_tasks, @sqlite_tasks = stub, stub, stub
+      ActiveRecord::Tasks::MySQLDatabaseTasks.stubs(:new).returns @mysql_tasks
+      ActiveRecord::Tasks::PostgreSQLDatabaseTasks.stubs(:new).
+        returns @postgresql_tasks
+      ActiveRecord::Tasks::SQLiteDatabaseTasks.stubs(:new).returns @sqlite_tasks
+    end
+
+    def test_mysql_structure_load
+      @mysql_tasks.expects(:structure_load).with("awesome-file.sql")
+
+      ActiveRecord::Tasks::DatabaseTasks.structure_load({'adapter' => 'mysql'}, "awesome-file.sql")
+    end
+
+    def test_mysql2_structure_load
+      @mysql_tasks.expects(:structure_load).with("awesome-file.sql")
+
+      ActiveRecord::Tasks::DatabaseTasks.structure_load({'adapter' => 'mysql2'}, "awesome-file.sql")
+    end
+
+    def test_postgresql_structure_load
+      @postgresql_tasks.expects(:structure_load).with("awesome-file.sql")
+
+      ActiveRecord::Tasks::DatabaseTasks.structure_load({'adapter' => 'postgresql'}, "awesome-file.sql")
+    end
+
+    def test_sqlite_structure_load
+      @sqlite_tasks.expects(:structure_load).with("awesome-file.sql")
+
+      ActiveRecord::Tasks::DatabaseTasks.structure_load({'adapter' => 'sqlite3'}, "awesome-file.sql")
+    end
+  end
 end
diff --git a/activerecord/test/cases/mysql_rake_test.rb b/activerecord/test/cases/mysql_rake_test.rb
index 7da63de282..42a11b0171 100644
--- a/activerecord/test/cases/mysql_rake_test.rb
+++ b/activerecord/test/cases/mysql_rake_test.rb
@@ -218,4 +218,29 @@ module ActiveRecord
       FileUtils.rm(filename)
     end
   end
+
+  class MySQLStructureLoadTest < ActiveRecord::TestCase
+    def setup
+      @connection    = stub
+      @configuration = {
+        'adapter'  => 'mysql',
+        'database' => 'test-db'
+      }
+
+      ActiveRecord::Base.stubs(:connection).returns(@connection)
+      ActiveRecord::Base.stubs(:establish_connection).returns(true)
+    end
+
+    def test_structure_load
+      filename = "awesome-file.sql"
+      ActiveRecord::Base.expects(:establish_connection).with(@configuration)
+      @connection.expects(:execute).twice
+
+      open(filename, 'w') { |f| f.puts("SELECT CURDATE();") }
+      ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
+    ensure
+      FileUtils.rm(filename)
+    end
+  end
+
 end
diff --git a/activerecord/test/cases/postgresql_rake_test.rb b/activerecord/test/cases/postgresql_rake_test.rb
index b055762a24..e8769bd4df 100644
--- a/activerecord/test/cases/postgresql_rake_test.rb
+++ b/activerecord/test/cases/postgresql_rake_test.rb
@@ -175,4 +175,26 @@ module ActiveRecord
       FileUtils.rm(filename)
     end
   end
+
+  class PostgreSQLStructureLoadTest < ActiveRecord::TestCase
+    def setup
+      @connection    = stub
+      @configuration = {
+        'adapter'  => 'postgresql',
+        'database' => 'my-app-db'
+      }
+
+      ActiveRecord::Base.stubs(:connection).returns(@connection)
+      ActiveRecord::Base.stubs(:establish_connection).returns(true)
+      Kernel.stubs(:system)
+    end
+
+    def test_structure_dump
+      filename = "awesome-file.sql"
+      Kernel.expects(:system).with("psql -f #{filename} my-app-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
+    end
+  end
+
 end
diff --git a/activerecord/test/cases/sqlite_rake_test.rb b/activerecord/test/cases/sqlite_rake_test.rb
index 69fb129326..b5557fc953 100644
--- a/activerecord/test/cases/sqlite_rake_test.rb
+++ b/activerecord/test/cases/sqlite_rake_test.rb
@@ -145,4 +145,26 @@ module ActiveRecord
       FileUtils.rm(dbfile)
     end
   end
+
+  class SqliteStructureLoadTest < ActiveRecord::TestCase
+    def setup
+      @database      = "db_create.sqlite3"
+      @configuration = {
+        'adapter'  => 'sqlite3',
+        'database' => @database
+      }
+    end
+
+    def test_structure_load
+      dbfile   = @database
+      filename = "awesome-file.sql"
+
+      open(filename, 'w') { |f| f.puts("select datetime('now', 'localtime');") }
+      ActiveRecord::Tasks::DatabaseTasks.structure_load @configuration, filename, '/rails/root'
+      assert File.exists?(dbfile)
+    ensure
+      FileUtils.rm(filename)
+      FileUtils.rm(dbfile)
+    end
+  end
 end