chore: faster referential integrity for tests

BuonOmo · rafiss · commit 8daf4fb0bc64 · 2025-01-06T10:53:18.000-05:00
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
