feat(referential-integrity): handle autocommit_before_ddl

If `#disable_referential_integrity` is ran within a
transaction, we may be able to still run the user's code.

Otherwise we warn the user why it failed.

Author: BuonOmo

lib/active_record/connection_adapters/cockroachdb/database_statements.rb

@@ -34,7 +34,6 @@ def insert_fixtures_set(fixture_set, tables_to_delete = [])
           # our statements by calling `#execute` instead of `#execute_batch`.
           #
           # [1]: https://github.com/rails/rails/pull/52428
-
           begin # much faster without disabling referential integrity, worth trying.
             transaction(requires_new: true) do
               execute(statements, "Fixtures Load")

lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb

@@ -34,25 +34,41 @@ def check_all_foreign_keys_valid!
         end
 
         def disable_referential_integrity
-          foreign_keys = all_foreign_keys
-
-          remove_foreign_keys(foreign_keys)
-
-          # Prefixes and suffixes are added in add_foreign_key
-          # in AR7+ so we need to temporarily disable them here,
-          # otherwise prefixes/suffixes will be erroneously added.
-          old_prefix = ActiveRecord::Base.table_name_prefix
-          old_suffix = ActiveRecord::Base.table_name_suffix
-
-          begin
-            yield
-          ensure
-            ActiveRecord::Base.table_name_prefix = ""
-            ActiveRecord::Base.table_name_suffix = ""
-            add_foreign_keys(foreign_keys) # Never raises.
-
-            ActiveRecord::Base.table_name_prefix = old_prefix if defined?(old_prefix)
-            ActiveRecord::Base.table_name_suffix = old_suffix if defined?(old_suffix)
+          if transaction_open? && query_value("SHOW autocommit_before_ddl") == "off"
+            begin
+              yield
+            rescue ActiveRecord::InvalidForeignKey => e
+              warn <<-WARNING
+WARNING: Rails was not able to disable referential integrity.
+
+This is due to CockroachDB's need of committing transactions
+before a schema change occurs. To bypass this, you can set
+`autocommit_before_ddl: "on"` in your database configuration.
+WARNING
+              raise e
+            end
+          else
+            foreign_keys = all_foreign_keys
+
+            remove_foreign_keys(foreign_keys)
+
+            # Prefixes and suffixes are added in add_foreign_key
+            # in AR7+ so we need to temporarily disable them here,
+            # otherwise prefixes/suffixes will be erroneously added.
+            old_prefix = ActiveRecord::Base.table_name_prefix
+            old_suffix = ActiveRecord::Base.table_name_suffix
+
+            begin
+              yield
+            ensure
+              ActiveRecord::Base.table_name_prefix = ""
+              ActiveRecord::Base.table_name_suffix = ""
+
+              add_foreign_keys(foreign_keys) # Never raises.
+
+              ActiveRecord::Base.table_name_prefix = old_prefix if defined?(old_prefix)
+              ActiveRecord::Base.table_name_suffix = old_suffix if defined?(old_suffix)
+            end
           end
         end
 

test/cases/adapters/cockroachdb/referential_integrity_test.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require "cases/helper_cockroachdb"
+require "support/connection_helper" # for #reset_connection
+require "support/copy_cat"
+
+class CockroachDBReferentialIntegrityTest < ActiveRecord::PostgreSQLTestCase
+  include ConnectionHelper
+
+  module ProgrammerMistake
+    def execute_batch(sql, name = nil)
+      raise ArgumentError, "something is not right." if name.match?(/referential integrity/)
+      super
+    end
+  end
+
+  def setup
+    @connection = ActiveRecord::Base.lease_connection
+  end
+
+  def teardown
+    reset_connection
+  end
+
+  exclude_from_transactional_tests :test_only_catch_active_record_errors_others_bubble_up
+  CopyCat.copy_methods(self, ::PostgreSQLReferentialIntegrityTest, :test_only_catch_active_record_errors_others_bubble_up)
+
+  def test_should_reraise_invalid_foreign_key_exception_and_show_warning
+    warning = capture(:stderr) do
+      e = assert_raises(ActiveRecord::InvalidForeignKey) do
+        @connection.disable_referential_integrity do
+          @connection.execute("INSERT INTO authors (name, author_address_id) VALUES ('Mona Chollet', 42)")
+        end
+      end
+      assert_match (/Key \(author_address_id\)=\(42\) is not present in table/), e.message
+    end
+    assert_match (/WARNING: Rails was not able to disable referential integrity/), warning
+    assert_match (/autocommit_before_ddl/), warning
+  end
+
+  def test_no_warning_nor_error_with_autocommit_before_ddl
+    @connection.execute("SET SESSION autocommit_before_ddl = 'on'")
+    warning = capture(:stderr) do
+      @connection.disable_referential_integrity do
+        @connection.execute("INSERT INTO authors (name, author_address_id) VALUES ('Mona Chollet', 42)")
+        @connection.truncate(:authors)
+      end
+    end
+    assert_predicate warning, :blank?, "expected no warnings but got:\n#{warning}"
+  end
+end

test/excludes/PostgreSQLReferentialIntegrityTest.rb

@@ -1,4 +1,14 @@
-exclude :test_only_catch_active_record_errors_others_bubble_up, "Skipping until we can triage further. See https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/48"
-exclude :test_should_reraise_invalid_foreign_key_exception_and_show_warning, "Skipping until we can triage further. See https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/48"
-exclude :test_does_not_break_transactions, "Skipping until we can triage further. See https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/48"
-exclude :test_does_not_break_nested_transactions, "Skipping until we can triage further. See https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/48"
+exclude :test_should_reraise_invalid_foreign_key_exception_and_show_warning,
+    "CockroachDB has a different limitation as there is no" \
+    "'DISABLE TRIGGER' statement."
+
+break_tx = "CockroachDB will always alter transactions when " \
+  "trying to disable referential integrity. Either it cannot " \
+  "work within transaction, or autocommit_before_ddl is set " \
+  "and transactions will be committed."
+exclude :test_does_not_break_transactions, break_tx
+exclude :test_does_not_break_nested_transactions, break_tx
+
+exclude :test_only_catch_active_record_errors_others_bubble_up,
+    "Reimplemented in test/cases/adapters/cockroachdb/referential_integrity_test.rb" \
+    " to use a different trigger for the error."