From 8daf4fb0bc6447b65be3060084e40f7a7b7998d1 Mon Sep 17 00:00:00 2001 From: Ulysse Buonomo Date: Sun, 5 Jan 2025 20:41:47 +0100 Subject: [PATCH] chore: faster referential integrity for tests --- CONTRIBUTING.md | 7 +- .../cockroachdb/referential_integrity.rb | 82 +++++++++++++++++-- .../cockroachdb/schema_statements.rb | 26 ------ test/cases/fixtures_test.rb | 2 +- test/cases/helper_cockroachdb.rb | 20 +++-- test/support/sql_logger.rb | 55 +++++++++++++ 6 files changed, 147 insertions(+), 45 deletions(-) create mode 100644 test/support/sql_logger.rb diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 63e744c3..58f8bf71 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -115,9 +115,10 @@ This section intent to help you with a checklist. - Verify the written text at the beginning of the test suite, there are likely some changes in excluded tests. - Check for some important methods, some will change for sure: - - [x] `def new_column_from_field(` - - [x] `def column_definitions(` - - [x] `def pk_and_sequence_for(` + - [ ] `def new_column_from_field(` + - [ ] `def column_definitions(` + - [ ] `def pk_and_sequence_for(` + - [ ] `def foreign_keys(` and `def all_foreign_keys(` - [ ] ... - Check for setups containing `drop_table` in the test suite. Especially if you have tons of failure, this is likely the cause. diff --git a/lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb b/lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb index 149e9ddf..bdbf3852 100644 --- a/lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb +++ b/lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb @@ -34,10 +34,16 @@ def check_all_foreign_keys_valid! end def disable_referential_integrity - foreign_keys = tables.map { |table| foreign_keys(table) }.flatten + foreign_keys = all_foreign_keys foreign_keys.each do |foreign_key| - remove_foreign_key(foreign_key.from_table, name: foreign_key.options[:name]) + # We do not use the `#remove_foreign_key` method here because it + # checks for foreign keys existance in the schema cache. This method + # is performance critical and we know the foreign key exist. + at = create_alter_table foreign_key.from_table + at.drop_foreign_key foreign_key.name + + execute schema_creation.accept(at) end yield @@ -52,11 +58,16 @@ def disable_referential_integrity ActiveRecord::Base.table_name_suffix = "" begin + # Avoid having PG:DuplicateObject error if a test is ran in transaction. + # TODO: verify that there is no cache issue related to running this (e.g: fk + # still in cache but not in db) + # + # We avoid using `foreign_key_exists?` here because it checks the schema cache + # for every key. This method is performance critical for the test suite, hence + # we use the `#all_foreign_keys` method that only make one query to the database. + already_inserted_foreign_keys = all_foreign_keys foreign_keys.each do |foreign_key| - # Avoid having PG:DuplicateObject error if a test is ran in transaction. - # TODO: verify that there is no cache issue related to running this (e.g: fk - # still in cache but not in db) - next if foreign_key_exists?(foreign_key.from_table, name: foreign_key.options[:name]) + next if already_inserted_foreign_keys.any? { |fk| fk.from_table == foreign_key.from_table && fk.options[:name] == foreign_key.options[:name] } add_foreign_key(foreign_key.from_table, foreign_key.to_table, **foreign_key.options) end @@ -65,6 +76,65 @@ def disable_referential_integrity ActiveRecord::Base.table_name_suffix = old_suffix end end + + private + + # Copy/paste of the `#foreign_keys(table)` method adapted to return every single + # foreign key in the database. + def all_foreign_keys + fk_info = internal_exec_query(<<~SQL, "SCHEMA") + SELECT CASE + WHEN n1.nspname = current_schema() + THEN '' + ELSE n1.nspname || '.' + END || t1.relname AS from_table, + CASE + WHEN n2.nspname = current_schema() + THEN '' + ELSE n2.nspname || '.' + END || t2.relname AS to_table, + a1.attname AS column, a2.attname AS primary_key, c.conname AS name, c.confupdtype AS on_update, c.confdeltype AS on_delete, c.convalidated AS valid, c.condeferrable AS deferrable, c.condeferred AS deferred, + c.conkey, c.confkey, c.conrelid, c.confrelid + FROM pg_constraint c + JOIN pg_class t1 ON c.conrelid = t1.oid + JOIN pg_class t2 ON c.confrelid = t2.oid + JOIN pg_attribute a1 ON a1.attnum = c.conkey[1] AND a1.attrelid = t1.oid + JOIN pg_attribute a2 ON a2.attnum = c.confkey[1] AND a2.attrelid = t2.oid + JOIN pg_namespace t3 ON c.connamespace = t3.oid + JOIN pg_namespace n1 ON t1.relnamespace = n1.oid + JOIN pg_namespace n2 ON t2.relnamespace = n2.oid + WHERE c.contype = 'f' + ORDER BY c.conname + SQL + + fk_info.map do |row| + from_table = PostgreSQL::Utils.unquote_identifier(row["from_table"]) + to_table = PostgreSQL::Utils.unquote_identifier(row["to_table"]) + conkey = row["conkey"].scan(/\d+/).map(&:to_i) + confkey = row["confkey"].scan(/\d+/).map(&:to_i) + + if conkey.size > 1 + column = column_names_from_column_numbers(row["conrelid"], conkey) + primary_key = column_names_from_column_numbers(row["confrelid"], confkey) + else + column = PostgreSQL::Utils.unquote_identifier(row["column"]) + primary_key = row["primary_key"] + end + + options = { + column: column, + name: row["name"], + primary_key: primary_key + } + options[:on_delete] = extract_foreign_key_action(row["on_delete"]) + options[:on_update] = extract_foreign_key_action(row["on_update"]) + options[:deferrable] = extract_constraint_deferrable(row["deferrable"], row["deferred"]) + + options[:validate] = row["valid"] + + ForeignKeyDefinition.new(from_table, to_table, options) + end + end end end end diff --git a/lib/active_record/connection_adapters/cockroachdb/schema_statements.rb b/lib/active_record/connection_adapters/cockroachdb/schema_statements.rb index 9bec6708..ba5f469a 100644 --- a/lib/active_record/connection_adapters/cockroachdb/schema_statements.rb +++ b/lib/active_record/connection_adapters/cockroachdb/schema_statements.rb @@ -216,7 +216,6 @@ def foreign_keys(table_name) options[:deferrable] = extract_constraint_deferrable(row["deferrable"], row["deferred"]) options[:validate] = row["valid"] - to_table = PostgreSQL::Utils.unquote_identifier(row["to_table"]) ForeignKeyDefinition.new(table_name, to_table, options) end @@ -296,31 +295,6 @@ def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, **) # sql end - # This overrides the method from PostegreSQL adapter - # Resets the sequence of a table's primary key to the maximum value. - def reset_pk_sequence!(table, pk = nil, sequence = nil) - unless pk && sequence - default_pk, default_sequence = pk_and_sequence_for(table) - - pk ||= default_pk - sequence ||= default_sequence - end - - if @logger && pk && !sequence - @logger.warn "#{table} has primary key #{pk} with no default sequence." - end - - if pk && sequence - quoted_sequence = quote_table_name(sequence) - max_pk = query_value("SELECT MAX(#{quote_column_name pk}) FROM #{quote_table_name(table)}", "SCHEMA") - if max_pk.nil? - minvalue = query_value("SELECT seqmin FROM pg_sequence WHERE seqrelid = #{quote(quoted_sequence)}::regclass", "SCHEMA") - end - - query_value("SELECT setval(#{quote(quoted_sequence)}, #{max_pk ? max_pk : minvalue}, #{max_pk ? true : false})", "SCHEMA") - end - end - # override def native_database_types # Add spatial types diff --git a/test/cases/fixtures_test.rb b/test/cases/fixtures_test.rb index 1a069b8a..faa5ca3c 100644 --- a/test/cases/fixtures_test.rb +++ b/test/cases/fixtures_test.rb @@ -205,7 +205,7 @@ def self.name # under test will have primary key sequences, and the connection is from ActiveRecord::Base. # Normally, the primary keys would use CockroachDB's unique_rowid(). def test_create_symbol_fixtures - fixtures = ActiveRecord::FixtureSet.create_fixtures(FIXTURES_ROOT, :collections, collections: Course) { Course.lease_connection } + fixtures = ActiveRecord::FixtureSet.create_fixtures(FIXTURES_ROOT, :collections, collections: Course) assert Course.find_by_name("Collection"), "course is not in the database" assert fixtures.detect { |f| f.name == "collections" }, "no fixtures named 'collections' in #{fixtures.map(&:name).inspect}" diff --git a/test/cases/helper_cockroachdb.rb b/test/cases/helper_cockroachdb.rb index 0d7e4a0b..04d7321e 100644 --- a/test/cases/helper_cockroachdb.rb +++ b/test/cases/helper_cockroachdb.rb @@ -192,15 +192,17 @@ def before_teardown MiniTest::Test.include(TraceLibPlugin) end -if ENV['AR_LOG'] - ActiveRecord::Base.logger = Logger.new(STDOUT) - ActiveRecord::Base.logger.level = Logger::DEBUG - ActiveRecord::LogSubscriber::IGNORE_PAYLOAD_NAMES.clear - ActiveRecord::Base.logger.formatter = proc { |severity, time, progname, msg| - th = Thread.current[:name] - th = "THREAD=#{th}" if th - Logger::Formatter.new.call(severity, time, progname || th, msg) - } +# Log all SQL queries and print total time spent in SQL. +if ENV["AR_LOG"] + require "support/sql_logger" + case ENV["AR_LOG"].strip + when "stdout" then SQLLogger.stdout_log + when "summary" then SQLLogger.summary_log + else + SQLLogger.stdout_log + SQLLogger.summary_log + end + end # Remove the header from the schema dump to clarify tests outputs. diff --git a/test/support/sql_logger.rb b/test/support/sql_logger.rb new file mode 100644 index 00000000..e548e05e --- /dev/null +++ b/test/support/sql_logger.rb @@ -0,0 +1,55 @@ +# frozen_string_literal: true + +module SQLLogger + module_function + + def stdout_log + ActiveRecord::Base.logger = Logger.new(STDOUT) + ActiveRecord::Base.logger.level = Logger::DEBUG + ActiveRecord::LogSubscriber::IGNORE_PAYLOAD_NAMES.clear + ActiveRecord::Base.logger.formatter = proc { |severity, time, progname, msg| + th = Thread.current[:name] + th = "THREAD=#{th}" if th + Logger::Formatter.new.call(severity, time, progname || th, msg) + } + end + + def summary_log + ActiveRecord::TotalTimeSubscriber.attach_to :active_record + Minitest.after_run { + detail = ActiveRecord::TotalTimeSubscriber.hash.map { |k,v| [k, [v.sum, v.sum / v.size, v.size]]}.sort_by { |_, (_total, avg, _)| -avg }.to_h + time = detail.values.sum { |(total, _, _)| total } / 1_000 + count = detail.values.sum { |(_, _, count)| count } + puts "Total time spent in SQL: #{time}s (#{count} queries)" + puts "Detail per query kind available in tmp/query_time.json (total time in ms, avg time in ms, query count). Sorted by avg time." + File.write( + "tmp/query_time.json", + JSON.pretty_generate(detail) + ) + } + end + + # Remove content between single quotes and double quotes from keys + # to have a clear idea of which queries are being executed. + def clean_sql(sql) + sql.gsub(/".*?"/m, "\"...\"").gsub("''", "").gsub(/'.*?'/m, "'...'") + end +end + +class ActiveRecord::TotalTimeSubscriber < ActiveRecord::LogSubscriber + def self.hash + @@hash + end + + def sql(event) + # NOTE: If you want to debug a specific query, you can use a 'binding.irb' here with + # a specific condition on 'event.payload[:sql]' content. + # + # binding.irb if event.payload[:sql].include?("attr.attname, nsp.nspname") + # + @@hash ||= {} + key = SQLLogger.clean_sql(event.payload[:sql]) + @@hash[key] ||= [] + @@hash[key].push event.duration + end +end