Merge pull request #20921 from pboling/fix-sql-colors-in-log-subscriber

Fix and Improve sql logging coloration in `ActiveRecord::LogSubscriber`.

3 files changed, 126 insertions, 8 deletions

diff --git a/activerecord/lib/active_record/log_subscriber.rb b/activerecord/lib/active_record/log_subscriber.rb
index 4d597a0ab1..b63caa4473 100644
--- a/activerecord/lib/active_record/log_subscriber.rb
+++ b/activerecord/lib/active_record/log_subscriber.rb
@@ -47,20 +47,40 @@ module ActiveRecord
         binds = "  " + payload[:binds].map { |attr| render_bind(attr) }.inspect
       end
 
-      name = color(name, nil, true)
+      name = colorize_payload_name(name, payload[:name])
       sql  = color(sql, sql_color(sql), true)
 
       debug "  #{name}  #{sql}#{binds}"
     end
 
+    private
+
+    def colorize_payload_name(name, payload_name)
+      if payload_name.blank? || payload_name == "SQL" # SQL vs Model Load/Exists
+        color(name, MAGENTA, true)
+      else
+        color(name, CYAN, true)
+      end
+    end
+
     def sql_color(sql)
       case sql
-        when /\s*\Ainsert/i      then GREEN
-        when /\s*\Aselect/i      then BLUE
-        when /\s*\Aupdate/i      then YELLOW
-        when /\s*\Adelete/i      then RED
-        when /transaction\s*\Z/i then CYAN
-        else MAGENTA
+        when /\A\s*rollback/mi
+          RED
+        when /\s*.*?select .*for update/mi, /\A\s*lock/mi
+          WHITE
+        when /\A\s*select/i
+          BLUE
+        when /\A\s*insert/i
+          GREEN
+        when /\A\s*update/i
+          YELLOW
+        when /\A\s*delete/i
+          RED
+        when /transaction\s*\Z/i
+          CYAN
+        else
+          MAGENTA
       end
     end
 
diff --git a/activerecord/test/cases/log_subscriber_test.rb b/activerecord/test/cases/log_subscriber_test.rb
index 4192d12ff4..3846ba8e7f 100644
--- a/activerecord/test/cases/log_subscriber_test.rb
+++ b/activerecord/test/cases/log_subscriber_test.rb
@@ -7,6 +7,20 @@ require "active_support/log_subscriber/test_helper"
 class LogSubscriberTest < ActiveRecord::TestCase
   include ActiveSupport::LogSubscriber::TestHelper
   include ActiveSupport::Logger::Severity
+  REGEXP_CLEAR = Regexp.escape(ActiveRecord::LogSubscriber::CLEAR)
+  REGEXP_BOLD = Regexp.escape(ActiveRecord::LogSubscriber::BOLD)
+  REGEXP_MAGENTA = Regexp.escape(ActiveRecord::LogSubscriber::MAGENTA)
+  REGEXP_CYAN = Regexp.escape(ActiveRecord::LogSubscriber::CYAN)
+  SQL_COLORINGS = {
+      SELECT: Regexp.escape(ActiveRecord::LogSubscriber::BLUE),
+      INSERT: Regexp.escape(ActiveRecord::LogSubscriber::GREEN),
+      UPDATE: Regexp.escape(ActiveRecord::LogSubscriber::YELLOW),
+      DELETE: Regexp.escape(ActiveRecord::LogSubscriber::RED),
+      LOCK: Regexp.escape(ActiveRecord::LogSubscriber::WHITE),
+      ROLLBACK: Regexp.escape(ActiveRecord::LogSubscriber::RED),
+      TRANSACTION: REGEXP_CYAN,
+      OTHER: REGEXP_MAGENTA
+  }
 
   class TestDebugLogSubscriber < ActiveRecord::LogSubscriber
     attr_reader :debugs
@@ -71,6 +85,90 @@ class LogSubscriberTest < ActiveRecord::TestCase
     assert_match(/SELECT .*?FROM .?developers.?/i, @logger.logged(:debug).last)
   end
 
+  def test_basic_query_logging_coloration
+    event = Struct.new(:duration, :payload)
+    logger = TestDebugLogSubscriber.new
+    logger.colorize_logging = true
+    SQL_COLORINGS.each do |verb, color_regex|
+      logger.sql(event.new(0, sql: verb.to_s))
+      assert_match(/#{REGEXP_BOLD}#{color_regex}#{verb}#{REGEXP_CLEAR}/i, logger.debugs.last)
+    end
+  end
+
+  def test_basic_payload_name_logging_coloration_generic_sql
+    event = Struct.new(:duration, :payload)
+    logger = TestDebugLogSubscriber.new
+    logger.colorize_logging = true
+    SQL_COLORINGS.each do |verb, _|
+      logger.sql(event.new(0, sql: verb.to_s))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_MAGENTA} \(0.0ms\)#{REGEXP_CLEAR}/i, logger.debugs.last)
+
+      logger.sql(event.new(0, {sql: verb.to_s, name: "SQL"}))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_MAGENTA}SQL \(0.0ms\)#{REGEXP_CLEAR}/i, logger.debugs.last)
+    end
+  end
+
+  def test_basic_payload_name_logging_coloration_named_sql
+    event = Struct.new(:duration, :payload)
+    logger = TestDebugLogSubscriber.new
+    logger.colorize_logging = true
+    SQL_COLORINGS.each do |verb, _|
+      logger.sql(event.new(0, {sql: verb.to_s, name: "Model Load"}))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_CYAN}Model Load \(0.0ms\)#{REGEXP_CLEAR}/i, logger.debugs.last)
+
+      logger.sql(event.new(0, {sql: verb.to_s, name: "Model Exists"}))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_CYAN}Model Exists \(0.0ms\)#{REGEXP_CLEAR}/i, logger.debugs.last)
+
+      logger.sql(event.new(0, {sql: verb.to_s, name: "ANY SPECIFIC NAME"}))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_CYAN}ANY SPECIFIC NAME \(0.0ms\)#{REGEXP_CLEAR}/i, logger.debugs.last)
+    end
+  end
+
+  def test_query_logging_coloration_with_nested_select
+    event = Struct.new(:duration, :payload)
+    logger = TestDebugLogSubscriber.new
+    logger.colorize_logging = true
+    SQL_COLORINGS.slice(:SELECT, :INSERT, :UPDATE, :DELETE).each do |verb, color_regex|
+      logger.sql(event.new(0, sql: "#{verb} WHERE ID IN SELECT"))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_MAGENTA} \(0.0ms\)#{REGEXP_CLEAR}  #{REGEXP_BOLD}#{color_regex}#{verb} WHERE ID IN SELECT#{REGEXP_CLEAR}/i, logger.debugs.last)
+    end
+  end
+
+  def test_query_logging_coloration_with_multi_line_nested_select
+    event = Struct.new(:duration, :payload)
+    logger = TestDebugLogSubscriber.new
+    logger.colorize_logging = true
+    SQL_COLORINGS.slice(:SELECT, :INSERT, :UPDATE, :DELETE).each do |verb, color_regex|
+      sql = <<-EOS
+        #{verb}
+        WHERE ID IN (
+          SELECT ID FROM THINGS
+        )
+      EOS
+      logger.sql(event.new(0, sql: sql))
+      assert_match(/#{REGEXP_BOLD}#{REGEXP_MAGENTA} \(0.0ms\)#{REGEXP_CLEAR}  #{REGEXP_BOLD}#{color_regex}.*#{verb}.*#{REGEXP_CLEAR}/mi, logger.debugs.last)
+    end
+  end
+
+  def test_query_logging_coloration_with_lock
+    event = Struct.new(:duration, :payload)
+    logger = TestDebugLogSubscriber.new
+    logger.colorize_logging = true
+    sql = <<-EOS
+      SELECT * FROM
+        (SELECT * FROM mytable FOR UPDATE) ss
+      WHERE col1 = 5;
+    EOS
+    logger.sql(event.new(0, sql: sql))
+    assert_match(/#{REGEXP_BOLD}#{REGEXP_MAGENTA} \(0.0ms\)#{REGEXP_CLEAR}  #{REGEXP_BOLD}#{SQL_COLORINGS[:LOCK]}.*FOR UPDATE.*#{REGEXP_CLEAR}/mi, logger.debugs.last)
+
+    sql = <<-EOS
+      LOCK TABLE films IN SHARE MODE;
+    EOS
+    logger.sql(event.new(0, sql: sql))
+    assert_match(/#{REGEXP_BOLD}#{REGEXP_MAGENTA} \(0.0ms\)#{REGEXP_CLEAR}  #{REGEXP_BOLD}#{SQL_COLORINGS[:LOCK]}.*LOCK TABLE.*#{REGEXP_CLEAR}/mi, logger.debugs.last)
+  end
+
   def test_exists_query_logging
     Developer.exists? 1
     wait
diff --git a/activesupport/lib/active_support/log_subscriber.rb b/activesupport/lib/active_support/log_subscriber.rb
index e95dc5a866..e782cd2d4b 100644
--- a/activesupport/lib/active_support/log_subscriber.rb
+++ b/activesupport/lib/active_support/log_subscriber.rb
@@ -95,7 +95,7 @@ module ActiveSupport
       METHOD
     end
 
-    # Set color by using a string or one of the defined constants. If a third
+    # Set color by using a symbol or one of the defined constants. If a third
     # option is set to +true+, it also adds bold to the string. This is based
     # on the Highline implementation and will automatically append CLEAR to the
     # end of the returned String.