From 8e7998faa883eac7f7eab7b66ea9eb27b8a27f4a Mon Sep 17 00:00:00 2001
From: Geoff Harcourt <geoff.harcourt@gmail.com>
Date: Tue, 12 Oct 2021 16:03:44 -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.
---
 lib/marginalia/comment.rb   |  5 +++++
 test/query_comments_test.rb | 12 ++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/lib/marginalia/comment.rb b/lib/marginalia/comment.rb
index f4a77de..52ac44a 100644
--- a/lib/marginalia/comment.rb
+++ b/lib/marginalia/comment.rb
@@ -162,6 +162,11 @@ def self.database
         end
       end
 
+      def self.db_role
+        return if marginalia_adapter.pool.nil?
+        ActiveRecord::Base.current_role
+      end
+
       if Gem::Version.new(ActiveRecord::VERSION::STRING) < Gem::Version.new('6.1')
         def self.connection_config
           return if marginalia_adapter.pool.nil?
diff --git a/test/query_comments_test.rb b/test/query_comments_test.rb
index a5cb9c9..2cb06fe 100644
--- a/test/query_comments_test.rb
+++ b/test/query_comments_test.rb
@@ -9,6 +9,10 @@ def pool_db_config?
   Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new('6.1')
 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'
@@ -259,6 +263,14 @@ def test_socket
     end
   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_request_id
     @env["action_dispatch.request_id"] = "some-uuid"
     Marginalia::Comment.components = [:request_id]