Feature query annotation

lib/readyset/railtie.rb

@@ -38,5 +38,31 @@ class Railtie < Rails::Railtie
     rake_tasks do
       Dir[File.join(File.dirname(__FILE__), '../tasks/*.rake')].each { |f| load f }
     end
+
+    initializer 'readyset.query_annotator' do |app|
+      setup_query_annotator
+    end
+
+    def setup_query_annotator
+      config.after_initialize do
+        if Rails.env.development? || Rails.env.test?
+          if Rails.configuration.active_record.query_log_tags_enabled
+            Rails.configuration.active_record.query_log_tags ||= []
+            Rails.configuration.active_record.query_log_tags << {
+              destination: ->(context) do
+                ActiveRecord::Base.connection_db_config.name
+              end,
+            }
+          else
+            Rails.logger.warn 'Query log tags are currently disabled.' \
+              'The ReadySet gem uses these tags to display information' \
+              'in the logs about whether a query was routed to ReadySet.' \
+              'It is highly recommended that you enable query log tags by setting' \
+              '`Rails.configuration.active_record.query_log_tags_enabled` to true to' \
+              'verify that queries are being routed to ReadySet as expected.'
+          end
+        end
+      end
+    end
   end
 end

spec/internal/config/environments/test.rb

@@ -0,0 +1,3 @@
+Rails.application.configure do
+  config.active_record.query_log_tags_enabled = true
+end

spec/railtie_spec.rb

@@ -1,6 +1,6 @@
 # spec/railtie_spec.rb
 
-require 'rails_helper'
+require 'spec_helper'
 
 RSpec.describe Readyset::Railtie do
   describe 'readyset.action_controller', type: :controller do
@@ -28,6 +28,131 @@ def index
     end
   end
 
+  describe 'readyset.query_annotator' do
+    context 'when Rails.env.development? is true' do
+      context 'when query log tags are enabled' do
+        it 'adds a query_log_tag for routing to Readyset' do
+          # Setup
+          rails_env = 'development'.inquiry # Allows it to respond to development?
+          allow(Rails).to receive(:env).and_return(rails_env)
+          allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+            and_return(true)
+          expected_tag = {
+            destination: ->(context) do
+              ActiveRecord::Base.connection_db_config.name
+            end,
+          }
+          allow(Rails.configuration.active_record.query_log_tags).to receive(:<<).with(expected_tag)
+          Readyset::Railtie.setup_query_annotator
+
+          # Verify
+          expect(Rails.configuration.active_record.query_log_tags).to have_received(:<<).
+            with(expected_tag)
+        end
+      end
+
+      context 'when query log tags are not enabled' do
+        it 'logs a warning about query log tags being disabled' do
+          # Setup
+          rails_env = 'development'.inquiry # Allows it to respond to development?
+          allow(Rails).to receive(:env).and_return(rails_env)
+          allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+            and_return(false)
+
+          allow(Rails.logger).to receive(:warn).with(anything)
+
+          # Execute
+          Readyset::Railtie.setup_query_annotator
+
+          # Verify
+          expect(Rails.logger).to have_received(:warn).with(anything)
+        end
+
+        it 'does not add a query_log_tag for routing to Readyset' do
+          # Setup
+          rails_env = 'development'.inquiry # Allows it to respond to development?
+          allow(Rails).to receive(:env).and_return(rails_env)
+          allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+            and_return(false)
+          allow(Rails.configuration.active_record.query_log_tags).to receive(:<<)
+          Readyset::Railtie.setup_query_annotator
+
+          # Verify
+          expect(Rails.configuration.active_record.query_log_tags).not_to have_received(:<<)
+        end
+      end
+    end
+
+    context 'when Rails.env.test? is true' do
+      context 'when query log tags are enabled' do
+        it 'adds a query_log_tag for routing to Readyset' do
+          # Setup
+          rails_env = 'test'.inquiry # Allows it to respond to test?
+          allow(Rails).to receive(:env).and_return(rails_env)
+          allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+            and_return(true)
+          expected_tag = {
+            destination: ->(context) do
+              ActiveRecord::Base.connection_db_config.name
+            end,
+          }
+          allow(Rails.configuration.active_record.query_log_tags).to receive(:<<).with(expected_tag)
+          Readyset::Railtie.setup_query_annotator
+
+          # Verify
+          expect(Rails.configuration.active_record.query_log_tags).to have_received(:<<).
+            with(expected_tag)
+        end
+      end
+
+      context 'when query log tags are not enabled' do
+        it 'logs a warning about query log tags being disabled' do
+          # Setup
+          rails_env = 'test'.inquiry # Allows it to respond to test?
+          allow(Rails).to receive(:env).and_return(rails_env)
+          allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+            and_return(false)
+
+          allow(Rails.logger).to receive(:warn).with(anything)
+
+          # Execute
+          Readyset::Railtie.setup_query_annotator
+
+          # Verify
+          expect(Rails.logger).to have_received(:warn).with(anything)
+        end
+
+        it 'does not add a query_log_tag for routing to Readyset' do
+          # Setup
+          rails_env = 'test'.inquiry # Allows it to respond to development?
+          allow(Rails).to receive(:env).and_return(rails_env)
+          allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+            and_return(false)
+          allow(Rails.configuration.active_record.query_log_tags).to receive(:<<)
+          Readyset::Railtie.setup_query_annotator
+
+          # Verify
+          expect(Rails.configuration.active_record.query_log_tags).not_to have_received(:<<)
+        end
+      end
+    end
+
+    context 'when Rails.env.development? and Rails.env.test? are both false' do
+      it 'does not add a query_log_tag for routing to Readyset' do
+        # Setup
+        rails_env = 'production'.inquiry # Allows it to respond to development?
+        allow(Rails).to receive(:env).and_return(rails_env)
+        allow(Rails.configuration.active_record).to receive(:query_log_tags_enabled).
+          and_return(true)
+        allow(Rails.configuration.active_record.query_log_tags).to receive(:<<)
+        Readyset::Railtie.setup_query_annotator
+
+        # Verify
+        expect(Rails.configuration.active_record.query_log_tags).not_to have_received(:<<)
+      end
+    end
+  end
+
   describe 'readyset.connection_pools' do
     it 'sets up connection pools for both the reading and writing roles' do
       pools = ActiveRecord::Base.connection_handler.connection_pools

spec/ready_set_spec.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'spec_helper'
+require 'rails_helper'
 
 RSpec.describe Readyset do
   it 'has a version number' do
@@ -325,5 +325,24 @@
         expect(proxied_queries).to be_empty
       end
     end
+
+    # NOTE: If query tags aren't available, it will annotate anything.
+    # Adding a feature toggle via config would be redundant.
+    context 'when query tags are available' do
+      it 'annotates queries passing through Readyset.route' do
+        # Setup - set up custom logger
+        log = StringIO.new
+        custom_logger = ActiveSupport::Logger.new(log)
+        allow(ActiveRecord::Base).to receive(:logger).and_return(custom_logger)
+
+        # Exercise - Run a query to be routed and tagged
+        query = Cat.where(name: 'whiskers')
+        Readyset.route { query }
+
+        # Verify - Check if the query log contains the expected annotation
+        log.rewind # To latest log.
+        expect(log.read).to match(/destination[:=]readyset/)
+      end
+    end
   end
 end