chore: faster referential integrity for tests

Author: BuonOmo

CONTRIBUTING.md

@@ -1,6 +1,5 @@
 # Getting started
 
-
 ## ActiveRecord adapters and you
 
 There are two repositories for the ActiveRecord adapter. The one you're in
@@ -121,7 +120,6 @@ master branch, with an alpha build of CockroachDB. it would be even
 better to be able to test multiple versions of the adapter, and do so
 against different versions of CockroachDB.
 
-
 ## Adding feature support
 
 As CockroachDB improves, so do the features that can be supported in
@@ -131,15 +129,34 @@ gates should be toggled. Something that would help this process would be
 linking those issues back to this adapter so that part of the feature
 completing includes updating the adapter.
 
-
 ## Execute only tests that run with a connection
 
-I have not investigated if this is already possible, but I would assume
-no.
-
-A possible way to approach this would be to add a shim to cause any
-tests that use it to fail, and grep the tests that pass and then skip
-them.
+- Check for TODO or NOTE tags that are referencing the old or new version of
+  rails.
+  ```bash
+  rg 'TODO|NOTE' --after-context=2
+  ```
+- Check postgresql_specific_schema.rb changelog in rails, and apply the changes
+  you want. Ex:
+  ```bash
+  git diff v7.1.4..v7.2.1 -- $(fd postgresql_specific_schema)
+  ```
+- 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:
+  - [ ] `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.
+- In the same way, run `test/cases/fixtures_test.rb` first, and check
+  if this corrupted the test database for other tests.
+- For both of the above, the diff of `schema.rb` can be useful:
+  ```bash
+  git diff v7.1.2..v7.2.1 -- activerecord/test/schema/schema.rb
+  ```
 
 ## Publishing to Rubygems
 
@@ -151,7 +168,6 @@ gem build ...
 gem publish <output file>
 ```
 
-
 # Notes
 
 When executing the test suite, each test file will reload fixtures. This
@@ -182,10 +198,8 @@ cleaned up, or skipped until passing.
 The purpose of these was to make the tests grep-able while going through
 all the failures.
 
-
 [cockroachdb/cockroach#20753]: https://github.com/cockroachdb/cockroach/issues/20753#issuecomment-352810425
 
-
 ## Tracked test failures
 
 Some of the skipped failures are:
@@ -224,7 +238,6 @@ most of the touchpoints including test failures and temporary monkey
 patches. Some monkey patches were made directly to Rails, which will
 need to be cleaned up.
 
-
 # Notes for the non-Rubyer
 
 rbenv is an environment manager that lets you manage and swap between

lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb

@@ -20,10 +20,16 @@ def 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
@@ -38,23 +44,83 @@ 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|
-              begin
-                add_foreign_key(foreign_key.from_table, foreign_key.to_table, **foreign_key.options)
-              rescue ActiveRecord::StatementInvalid => error
-                if error.cause.class == PG::DuplicateObject
-                  # This error is safe to ignore because the yielded caller
-                  # already re-added the foreign key constraint.
-                else
-                  raise error
-                end
-              end
+              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
           ensure
             ActiveRecord::Base.table_name_prefix = old_prefix
             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

lib/active_record/connection_adapters/cockroachdb/schema_statements.rb

@@ -245,31 +245,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

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