Add database_role option to scope checks for replicas (#57)

* Add `database_role` option to scope checks for replicas

It can be useful to confirm that certain database queries are being made
against a primary database or a replica database. This change adds a new
option `database_role`, which makes the query counter only count queries
that were made against a database with the specified role.

We've used this logic at CommonLit for several years to ensure we're
querying the right database. This works in test environments even if
your primary and replica database configurations point at the same
database.

* remove explicit sqlite declaration

* allow rails 7.1.0

* add Rails 7.1.0 to CI Build

---------

Co-authored-by: Jonathan Rochkind <[email protected]>

gemfiles/rails_6_1.gemfile

@@ -3,6 +3,5 @@
 source "https://rubygems.org"
 
 gem "activesupport", "~> 6.1.0"
-gem "sqlite3", "~> 1.4"
 
 gemspec path: "../"

gemfiles/rails_7_0.gemfile

@@ -3,6 +3,5 @@
 source "https://rubygems.org"
 
 gem "activesupport", "~> 7.0.0"
-gem "sqlite3", "~> 1.4"
 
 gemspec path: "../"

lib/db_query_matchers/make_database_queries.rb

@@ -69,6 +69,8 @@ def pluralize(count, singular, plural = nil)
         counter_options[:matches] << Regexp.new(Regexp.escape(options[:matching]))
       end
     end
+
+    counter_options[:database_role] = options[:database_role]
     @counter = DBQueryMatchers::QueryCounter.new(counter_options)
     ActiveSupport::Notifications.subscribed(@counter.to_proc,
                                             DBQueryMatchers.configuration.db_event,

lib/db_query_matchers/query_counter.rb

@@ -18,6 +18,7 @@ class QueryCounter
 
     def initialize(options = {})
       @matches = options[:matches]
+      @database_role = options[:database_role]
       @count = 0
       @log   = []
     end
@@ -39,6 +40,7 @@ def to_proc
     # @param _message_id [String] unique ID for this notification
     # @param payload    [Hash]   the payload
     def callback(_name, _start,  _finish, _message_id, payload)
+      return if @database_role && (ActiveRecord::Base.current_role != @database_role)
       return if @matches && !any_match?(@matches, payload[:sql])
       return if any_match?(DBQueryMatchers.configuration.ignores, payload[:sql])
       return if DBQueryMatchers.configuration.ignore_cached && payload[:cached]

spec/db_query_matchers/make_database_queries_spec.rb

@@ -417,6 +417,27 @@
       end
     end
 
+    if ActiveRecord::VERSION::MAJOR > 6 ||
+        (ActiveRecord::VERSION::MAJOR == 6 && ActiveRecord::VERSION::MINOR > 0)
+      context 'when the database_role option is used' do
+        context 'and a query is using the specified role' do
+          subject { Cat.create }
+          it 'matches true' do
+            expect { subject }.to make_database_queries(database_role: :writing)
+          end
+        end
+
+        context 'and no queries are made matching the role' do
+          it 'raises an error' do
+            expect do
+              expect { subject }.to make_database_queries(database_role: :reading)
+            end.to raise_error(RSpec::Expectations::ExpectationNotMetError,
+                    /expected queries, but none were made/)
+          end
+        end
+      end
+    end
+
     context 'when a `schemaless` option is true' do
       before do
         DBQueryMatchers.configure do |config|

spec/spec_helper.rb

@@ -4,8 +4,14 @@
 Dir[File.dirname(__FILE__) + '/support/**/*.rb'].each { |f| require f }
 
 RSpec.configure do |config|
-  ActiveRecord::Base.establish_connection adapter:  'sqlite3',
-                                          database: ':memory:'
+  if ActiveRecord::VERSION::MAJOR > 5
+    ActiveRecord::Base.establish_connection adapter:  'sqlite3',
+                                            database: ':memory:',
+                                            role: :writing
+  else
+    ActiveRecord::Base.establish_connection adapter:  'sqlite3',
+                                            database: ':memory:'
+  end
 
   ActiveRecord::Schema.define do
     self.verbose = false