Merge pull request #6838 from kennyj/added_registration_task

Allow to register database tasks from different adapters

2 files changed, 27 insertions, 7 deletions

diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index f1241502f5..fb3dfc2730 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -3,13 +3,17 @@ module ActiveRecord
     module DatabaseTasks # :nodoc:
       extend self
 
-      TASKS_PATTERNS = {
-        /mysql/      => ActiveRecord::Tasks::MySQLDatabaseTasks,
-        /postgresql/ => ActiveRecord::Tasks::PostgreSQLDatabaseTasks,
-        /sqlite/     => ActiveRecord::Tasks::SQLiteDatabaseTasks
-      }
       LOCAL_HOSTS    = ['127.0.0.1', 'localhost']
 
+      def register_task(pattern, task)
+        @tasks ||= {}
+        @tasks[pattern] = task
+      end
+
+      register_task(/mysql/, ActiveRecord::Tasks::MySQLDatabaseTasks)
+      register_task(/postgresql/, ActiveRecord::Tasks::PostgreSQLDatabaseTasks)
+      register_task(/sqlite/, ActiveRecord::Tasks::SQLiteDatabaseTasks)
+
       def create(*arguments)
         configuration = arguments.first
         class_for_adapter(configuration['adapter']).new(*arguments).create
@@ -84,8 +88,8 @@ module ActiveRecord
       private
 
       def class_for_adapter(adapter)
-        key = TASKS_PATTERNS.keys.detect { |pattern| adapter[pattern] }
-        TASKS_PATTERNS[key]
+        key = @tasks.keys.detect { |pattern| adapter[pattern] }
+        @tasks[key]
       end
 
       def each_current_configuration(environment)
diff --git a/activerecord/test/cases/tasks/database_tasks_test.rb b/activerecord/test/cases/tasks/database_tasks_test.rb
index 8c96a8aaa1..4f3489b7a5 100644
--- a/activerecord/test/cases/tasks/database_tasks_test.rb
+++ b/activerecord/test/cases/tasks/database_tasks_test.rb
@@ -16,6 +16,22 @@ module ActiveRecord
      :postgresql => :postgresql_tasks,
      :sqlite3    => :sqlite_tasks
   }
+
+  class DatabaseTasksRegisterTask < ActiveRecord::TestCase
+    def test_register_task
+      klazz = Class.new do
+        def initialize(*arguments); end
+        def structure_dump(filename); end
+      end
+      instance = klazz.new
+
+      klazz.stubs(:new).returns instance
+      instance.expects(:structure_dump).with("awesome-file.sql")
+
+      ActiveRecord::Tasks::DatabaseTasks.register_task(/foo/, klazz)
+      ActiveRecord::Tasks::DatabaseTasks.structure_dump({'adapter' => :foo}, "awesome-file.sql")
+    end
+  end
  
   class DatabaseTasksCreateTest < ActiveRecord::TestCase
     include DatabaseTasksSetupper