From ed91385c18cdc831eea4fc5178a0829f640990fd Mon Sep 17 00:00:00 2001
From: David Taylor <david@taylorhq.com>
Date: Tue, 8 Dec 2020 00:03:31 +0000
Subject: [PATCH] DEV: Update `DB.after_commit` to be compatible with 'real'
 transactions (#11294)

Previously it matched the behavior of standard ActiveRecord after_commit callbacks. They do not work well within `joinable: false` nested transactions. Now `DB.after_commit` callbacks will only be run when the outermost transaction has been committed.

Tests always run inside transactions, so this also introduces some logic to run callbacks once the test-wrapping transaction is reached.
---
 lib/mini_sql_multisite_connection.rb          | 28 +++++++++----
 .../lib/mini_sql_multisite_connection_spec.rb | 39 ++++++++++++-------
 spec/rails_helper.rb                          |  5 +++
 3 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/lib/mini_sql_multisite_connection.rb b/lib/mini_sql_multisite_connection.rb
index 202cf87b4e1..2cf50ca61e9 100644
--- a/lib/mini_sql_multisite_connection.rb
+++ b/lib/mini_sql_multisite_connection.rb
@@ -32,7 +32,12 @@ class MiniSqlMultisiteConnection < MiniSql::Postgres::Connection
     end
 
     def committed!(*)
-      @callback.call
+      if DB.transaction_open?
+        # Nested transaction. Pass the callback to the parent
+        ActiveRecord::Base.connection.add_transaction_record(self)
+      else
+        @callback.call
+      end
     end
 
     def before_committed!(*); end
@@ -42,16 +47,25 @@ class MiniSqlMultisiteConnection < MiniSql::Postgres::Connection
     end
   end
 
+  def transaction_open?
+    ActiveRecord::Base.connection.transaction_open?
+  end
+
+  if Rails.env.test?
+    def test_transaction=(transaction)
+      @test_transaction = transaction
+    end
+
+    def transaction_open?
+      ActiveRecord::Base.connection.current_transaction != @test_transaction
+    end
+  end
+
   # Allows running arbitrary code after the current transaction has been committed.
   # Works with nested ActiveRecord transaction blocks. Useful for scheduling sidekiq jobs.
   # If not currently in a transaction, will execute immediately
   def after_commit(&blk)
-    return blk.call if !ActiveRecord::Base.connection.transaction_open?
-
-    # In tests, everything is run inside a transaction.
-    # To run immediately, check for joinable? transaction
-    # This mimics core rails behavior: https://github.com/rails/rails/blob/348e142b/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb#L211
-    return blk.call if Rails.env.test? && !ActiveRecord::Base.connection.current_transaction.joinable?
+    return blk.call if !transaction_open?
 
     ActiveRecord::Base.connection.add_transaction_record(
       AfterCommitWrapper.new(&blk)
diff --git a/spec/lib/mini_sql_multisite_connection_spec.rb b/spec/lib/mini_sql_multisite_connection_spec.rb
index 9cfa452178d..47e8f558371 100644
--- a/spec/lib/mini_sql_multisite_connection_spec.rb
+++ b/spec/lib/mini_sql_multisite_connection_spec.rb
@@ -5,27 +5,38 @@ require 'rails_helper'
 describe MiniSqlMultisiteConnection do
 
   describe "after_commit" do
-    it "runs callbacks after outermost transaction is committed" do
+    it "works for 'fake' (joinable) transactions" do
       outputString = "1"
 
-      # Main transaction
       ActiveRecord::Base.transaction do
         outputString += "2"
+        DB.exec("SELECT 1")
+        ActiveRecord::Base.transaction do
+          DB.exec("SELECT 2")
+          outputString += "3"
+          DB.after_commit { outputString += "6" }
+          outputString += "4"
+        end
+        DB.after_commit { outputString += "7" }
+        outputString += "5"
+      end
 
-          # Nested transaction
-          ActiveRecord::Base.transaction do
+      expect(outputString).to eq("1234567")
+    end
+
+    it "works for real (non-joinable) transactions" do
+      outputString = "1"
+
+      ActiveRecord::Base.transaction(requires_new: true, joinable: false) do
+        outputString += "2"
+        DB.exec("SELECT 1")
+          ActiveRecord::Base.transaction(requires_new: true) do
+            DB.exec("SELECT 2")
             outputString += "3"
-
-              DB.after_commit do
-                outputString += "6"
-              end
-              outputString += "4"
+            DB.after_commit { outputString += "6" }
+            outputString += "4"
           end
-
-          DB.after_commit do
-            outputString += "7"
-          end
-
+          DB.after_commit { outputString += "7" }
           outputString += "5"
       end
 
diff --git a/spec/rails_helper.rb b/spec/rails_helper.rb
index 42f7d6f69a6..6286218a94c 100644
--- a/spec/rails_helper.rb
+++ b/spec/rails_helper.rb
@@ -273,6 +273,11 @@ RSpec.configure do |config|
     expect(before_event_count).to eq(after_event_count), "DiscourseEvent registrations were not cleaned up"
   end
 
+  config.before :each do
+    # This allows DB.transaction_open? to work in tests. See lib/mini_sql_multisite_connection.rb
+    DB.test_transaction = ActiveRecord::Base.connection.current_transaction
+  end
+
   config.before(:each, type: :multisite) do
     Rails.configuration.multisite = true # rubocop:disable Discourse/NoDirectMultisiteManipulation