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

Build a multi-statement query when inserting fixtures

4 files changed, 139 insertions, 55 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
diff --git a/activerecord/lib/active_record/fixtures.rb b/activerecord/lib/active_record/fixtures.rb
index 86f13d75d5..896d51c0fe 100644
--- a/activerecord/lib/active_record/fixtures.rb
+++ b/activerecord/lib/active_record/fixtures.rb
@@ -540,47 +540,38 @@ module ActiveRecord
       }
 
       unless files_to_read.empty?
-        connection.disable_referential_integrity do
-          fixtures_map = {}
-
-          fixture_sets = files_to_read.map do |fs_name|
-            klass = class_names[fs_name]
-            conn = klass ? klass.connection : connection
-            fixtures_map[fs_name] = new( # ActiveRecord::FixtureSet.new
-              conn,
-              fs_name,
-              klass,
-              ::File.join(fixtures_directory, fs_name))
-          end
-
-          update_all_loaded_fixtures fixtures_map
-
-          connection.transaction(requires_new: true) do
-            deleted_tables = Hash.new { |h, k| h[k] = Set.new }
-            fixture_sets.each do |fs|
-              conn = fs.model_class.respond_to?(:connection) ? fs.model_class.connection : connection
-              table_rows = fs.table_rows
+        fixtures_map = {}
+
+        fixture_sets = files_to_read.map do |fs_name|
+          klass = class_names[fs_name]
+          conn = klass ? klass.connection : connection
+          fixtures_map[fs_name] = new( # ActiveRecord::FixtureSet.new
+            conn,
+            fs_name,
+            klass,
+            ::File.join(fixtures_directory, fs_name))
+        end
 
-              table_rows.each_key do |table|
-                unless deleted_tables[conn].include? table
-                  conn.delete "DELETE FROM #{conn.quote_table_name(table)}", "Fixture Delete"
-                end
-                deleted_tables[conn] << table
-              end
+        update_all_loaded_fixtures fixtures_map
+        fixture_sets_by_connection = fixture_sets.group_by { |fs| fs.model_class ? fs.model_class.connection : connection }
 
-              table_rows.each do |fixture_set_name, rows|
-                conn.insert_fixtures(rows, fixture_set_name)
-              end
+        fixture_sets_by_connection.each do |conn, set|
+          table_rows_for_connection = Hash.new { |h, k| h[k] = [] }
 
-              # Cap primary key sequences to max(pk).
-              if conn.respond_to?(:reset_pk_sequence!)
-                conn.reset_pk_sequence!(fs.table_name)
-              end
+          set.each do |fs|
+            fs.table_rows.each do |table, rows|
+              table_rows_for_connection[table].unshift(*rows)
             end
           end
+          conn.insert_fixtures_set(table_rows_for_connection, table_rows_for_connection.keys)
 
-          cache_fixtures(connection, fixtures_map)
+          # Cap primary key sequences to max(pk).
+          if conn.respond_to?(:reset_pk_sequence!)
+            set.each { |fs| conn.reset_pk_sequence!(fs.table_name) }
+          end
         end
+
+        cache_fixtures(connection, fixtures_map)
       end
       cached_fixtures(connection, fixture_set_names)
     end