Merge pull request #31422 from Edouard-chin/multistatement-fixtures

Build a multi-statement query when inserting fixtures

3 files changed, 114 insertions, 21 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
index 36048bee03..d663b59444 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
@@ -356,35 +356,33 @@ module ActiveRecord
       # Inserts a set of fixtures into the table. Overridden in adapters that require
       # something beyond a simple insert (eg. Oracle).
       def insert_fixtures(fixtures, table_name)
+        ActiveSupport::Deprecation.warn(<<-MSG.squish)
+          `insert_fixtures` is deprecated and will be removed in the next version of Rails.
+          Consider using `insert_fixtures_set` for performance improvement.
+        MSG
         return if fixtures.empty?
 
-        columns = schema_cache.columns_hash(table_name)
+        execute(build_fixture_sql(fixtures, table_name), "Fixtures Insert")
+      end
 
-        values = fixtures.map do |fixture|
-          fixture = fixture.stringify_keys
+      def insert_fixtures_set(fixture_set, tables_to_delete = [])
+        fixture_inserts = fixture_set.map do |table_name, fixtures|
+          next if fixtures.empty?
 
-          unknown_columns = fixture.keys - columns.keys
-          if unknown_columns.any?
-            raise Fixture::FixtureError, %(table "#{table_name}" has no columns named #{unknown_columns.map(&:inspect).join(', ')}.)
-          end
+          build_fixture_sql(fixtures, table_name)
+        end.compact
 
-          columns.map do |name, column|
-            if fixture.key?(name)
-              type = lookup_cast_type_from_column(column)
-              bind = Relation::QueryAttribute.new(name, fixture[name], type)
-              with_yaml_fallback(bind.value_for_database)
-            else
-              Arel.sql("DEFAULT")
+        table_deletes = tables_to_delete.map { |table| "DELETE FROM #{quote_table_name table}".dup }
+        total_sql = Array.wrap(combine_multi_statements(table_deletes + fixture_inserts))
+
+        disable_referential_integrity do
+          transaction(requires_new: true) do
+            total_sql.each do |sql|
+              execute sql, "Fixtures Load"
+              yield if block_given?
             end
           end
         end
-
-        table = Arel::Table.new(table_name)
-        manager = Arel::InsertManager.new
-        manager.into(table)
-        columns.each_key { |column| manager.columns << table[column] }
-        manager.values = manager.create_values_list(values)
-        execute manager.to_sql, "Fixtures Insert"
       end
 
       def empty_insert_statement_value
@@ -417,6 +415,41 @@ module ActiveRecord
 
       private
 
+        def build_fixture_sql(fixtures, table_name)
+          columns = schema_cache.columns_hash(table_name)
+
+          values = fixtures.map do |fixture|
+            fixture = fixture.stringify_keys
+
+            unknown_columns = fixture.keys - columns.keys
+            if unknown_columns.any?
+              raise Fixture::FixtureError, %(table "#{table_name}" has no columns named #{unknown_columns.map(&:inspect).join(', ')}.)
+            end
+
+            columns.map do |name, column|
+              if fixture.key?(name)
+                type = lookup_cast_type_from_column(column)
+                bind = Relation::QueryAttribute.new(name, fixture[name], type)
+                with_yaml_fallback(bind.value_for_database)
+              else
+                Arel.sql("DEFAULT")
+              end
+            end
+          end
+
+          table = Arel::Table.new(table_name)
+          manager = Arel::InsertManager.new
+          manager.into(table)
+          columns.each_key { |column| manager.columns << table[column] }
+          manager.values = manager.create_values_list(values)
+
+          manager.to_sql
+        end
+
+        def combine_multi_statements(total_sql)
+          total_sql.join(";\n")
+        end
+
         # Returns a subquery for the given key using the join information.
         def subquery_for(key, select)
           subselect = select.clone
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
index b394a8969f..d1a3b6de40 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -530,8 +530,56 @@ module ActiveRecord
         without_sql_mode("NO_AUTO_VALUE_ON_ZERO") { super }
       end
 
+      def insert_fixtures_set(fixture_set, tables_to_delete = [])
+        iterate_over_results = -> { while raw_connection.next_result; end; }
+
+        with_multi_statements do
+          without_sql_mode("NO_AUTO_VALUE_ON_ZERO") do
+            super(fixture_set, tables_to_delete, &iterate_over_results)
+          end
+        end
+      end
+
       private
 
+        def combine_multi_statements(total_sql)
+          total_sql.each_with_object([]) do |sql, total_sql_chunks|
+            previous_packet = total_sql_chunks.last
+            sql << ";\n"
+            if max_allowed_packet_reached?(sql, previous_packet) || total_sql_chunks.empty?
+              total_sql_chunks << sql
+            else
+              previous_packet << sql
+            end
+          end
+        end
+
+        def max_allowed_packet_reached?(current_packet, previous_packet)
+          if current_packet.bytesize > max_allowed_packet
+            raise ActiveRecordError, "Fixtures set is too large #{current_packet.bytesize}. Consider increasing the max_allowed_packet variable."
+          elsif previous_packet.nil?
+            false
+          else
+            (current_packet.bytesize + previous_packet.bytesize) > max_allowed_packet
+          end
+        end
+
+        def max_allowed_packet
+          bytes_margin = 2
+          @max_allowed_packet ||= (show_variable("max_allowed_packet") - bytes_margin)
+        end
+
+        def with_multi_statements
+          previous_flags = @config[:flags]
+          @config[:flags] = Mysql2::Client::MULTI_STATEMENTS
+          reconnect!
+
+          yield
+        ensure
+          @config[:flags] = previous_flags
+          reconnect!
+        end
+
         def without_sql_mode(mode)
           result = execute("SELECT @@SESSION.sql_mode")
           current_mode = result.first[0]
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
index 9c69d8f13b..c66cada07a 100644
--- a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
@@ -372,6 +372,18 @@ module ActiveRecord
         end
       end
 
+      def insert_fixtures_set(fixture_set, tables_to_delete = [])
+        disable_referential_integrity do
+          transaction(requires_new: true) do
+            tables_to_delete.each { |table| delete "DELETE FROM #{quote_table_name(table)}", "Fixture Delete" }
+
+            fixture_set.each do |table_name, rows|
+              rows.each { |row| insert_fixture(row, table_name) }
+            end
+          end
+        end
+      end
+
       private
         def initialize_type_map(m = type_map)
           super