Merge pull request #27437 from kirs/structure-load-dump-flags

Make ActiveRecord structure load/dump configurable

activerecord/lib/active_record/tasks/database_tasks.rb

@@ -35,6 +35,16 @@ module ActiveRecord
     #
     #   DatabaseTasks.create_current('production')
     module DatabaseTasks
+      ##
+      # :singleton-method:
+      # Extra flags passed to database CLI tool (mysqldump/pg_dump) when calling db:structure:dump
+      mattr_accessor :structure_dump_flags, instance_accessor: false
+
+      ##
+      # :singleton-method:
+      # Extra flags passed to database CLI tool when calling db:structure:load
+      mattr_accessor :structure_load_flags, instance_accessor: false
+
       extend self
 
       attr_writer :current_config, :db_dir, :migrations_paths, :fixtures_path, :root, :env, :seed_loader
@@ -204,13 +214,13 @@ module ActiveRecord
       def structure_dump(*arguments)
         configuration = arguments.first
         filename = arguments.delete_at 1
-        class_for_adapter(configuration["adapter"]).new(*arguments).structure_dump(filename)
+        class_for_adapter(configuration["adapter"]).new(*arguments).structure_dump(filename, structure_dump_flags)
       end
 
       def structure_load(*arguments)
         configuration = arguments.first
         filename = arguments.delete_at 1
-        class_for_adapter(configuration["adapter"]).new(*arguments).structure_load(filename)
+        class_for_adapter(configuration["adapter"]).new(*arguments).structure_load(filename, structure_load_flags)
       end
 
       def load_schema(configuration, format = ActiveRecord::Base.schema_format, file = nil) # :nodoc:

activerecord/lib/active_record/tasks/mysql_database_tasks.rb

@@ -53,21 +53,23 @@ module ActiveRecord
         connection.collation
       end
 
-      def structure_dump(filename)
+      def structure_dump(filename, extra_flags)
         args = prepare_command_options
         args.concat(["--result-file", "#{filename}"])
         args.concat(["--no-data"])
         args.concat(["--routines"])
         args.concat(["--skip-comments"])
+        args.concat(Array(extra_flags)) if extra_flags
         args.concat(["#{configuration['database']}"])
 
         run_cmd("mysqldump", args, "dumping")
       end
 
-      def structure_load(filename)
+      def structure_load(filename, extra_flags)
         args = prepare_command_options
         args.concat(["--execute", %{SET FOREIGN_KEY_CHECKS = 0; SOURCE #{filename}; SET FOREIGN_KEY_CHECKS = 1}])
         args.concat(["--database", "#{configuration['database']}"])
+        args.concat(Array(extra_flags)) if extra_flags
 
         run_cmd("mysql", args, "loading")
       end

activerecord/lib/active_record/tasks/postgresql_database_tasks.rb

@@ -43,7 +43,7 @@ module ActiveRecord
         create true
       end
 
-      def structure_dump(filename)
+      def structure_dump(filename, extra_flags)
         set_psql_env
 
         search_path = \
@@ -57,6 +57,7 @@ module ActiveRecord
           end
 
         args = ["-s", "-x", "-O", "-f", filename]
+        args.concat(Array(extra_flags)) if extra_flags
         unless search_path.blank?
           args += search_path.split(",").map do |part|
             "--schema=#{part.strip}"
@@ -67,9 +68,11 @@ module ActiveRecord
         File.open(filename, "a") { |f| f << "SET search_path TO #{connection.schema_search_path};\n\n" }
       end
 
-      def structure_load(filename)
+      def structure_load(filename, extra_flags)
         set_psql_env
-        args = [ "-v", ON_ERROR_STOP_1, "-q", "-f", filename, configuration["database"] ]
+        args = ["-v", ON_ERROR_STOP_1, "-q", "-f", filename]
+        args.concat(Array(extra_flags)) if extra_flags
+        args << configuration["database"]
         run_cmd("psql", args, "loading")
       end
 

activerecord/lib/active_record/tasks/sqlite_database_tasks.rb

@@ -35,14 +35,16 @@ module ActiveRecord
         connection.encoding
       end
 
-      def structure_dump(filename)
+      def structure_dump(filename, extra_flags)
         dbfile = configuration["database"]
-        `sqlite3 #{dbfile} .schema > #{filename}`
+        flags = extra_flags.try!(:join, " ")
+        `sqlite3 #{flags} #{dbfile} .schema > #{filename}`
       end
 
-      def structure_load(filename)
+      def structure_load(filename, extra_flags)
         dbfile = configuration["database"]
-        `sqlite3 #{dbfile} < "#{filename}"`
+        flags = extra_flags.try!(:join, " ")
+        `sqlite3 #{flags} #{dbfile} < "#{filename}"`
       end
 
       private

activerecord/test/cases/tasks/database_tasks_test.rb

@@ -61,7 +61,7 @@ module ActiveRecord
       instance = klazz.new
 
       klazz.stubs(:new).returns instance
-      instance.expects(:structure_dump).with("awesome-file.sql")
+      instance.expects(:structure_dump).with("awesome-file.sql", nil)
 
       ActiveRecord::Tasks::DatabaseTasks.register_task(/foo/, klazz)
       ActiveRecord::Tasks::DatabaseTasks.structure_dump({ "adapter" => :foo }, "awesome-file.sql")
@@ -411,7 +411,7 @@ module ActiveRecord
 
     ADAPTERS_TASKS.each do |k, v|
       define_method("test_#{k}_structure_dump") do
-        eval("@#{v}").expects(:structure_dump).with("awesome-file.sql")
+        eval("@#{v}").expects(:structure_dump).with("awesome-file.sql", nil)
         ActiveRecord::Tasks::DatabaseTasks.structure_dump({ "adapter" => k }, "awesome-file.sql")
       end
     end
@@ -422,7 +422,7 @@ module ActiveRecord
 
     ADAPTERS_TASKS.each do |k, v|
       define_method("test_#{k}_structure_load") do
-        eval("@#{v}").expects(:structure_load).with("awesome-file.sql")
+        eval("@#{v}").expects(:structure_load).with("awesome-file.sql", nil)
         ActiveRecord::Tasks::DatabaseTasks.structure_load({ "adapter" => k }, "awesome-file.sql")
       end
     end

activerecord/test/cases/tasks/mysql_rake_test.rb

@@ -294,6 +294,17 @@ if current_adapter?(:Mysql2Adapter)
         ActiveRecord::Tasks::DatabaseTasks.structure_dump(@configuration, filename)
       end
 
+      def test_structure_dump_with_extra_flags
+        filename = "awesome-file.sql"
+        expected_command = ["mysqldump", "--result-file", filename, "--no-data", "--routines", "--skip-comments", "--noop", "test-db"]
+
+        assert_called_with(Kernel, :system, expected_command, returns: true) do
+          with_structure_dump_flags(["--noop"]) do
+            ActiveRecord::Tasks::DatabaseTasks.structure_dump(@configuration, filename)
+          end
+        end
+      end
+
       def test_warn_when_external_structure_dump_command_execution_fails
         filename = "awesome-file.sql"
         Kernel.expects(:system)
@@ -323,6 +334,15 @@ if current_adapter?(:Mysql2Adapter)
           @configuration.merge("sslca" => "ca.crt"),
           filename)
       end
+
+      private
+        def with_structure_dump_flags(flags)
+          old = ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags
+          ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags = flags
+          yield
+        ensure
+          ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags = old
+        end
     end
 
     class MySQLStructureLoadTest < ActiveRecord::TestCase
@@ -340,6 +360,26 @@ if current_adapter?(:Mysql2Adapter)
 
         ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
       end
+
+      def test_structure_load
+        filename = "awesome-file.sql"
+        expected_command = ["mysql", "--execute", %{SET FOREIGN_KEY_CHECKS = 0; SOURCE #{filename}; SET FOREIGN_KEY_CHECKS = 1}, "--database", "test-db", "--noop"]
+
+        assert_called_with(Kernel, :system, expected_command, returns: true) do
+          with_structure_load_flags(["--noop"]) do
+            ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
+          end
+        end
+      end
+
+      private
+        def with_structure_load_flags(flags)
+          old = ActiveRecord::Tasks::DatabaseTasks.structure_load_flags
+          ActiveRecord::Tasks::DatabaseTasks.structure_load_flags = flags
+          yield
+        ensure
+          ActiveRecord::Tasks::DatabaseTasks.structure_load_flags = old
+        end
     end
   end
 end

activerecord/test/cases/tasks/postgresql_rake_test.rb

@@ -236,6 +236,16 @@ if current_adapter?(:PostgreSQLAdapter)
         ActiveRecord::Tasks::DatabaseTasks.structure_dump(@configuration, @filename)
       end
 
+      def test_structure_dump_with_extra_flags
+        expected_command = ["pg_dump", "-s", "-x", "-O", "-f", @filename, "--noop", "my-app-db"]
+
+        assert_called_with(Kernel, :system, expected_command, returns: true) do
+          with_structure_dump_flags(["--noop"]) do
+            ActiveRecord::Tasks::DatabaseTasks.structure_dump(@configuration, @filename)
+          end
+        end
+      end
+
       def test_structure_dump_with_schema_search_path
         @configuration["schema_search_path"] = "foo,bar"
 
@@ -263,7 +273,6 @@ if current_adapter?(:PostgreSQLAdapter)
       end
 
       private
-
         def with_dump_schemas(value, &block)
           old_dump_schemas = ActiveRecord::Base.dump_schemas
           ActiveRecord::Base.dump_schemas = value
@@ -271,6 +280,14 @@ if current_adapter?(:PostgreSQLAdapter)
         ensure
           ActiveRecord::Base.dump_schemas = old_dump_schemas
         end
+
+        def with_structure_dump_flags(flags)
+          old = ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags
+          ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags = flags
+          yield
+        ensure
+          ActiveRecord::Tasks::DatabaseTasks.structure_dump_flags = old
+        end
     end
 
     class PostgreSQLStructureLoadTest < ActiveRecord::TestCase
@@ -293,12 +310,32 @@ if current_adapter?(:PostgreSQLAdapter)
         ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
       end
 
+      def test_structure_load_with_extra_flags
+        filename = "awesome-file.sql"
+        expected_command = ["psql", "-v", "ON_ERROR_STOP=1", "-q", "-f", filename, "--noop", @configuration["database"]]
+
+        assert_called_with(Kernel, :system, expected_command, returns: true) do
+          with_structure_load_flags(["--noop"]) do
+            ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
+          end
+        end
+      end
+
       def test_structure_load_accepts_path_with_spaces
         filename = "awesome file.sql"
         Kernel.expects(:system).with("psql", "-v", "ON_ERROR_STOP=1", "-q", "-f", filename, @configuration["database"]).returns(true)
 
         ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
       end
+
+      private
+        def with_structure_load_flags(flags)
+          old = ActiveRecord::Tasks::DatabaseTasks.structure_load_flags
+          ActiveRecord::Tasks::DatabaseTasks.structure_load_flags = flags
+          yield
+        ensure
+          ActiveRecord::Tasks::DatabaseTasks.structure_load_flags = old
+        end
     end
   end
 end