From f2efa0bf6ff04c84f59bbd376a6bfeca794f7931 Mon Sep 17 00:00:00 2001
From: Geoff Harcourt <geoff.harcourt@gmail.com>
Date: Sat, 23 May 2020 14:24:28 -0400
Subject: [PATCH] Add component to comment with the database role

With multiple databases in Rails 6, it can be difficult to tell from
logging if your database query is being made against the leader or
follower database. This change adds a new component, `db_role`, which
can be used with the database host and name to annotate the query with
the current ActiveRecord role.
---
 .travis.yml                 |  1 +
 gemfiles/6.0.gemfile        |  8 ++++++++
 lib/marginalia/comment.rb   |  5 +++++
 test/query_comments_test.rb | 12 ++++++++++++
 4 files changed, 26 insertions(+)
 create mode 100644 gemfiles/6.0.gemfile

diff --git a/.travis.yml b/.travis.yml
index fd1bb09..076ac42 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -22,3 +22,4 @@ gemfile:
   - gemfiles/5.0.gemfile
   - gemfiles/5.1.gemfile
   - gemfiles/5.2.gemfile
+  - gemfiles/6.0.gemfile
diff --git a/gemfiles/6.0.gemfile b/gemfiles/6.0.gemfile
new file mode 100644
index 0000000..e0f981e
--- /dev/null
+++ b/gemfiles/6.0.gemfile
@@ -0,0 +1,8 @@
+source "https://rubygems.org"
+
+gem "mysql2", "~> 0.4.10"
+gem "pg", "~> 0.15"
+gem "sqlite3", "~> 1.3.6"
+gem "rails", "= 6.0.3"
+
+gemspec :path => "../"
diff --git a/lib/marginalia/comment.rb b/lib/marginalia/comment.rb
index 783d4cd..0a5d779 100644
--- a/lib/marginalia/comment.rb
+++ b/lib/marginalia/comment.rb
@@ -155,6 +155,11 @@ def self.db_host
           end
         end
 
+        def self.db_role
+          return if marginalia_adapter.pool.nil?
+          ActiveRecord::Base.current_role
+        end
+
         def self.database
           if self.connection_config.present?
             self.connection_config[:database]
diff --git a/test/query_comments_test.rb b/test/query_comments_test.rb
index 525b675..1b524f5 100644
--- a/test/query_comments_test.rb
+++ b/test/query_comments_test.rb
@@ -17,6 +17,10 @@ def adapter_pool_available?
   Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new('3.2.19')
 end
 
+def database_role_available?
+  Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new('6.0.0')
+end
+
 require "minitest/autorun"
 require "mocha/minitest"
 require 'logger'
@@ -234,6 +238,14 @@ def test_db_host
       assert_match %r{/\*db_host:localhost}, @queries.first
     end
 
+    if database_role_available?
+      def test_db_role
+        Marginalia::Comment.components = [:db_role]
+        API::V1::PostsController.action(:driver_only).call(@env)
+        assert_match %r{/\*db_role:writing}, @queries.first
+      end
+    end
+
     def test_database
       Marginalia::Comment.components = [:database]
       API::V1::PostsController.action(:driver_only).call(@env)