Add ActiveRecordProxyAdapters::LogSubscriber (#8)

* Add ActiveRecordProxyAdapters::LogSubscriber
* Require log subscriber in unit tests

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## [Unreleased]
 
+- Add custom log subscriber to tag queries based on the adapter being used
 - Fix replica connection pool getter when database configurations have multiple replicas
 - Retrieve replica pool without checking out a connection
 

docker-compose.yml

@@ -50,6 +50,7 @@ services:
         - POSTGRES_LOG_STATEMENT=${POSTGRES_LOG_STATEMENT:-}
         - REPLICA_USER=replicator
         - REPLICA_PASSWORD=replicator
+    container_name: postgres_primary
     environment:
       POSTGRES_DB: postgres
       POSTGRES_USER: postgres_primary_test

lib/active_record_proxy_adapters/configuration.rb

@@ -5,8 +5,10 @@
 module ActiveRecordProxyAdapters
   # Provides a global configuration object to configure how the proxy should behave.
   class Configuration
-    PROXY_DELAY      = 2.seconds.freeze
-    CHECKOUT_TIMEOUT = 2.seconds.freeze
+    PROXY_DELAY                   = 2.seconds.freeze
+    CHECKOUT_TIMEOUT              = 2.seconds.freeze
+    LOG_SUBSCRIBER_PRIMARY_PREFIX = proc { |event| "#{event.payload[:connection].class::ADAPTER_NAME} Primary" }.freeze
+    LOG_SUBSCRIBER_REPLICA_PREFIX = proc { |event| "#{event.payload[:connection].class::ADAPTER_NAME} Replica" }.freeze
 
     # @return [ActiveSupport::Duration] How long the proxy should reroute all read requests to the primary database
     #   since the latest write. Defaults to PROXY_DELAY.
@@ -15,9 +17,25 @@ class Configuration
     #   Defaults to CHECKOUT_TIMEOUT.
     attr_accessor :checkout_timeout
 
+    # @return [Proc] Prefix for the log subscriber when the primary database is used.
+    attr_reader :log_subscriber_primary_prefix
+
+    # @return [Proc] Prefix for the log subscriber when the replica database is used.
+    attr_reader :log_subscriber_replica_prefix
+
     def initialize
-      self.proxy_delay      = PROXY_DELAY
-      self.checkout_timeout = CHECKOUT_TIMEOUT
+      self.proxy_delay                   = PROXY_DELAY
+      self.checkout_timeout              = CHECKOUT_TIMEOUT
+      self.log_subscriber_primary_prefix = LOG_SUBSCRIBER_PRIMARY_PREFIX
+      self.log_subscriber_replica_prefix = LOG_SUBSCRIBER_REPLICA_PREFIX
+    end
+
+    def log_subscriber_primary_prefix=(prefix)
+      @log_subscriber_primary_prefix = prefix.is_a?(Proc) ? prefix : proc { prefix.to_s }
+    end
+
+    def log_subscriber_replica_prefix=(prefix)
+      @log_subscriber_replica_prefix = prefix.is_a?(Proc) ? prefix : proc { prefix.to_s }
     end
   end
 end

lib/active_record_proxy_adapters/log_subscriber.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module ActiveRecordProxyAdapters
+  class LogSubscriber < ActiveRecord::LogSubscriber # rubocop:disable Style/Documentation
+    attach_to :active_record
+
+    IGNORE_PAYLOAD_NAMES = %w[SCHEMA EXPLAIN].freeze
+
+    def sql(event)
+      payload = event.payload
+      name = payload[:name]
+      unless IGNORE_PAYLOAD_NAMES.include?(name)
+        name = [database_instance_prefix_for(event), name].compact.join(" ")
+        payload[:name] = name
+      end
+      super
+    end
+
+    protected
+
+    def database_instance_prefix_for(event)
+      connection = event.payload[:connection]
+      config = connection.instance_variable_get(:@config)
+      prefix = if config[:replica] || config["replica"]
+                 log_subscriber_replica_prefix
+               else
+                 log_subscriber_primary_prefix
+               end
+
+      "[#{prefix.call(event)}]"
+    end
+
+    private
+
+    delegate :log_subscriber_primary_prefix, :log_subscriber_replica_prefix, to: :config
+
+    def config
+      ActiveRecordProxyAdapters.config
+    end
+  end
+end

spec/spec_helper.rb

@@ -30,9 +30,11 @@
 
 require "active_record_proxy_adapters"
 require "active_record_proxy_adapters/connection_handling"
+require "active_record_proxy_adapters/log_subscriber"
 require_relative "test_helper"
 
 ActiveRecord::Base.extend ActiveRecordProxyAdapters::ConnectionHandling
+ActiveRecord::Base.logger = Logger.new(Tempfile.create)
 
 ENV["RAILS_ENV"] ||= TestHelper.env_name