From 50c0f33bda33008379f3122191de27b8423ef08f Mon Sep 17 00:00:00 2001 From: Cody Cutrer Date: Wed, 26 Jun 2019 14:27:33 -0600 Subject: [PATCH] allow temp tables in deploy env also ensure we unshackle the db server automatically Change-Id: I709c92b0568783fd0b77b5b0518a8bf8783a1d74 Reviewed-on: https://gerrit.instructure.com/199223 Tested-by: Jenkins Reviewed-by: Rob Orton QA-Review: Rob Orton Product-Review: Rob Orton --- config/initializers/active_record.rb | 162 ++++++++++++++------------- 1 file changed, 83 insertions(+), 79 deletions(-) diff --git a/config/initializers/active_record.rb b/config/initializers/active_record.rb index 7e4cee6c7d1..7f11cf47804 100644 --- a/config/initializers/active_record.rb +++ b/config/initializers/active_record.rb @@ -860,97 +860,101 @@ ActiveRecord::Relation.class_eval do end def find_in_batches_with_temp_table(options = {}) - can_do_it = Rails.env.production? || - ActiveRecord::Base.in_migration || - (!Rails.env.test? && connection.open_transactions > 0) || - ActiveRecord::Base.in_transaction_in_test? - raise "find_in_batches_with_temp_table probably won't work outside a migration - and outside a transaction. Unfortunately, it's impossible to automatically - determine a better way to do it that will work correctly. You can try - switching to slave first (then switching to master if you modify anything - inside your loop), wrapping in a transaction (but be wary of locking records - for the duration of your query if you do any writes in your loop), or not - forcing find_in_batches to use a temp table (avoiding custom selects, - group, or order)." unless can_do_it + Shard.current.database_server.unshackle do + can_do_it = Rails.env.production? || + ActiveRecord::Base.in_migration || + Shackles.environment == :deploy || + (!Rails.env.test? && connection.open_transactions > 0) || + ActiveRecord::Base.in_transaction_in_test? + raise "find_in_batches_with_temp_table probably won't work outside a migration + and outside a transaction. Unfortunately, it's impossible to automatically + determine a better way to do it that will work correctly. You can try + switching to slave first (then switching to master if you modify anything + inside your loop), wrapping in a transaction (but be wary of locking records + for the duration of your query if you do any writes in your loop), or not + forcing find_in_batches to use a temp table (avoiding custom selects, + group, or order)." unless can_do_it - if options[:pluck] - pluck = Array(options[:pluck]) - pluck_for_select = pluck.map do |column_name| - if column_name.is_a?(Symbol) && column_names.include?(column_name.to_s) - "#{connection.quote_local_table_name(table_name)}.#{connection.quote_column_name(column_name)}" - else - column_name.to_s - end - end - pluck = pluck.map(&:to_s) - end - batch_size = options[:batch_size] || 1000 - if pluck - sql = select(pluck_for_select).to_sql - else - sql = to_sql - end - table = "#{table_name}_find_in_batches_temp_table_#{sql.hash.abs.to_s(36)}" - table = table[-63..-1] if table.length > 63 - connection.execute "CREATE TEMPORARY TABLE #{table} AS #{sql}" - begin - index = "temp_primary_key" - case connection.adapter_name - when 'PostgreSQL' - begin - old_proc = connection.raw_connection.set_notice_processor {} - if pluck && pluck.any?{|p| p == primary_key.to_s} - connection.execute("CREATE INDEX #{connection.quote_local_table_name(index)} ON #{connection.quote_local_table_name(table)}(#{connection.quote_column_name(primary_key)})") - index = primary_key.to_s - else - pluck.unshift(index) if pluck - connection.execute "ALTER TABLE #{table} - ADD temp_primary_key SERIAL PRIMARY KEY" - end - ensure - connection.raw_connection.set_notice_processor(&old_proc) if old_proc + if options[:pluck] + pluck = Array(options[:pluck]) + pluck_for_select = pluck.map do |column_name| + if column_name.is_a?(Symbol) && column_names.include?(column_name.to_s) + "#{connection.quote_local_table_name(table_name)}.#{connection.quote_column_name(column_name)}" + else + column_name.to_s end - else - raise "Temp tables not supported!" + end + pluck = pluck.map(&:to_s) end + batch_size = options[:batch_size] || 1000 + if pluck + sql = select(pluck_for_select).to_sql + else + sql = to_sql + end + table = "#{table_name}_find_in_batches_temp_table_#{sql.hash.abs.to_s(36)}" + table = table[-63..-1] if table.length > 63 - includes = includes_values + preload_values - klass.unscoped do - - quoted_plucks = pluck && pluck.map do |column_name| - # Rails 4.2 is going to try to quote them anyway but unfortunately not to the temp table, so just make it explicit - column_names.include?(column_name) ? - Arel.sql("#{connection.quote_local_table_name(table)}.#{connection.quote_column_name(column_name)}") : column_name + connection.execute "CREATE TEMPORARY TABLE #{table} AS #{sql}" + begin + index = "temp_primary_key" + case connection.adapter_name + when 'PostgreSQL' + begin + old_proc = connection.raw_connection.set_notice_processor {} + if pluck && pluck.any?{|p| p == primary_key.to_s} + connection.execute("CREATE INDEX #{connection.quote_local_table_name(index)} ON #{connection.quote_local_table_name(table)}(#{connection.quote_column_name(primary_key)})") + index = primary_key.to_s + else + pluck.unshift(index) if pluck + connection.execute "ALTER TABLE #{table} + ADD temp_primary_key SERIAL PRIMARY KEY" + end + ensure + connection.raw_connection.set_notice_processor(&old_proc) if old_proc + end + else + raise "Temp tables not supported!" end - if pluck - batch = klass.from(table).order(Arel.sql(index)).limit(batch_size).pluck(*quoted_plucks) - else - sql = "SELECT * FROM #{table} ORDER BY #{index} LIMIT #{batch_size}" - batch = klass.find_by_sql(sql) - end - while !batch.empty? - ActiveRecord::Associations::Preloader.new.preload(batch, includes) if includes - yield batch - break if batch.size < batch_size + includes = includes_values + preload_values + klass.unscoped do + + quoted_plucks = pluck && pluck.map do |column_name| + # Rails 4.2 is going to try to quote them anyway but unfortunately not to the temp table, so just make it explicit + column_names.include?(column_name) ? + Arel.sql("#{connection.quote_local_table_name(table)}.#{connection.quote_column_name(column_name)}") : column_name + end if pluck - last_value = pluck.length == 1 ? batch.last : batch.last[pluck.index(index)] - batch = klass.from(table).order(Arel.sql(index)).where("#{index} > ?", last_value).limit(batch_size).pluck(*quoted_plucks) + batch = klass.from(table).order(Arel.sql(index)).limit(batch_size).pluck(*quoted_plucks) else - last_value = batch.last[index] - sql = "SELECT * - FROM #{table} - WHERE #{index} > #{last_value} - ORDER BY #{index} ASC - LIMIT #{batch_size}" + sql = "SELECT * FROM #{table} ORDER BY #{index} LIMIT #{batch_size}" batch = klass.find_by_sql(sql) end + while !batch.empty? + ActiveRecord::Associations::Preloader.new.preload(batch, includes) if includes + yield batch + break if batch.size < batch_size + + if pluck + last_value = pluck.length == 1 ? batch.last : batch.last[pluck.index(index)] + batch = klass.from(table).order(Arel.sql(index)).where("#{index} > ?", last_value).limit(batch_size).pluck(*quoted_plucks) + else + last_value = batch.last[index] + sql = "SELECT * + FROM #{table} + WHERE #{index} > #{last_value} + ORDER BY #{index} ASC + LIMIT #{batch_size}" + batch = klass.find_by_sql(sql) + end + end + end + ensure + if !$!.is_a?(ActiveRecord::StatementInvalid) || connection.open_transactions == 0 + connection.execute "DROP TABLE #{table}" end - end - ensure - if !$!.is_a?(ActiveRecord::StatementInvalid) || connection.open_transactions == 0 - connection.execute "DROP TABLE #{table}" end end end