Merge pull request #43097 from Shopify/ar-query-transformers

Refactor ActiveRecord::QueryLogs hook point

activerecord/lib/active_record.rb

@@ -329,6 +329,9 @@ def self.global_executor_concurrency # :nodoc:
   singleton_class.attr_accessor :verify_foreign_keys_for_fixtures
   self.verify_foreign_keys_for_fixtures = false
 
+  singleton_class.attr_accessor :query_transformers
+  self.query_transformers = []
+
   def self.eager_load!
     super
     ActiveRecord::Locking.eager_load!

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -741,6 +741,13 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name =
           end
         end
 
+        def transform_query(sql)
+          ActiveRecord.query_transformers.each do |transformer|
+            sql = transformer.call(sql)
+          end
+          sql
+        end
+
         def translate_exception(exception, message:, sql:, binds:)
           # override in derived class
           case exception

activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb

@@ -37,15 +37,26 @@ def explain(arel, binds = [])
           MySQL::ExplainPrettyPrinter.new.pp(result, elapsed)
         end
 
+        def execute(sql, name = nil, async: false)
+          # make sure we carry over any changes to ActiveRecord.default_timezone that have been
+          # made since we established the connection
+          @connection.query_options[:database_timezone] = ActiveRecord.default_timezone
+
+          super
+        end
+        alias_method :raw_execute, :execute
+        private :raw_execute
+
         # Executes the SQL statement in the context of this connection.
         def execute(sql, name = nil, async: false)
+          sql = transform_query(sql)
           check_if_write_query(sql)
 
           # make sure we carry over any changes to ActiveRecord.default_timezone that have been
           # made since we established the connection
           @connection.query_options[:database_timezone] = ActiveRecord.default_timezone
 
-          super
+          raw_execute(sql, name, async: async)
         end
 
         def exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc:
@@ -81,8 +92,9 @@ def exec_delete(sql, name = nil, binds = []) # :nodoc:
 
         private
           def execute_batch(statements, name = nil)
+            statements = statements.map { |sql| transform_query(sql) }
             combine_multi_statements(statements).each do |statement|
-              execute(statement, name)
+              raw_execute(statement, name)
             end
             @connection.abandon_results!
           end
@@ -147,6 +159,7 @@ def max_allowed_packet
           end
 
           def exec_stmt_and_free(sql, name, binds, cache_stmt: false, async: false)
+            sql = transform_query(sql)
             check_if_write_query(sql)
 
             materialize_transactions

activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb

@@ -35,6 +35,7 @@ def write_query?(sql) # :nodoc:
         # Note: the PG::Result object is manually memory managed; if you don't
         # need it specifically, you may want consider the <tt>exec_query</tt> wrapper.
         def execute(sql, name = nil)
+          sql = transform_query(sql)
           check_if_write_query(sql)
 
           materialize_transactions

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -696,6 +696,7 @@ def load_types_queries(initializer, oids)
         FEATURE_NOT_SUPPORTED = "0A000" # :nodoc:
 
         def execute_and_clear(sql, name, binds, prepare: false, async: false)
+          sql = transform_query(sql)
           check_if_write_query(sql)
 
           if !prepare || without_prepared_statement?(binds)

activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb

@@ -19,6 +19,7 @@ def explain(arel, binds = [])
         end
 
         def execute(sql, name = nil) # :nodoc:
+          sql = transform_query(sql)
           check_if_write_query(sql)
 
           materialize_transactions
@@ -32,6 +33,7 @@ def execute(sql, name = nil) # :nodoc:
         end
 
         def exec_query(sql, name = nil, binds = [], prepare: false, async: false) # :nodoc:
+          sql = transform_query(sql)
           check_if_write_query(sql)
 
           materialize_transactions
@@ -104,6 +106,7 @@ def reset_read_uncommitted
           end
 
           def execute_batch(statements, name = nil)
+            statements = statements.map { |sql| transform_query(sql) }
             sql = combine_multi_statements(statements)
 
             check_if_write_query(sql)

activerecord/lib/active_record/query_logs.rb

@@ -118,13 +118,14 @@ def with_tag(tag, &block)
         inline_tags.pop
       end
 
-      def add_query_log_tags_to_sql(sql) # :nodoc:
-        comments.each do |comment|
-          unless sql.include?(comment)
-            sql = prepend_comment ? "#{comment} #{sql}" : "#{sql} #{comment}"
-          end
+      def call(sql) # :nodoc:
+        parts = self.comments
+        if prepend_comment
+          parts << sql
+        else
+          parts.unshift(sql)
         end
-        sql
+        parts.join(" ")
       end
 
       private
@@ -198,15 +199,5 @@ def inline_tag_content
           inline_tags.join
         end
     end
-
-    module ExecutionMethods
-      def execute(sql, *args, **kwargs)
-        super(ActiveRecord::QueryLogs.add_query_log_tags_to_sql(sql), *args, **kwargs)
-      end
-
-      def exec_query(sql, *args, **kwargs)
-        super(ActiveRecord::QueryLogs.add_query_log_tags_to_sql(sql), *args, **kwargs)
-      end
-    end
   end
 end

activerecord/lib/active_record/railtie.rb

@@ -360,6 +360,7 @@ class Railtie < Rails::Railtie # :nodoc:
     initializer "active_record.query_log_tags_config" do |app|
       config.after_initialize do
         if app.config.active_record.query_log_tags_enabled
+          ActiveRecord.query_transformers << ActiveRecord::QueryLogs
           ActiveRecord::QueryLogs.taggings.merge!(
             application:  Rails.application.class.name.split("::").first,
             pid:          -> { Process.pid },
@@ -375,12 +376,6 @@ class Railtie < Rails::Railtie # :nodoc:
           if app.config.active_record.cache_query_log_tags
             ActiveRecord::QueryLogs.cache_query_log_tags = true
           end
-
-          ActiveSupport.on_load(:active_record) do
-            ConnectionAdapters::AbstractAdapter.descendants.each do |klass|
-              klass.prepend(QueryLogs::ExecutionMethods) if klass.descendants.empty?
-            end
-          end
         end
       end
     end

activerecord/test/cases/query_logs_test.rb

@@ -12,16 +12,16 @@ class QueryLogsTest < ActiveRecord::TestCase
 
   def setup
     # Enable the query tags logging
-    ActiveRecord::Base.connection.class_eval do
-      prepend(ActiveRecord::QueryLogs::ExecutionMethods)
-    end
+    @original_transformers = ActiveRecord.query_transformers
     @original_prepend = ActiveRecord::QueryLogs.prepend_comment
+    ActiveRecord.query_transformers += [ActiveRecord::QueryLogs]
     ActiveRecord::QueryLogs.prepend_comment = false
     ActiveRecord::QueryLogs.cache_query_log_tags = false
     ActiveRecord::QueryLogs.cached_comment = nil
   end
 
   def teardown
+    ActiveRecord.query_transformers = @original_transformers
     ActiveRecord::QueryLogs.prepend_comment = @original_prepend
     ActiveRecord::QueryLogs.tags = []
   end
@@ -100,10 +100,10 @@ def test_retrieves_comment_from_cache_when_enabled_and_set
     ActiveRecord::QueryLogs.cache_query_log_tags = true
     ActiveRecord::QueryLogs.tags = [ :application ]
 
-    assert_equal " /*application:active_record*/", ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
+    assert_equal " /*application:active_record*/", ActiveRecord::QueryLogs.call("")
 
     ActiveRecord::QueryLogs.stub(:cached_comment, "/*cached_comment*/") do
-      assert_equal " /*cached_comment*/", ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
+      assert_equal " /*cached_comment*/", ActiveRecord::QueryLogs.call("")
     end
   ensure
     ActiveRecord::QueryLogs.cached_comment = nil
@@ -115,12 +115,12 @@ def test_resets_cache_on_context_update
     ActiveRecord::QueryLogs.update_context(temporary: "value")
     ActiveRecord::QueryLogs.tags = [ temporary_tag: ->(context) { context[:temporary] } ]
 
-    assert_equal " /*temporary_tag:value*/", ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
+    assert_equal " /*temporary_tag:value*/", ActiveRecord::QueryLogs.call("")
 
     ActiveRecord::QueryLogs.update_context(temporary: "new_value")
 
     assert_nil ActiveRecord::QueryLogs.cached_comment
-    assert_equal " /*temporary_tag:new_value*/", ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
+    assert_equal " /*temporary_tag:new_value*/", ActiveRecord::QueryLogs.call("")
   ensure
     ActiveRecord::QueryLogs.cached_comment = nil
     ActiveRecord::QueryLogs.cache_query_log_tags = false
@@ -201,15 +201,6 @@ def test_bad_inline_tags
     end
   end
 
-  def test_inline_tags_are_deduped
-    ActiveRecord::QueryLogs.tags = [ :application ]
-    assert_sql(%r{select id from posts /\*foo\*/ /\*application:active_record\*/$}) do
-      ActiveRecord::QueryLogs.with_tag("foo") do
-        ActiveRecord::Base.connection.execute "select id from posts /*foo*/"
-      end
-    end
-  end
-
   def test_empty_comments_are_not_added
     original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ empty: -> { nil } ]

railties/test/application/query_logs_test.rb

@@ -18,7 +18,7 @@ class User < ActiveRecord::Base
       app_file "app/controllers/users_controller.rb", <<-RUBY
         class UsersController < ApplicationController
           def index
-            render inline: ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
+            render inline: ActiveRecord::QueryLogs.call("")
           end
 
           def dynamic_content
@@ -30,7 +30,7 @@ def dynamic_content
       app_file "app/jobs/user_job.rb", <<-RUBY
         class UserJob < ActiveJob::Base
           def perform
-            ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
+            ActiveRecord::QueryLogs.call("")
           end
 
           def dynamic_content
@@ -57,15 +57,15 @@ def app
     test "does not modify the query execution path by default" do
       boot_app
 
-      assert_equal ActiveRecord::Base.connection.method(:execute).owner, ActiveRecord::ConnectionAdapters::SQLite3::DatabaseStatements
+      assert_not_includes ActiveRecord.query_transformers, ActiveRecord::QueryLogs
     end
 
     test "prepends the query execution path when enabled" do
       add_to_config "config.active_record.query_log_tags_enabled = true"
 
       boot_app
 
-      assert_equal ActiveRecord::Base.connection.method(:execute).owner, ActiveRecord::QueryLogs::ExecutionMethods
+      assert_includes ActiveRecord.query_transformers, ActiveRecord::QueryLogs
     end
 
     test "controller and job tags are defined by default" do