From 6973746ae24c30f424abf8f9587aaa37ad390e92 Mon Sep 17 00:00:00 2001
From: benk-gc <bkyriakou@gocardless.com>
Date: Thu, 23 Nov 2023 15:26:44 +0000
Subject: [PATCH] First pass of multi-db support.

---
 lib/generators/statesman/generator_helpers.rb |  2 +-
 lib/statesman.rb                              |  6 +-
 lib/statesman/adapters/active_record.rb       | 50 ++++++++--------
 .../adapters/active_record_queries.rb         |  2 +-
 lib/statesman/config.rb                       | 11 +---
 lib/tasks/statesman.rake                      |  4 +-
 ...active_record_transition_generator_spec.rb |  7 +++
 .../statesman/migration_generator_spec.rb     |  5 ++
 spec/spec_helper.rb                           | 38 ++++++++++--
 spec/statesman/adapters/active_record_spec.rb | 54 +++++++++++++----
 spec/support/active_record.rb                 | 58 ++++++++++++++++---
 spec/support/exactly_query_databases.rb       | 35 +++++++++++
 12 files changed, 205 insertions(+), 67 deletions(-)
 create mode 100644 spec/support/exactly_query_databases.rb

diff --git a/lib/generators/statesman/generator_helpers.rb b/lib/generators/statesman/generator_helpers.rb
index 147d5d71..b6c1a546 100644
--- a/lib/generators/statesman/generator_helpers.rb
+++ b/lib/generators/statesman/generator_helpers.rb
@@ -52,7 +52,7 @@ def configuration
     end
 
     def database_supports_partial_indexes?
-      Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+      Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(klass.constantize)
     end
 
     def metadata_default_value
diff --git a/lib/statesman.rb b/lib/statesman.rb
index 8890039f..57b7821c 100644
--- a/lib/statesman.rb
+++ b/lib/statesman.rb
@@ -34,10 +34,8 @@ def self.storage_adapter
     @storage_adapter || Adapters::Memory
   end
 
-  def self.mysql_gaplock_protection?
-    return @mysql_gaplock_protection unless @mysql_gaplock_protection.nil?
-
-    @mysql_gaplock_protection = config.mysql_gaplock_protection?
+  def self.mysql_gaplock_protection?(connection)
+    config.mysql_gaplock_protection?(connection)
   end
 
   def self.config
diff --git a/lib/statesman/adapters/active_record.rb b/lib/statesman/adapters/active_record.rb
index ff7b3b9f..2d42d2a0 100644
--- a/lib/statesman/adapters/active_record.rb
+++ b/lib/statesman/adapters/active_record.rb
@@ -7,19 +7,15 @@ module Adapters
     class ActiveRecord
       JSON_COLUMN_TYPES = %w[json jsonb].freeze
 
-      def self.database_supports_partial_indexes?
+      def self.database_supports_partial_indexes?(model)
         # Rails 3 doesn't implement `supports_partial_index?`
-        if ::ActiveRecord::Base.connection.respond_to?(:supports_partial_index?)
-          ::ActiveRecord::Base.connection.supports_partial_index?
+        if model.connection.respond_to?(:supports_partial_index?)
+          model.connection.supports_partial_index?
         else
-          ::ActiveRecord::Base.connection.adapter_name == "PostgreSQL"
+          model.connection.adapter_name == "PostgreSQL"
         end
       end
 
-      def self.adapter_name
-        ::ActiveRecord::Base.connection.adapter_name.downcase
-      end
-
       def initialize(transition_class, parent_model, observer, options = {})
         serialized = serialized?(transition_class)
         column_type = transition_class.columns_hash["metadata"].sql_type
@@ -88,10 +84,10 @@ def create_transition(from, to, metadata)
           default_transition_attributes(to, metadata),
         )
 
-        ::ActiveRecord::Base.transaction(requires_new: true) do
+        transition_class.transaction(requires_new: true) do
           @observer.execute(:before, from, to, transition)
 
-          if mysql_gaplock_protection?
+          if mysql_gaplock_protection?(transition_class.connection)
             # We save the transition first with most_recent falsy, then mark most_recent
             # true after to avoid letting MySQL acquire a next-key lock which can cause
             # deadlocks.
@@ -130,8 +126,8 @@ def default_transition_attributes(to, metadata)
       end
 
       def add_after_commit_callback(from, to, transition)
-        ::ActiveRecord::Base.connection.add_transaction_record(
-          ActiveRecordAfterCommitWrap.new do
+        transition_class.connection.add_transaction_record(
+          ActiveRecordAfterCommitWrap.new(transition_class.connection) do
             @observer.execute(:after_commit, from, to, transition)
           end,
         )
@@ -144,7 +140,7 @@ def transitions_for_parent
       # Sets the given transition most_recent = t while unsetting the most_recent of any
       # previous transitions.
       def update_most_recents(most_recent_id = nil)
-        update = build_arel_manager(::Arel::UpdateManager)
+        update = build_arel_manager(::Arel::UpdateManager, transition_class)
         update.table(transition_table)
         update.where(most_recent_transitions(most_recent_id))
         update.set(build_most_recents_update_all_values(most_recent_id))
@@ -152,9 +148,11 @@ def update_most_recents(most_recent_id = nil)
         # MySQL will validate index constraints across the intermediate result of an
         # update. This means we must order our update to deactivate the previous
         # most_recent before setting the new row to be true.
-        update.order(transition_table[:most_recent].desc) if mysql_gaplock_protection?
+        if mysql_gaplock_protection?(transition_class.connection)
+          update.order(transition_table[:most_recent].desc)
+        end
 
-        ::ActiveRecord::Base.connection.update(update.to_sql)
+        transition_class.connection.update(update.to_sql(transition_class))
       end
 
       def most_recent_transitions(most_recent_id = nil)
@@ -223,7 +221,7 @@ def most_recent_value(most_recent_id)
         if most_recent_id
           Arel::Nodes::Case.new.
             when(transition_table[:id].eq(most_recent_id)).then(db_true).
-            else(not_most_recent_value).to_sql
+            else(not_most_recent_value).to_sql(transition_class)
         else
           Arel::Nodes::SqlLiteral.new(not_most_recent_value)
         end
@@ -233,11 +231,11 @@ def most_recent_value(most_recent_id)
       # change in Arel as we move into Rails >6.0.
       #
       # https://github.com/rails/rails/commit/7508284800f67b4611c767bff9eae7045674b66f
-      def build_arel_manager(manager)
+      def build_arel_manager(manager, engine)
         if manager.instance_method(:initialize).arity.zero?
           manager.new
         else
-          manager.new(::ActiveRecord::Base)
+          manager.new(engine)
         end
       end
 
@@ -258,7 +256,7 @@ def transition_conflict_error?(err)
       end
 
       def unique_indexes
-        ::ActiveRecord::Base.connection.
+        transition_class.connection.
           indexes(transition_class.table_name).
           select do |index|
             next unless index.unique
@@ -329,16 +327,16 @@ def default_timezone
         ::ActiveRecord::Base.default_timezone
       end
 
-      def mysql_gaplock_protection?
-        Statesman.mysql_gaplock_protection?
+      def mysql_gaplock_protection?(connection)
+        Statesman.mysql_gaplock_protection?(connection)
       end
 
       def db_true
-        ::ActiveRecord::Base.connection.quote(type_cast(true))
+        transition_class.connection.quote(type_cast(true))
       end
 
       def db_false
-        ::ActiveRecord::Base.connection.quote(type_cast(false))
+        transition_class.connection.quote(type_cast(false))
       end
 
       def db_null
@@ -348,7 +346,7 @@ def db_null
       # Type casting against a column is deprecated and will be removed in Rails 6.2.
       # See https://github.com/rails/arel/commit/6160bfbda1d1781c3b08a33ec4955f170e95be11
       def type_cast(value)
-        ::ActiveRecord::Base.connection.type_cast(value)
+        transition_class.connection.type_cast(value)
       end
 
       # Check whether the `most_recent` column allows null values. If it doesn't, set old
@@ -368,9 +366,9 @@ def not_most_recent_value(db_cast: true)
     end
 
     class ActiveRecordAfterCommitWrap
-      def initialize(&block)
+      def initialize(connection, &block)
         @callback = block
-        @connection = ::ActiveRecord::Base.connection
+        @connection = connection
       end
 
       def self.trigger_transactional_callbacks?
diff --git a/lib/statesman/adapters/active_record_queries.rb b/lib/statesman/adapters/active_record_queries.rb
index d31d0058..35f6a340 100644
--- a/lib/statesman/adapters/active_record_queries.rb
+++ b/lib/statesman/adapters/active_record_queries.rb
@@ -153,7 +153,7 @@ def most_recent_transition_alias
         end
 
         def db_true
-          ::ActiveRecord::Base.connection.quote(true)
+          model.connection.quote(true)
         end
       end
     end
diff --git a/lib/statesman/config.rb b/lib/statesman/config.rb
index c18bcf1c..2814c1bf 100644
--- a/lib/statesman/config.rb
+++ b/lib/statesman/config.rb
@@ -15,17 +15,10 @@ def storage_adapter(adapter_class)
       @adapter_class = adapter_class
     end
 
-    def mysql_gaplock_protection?
-      return @mysql_gaplock_protection unless @mysql_gaplock_protection.nil?
-
+    def mysql_gaplock_protection?(connection)
       # If our adapter class suggests we're using mysql, enable gaplock protection by
       # default.
-      enable_mysql_gaplock_protection if mysql_adapter?(adapter_class)
-      @mysql_gaplock_protection
-    end
-
-    def enable_mysql_gaplock_protection
-      @mysql_gaplock_protection = true
+      mysql_adapter?(connection)
     end
 
     private
diff --git a/lib/tasks/statesman.rake b/lib/tasks/statesman.rake
index 47fa738b..94636869 100644
--- a/lib/tasks/statesman.rake
+++ b/lib/tasks/statesman.rake
@@ -21,8 +21,8 @@ namespace :statesman do
     batch_size = 500
 
     parent_class.find_in_batches(batch_size: batch_size) do |models|
-      ActiveRecord::Base.transaction(requires_new: true) do
-        if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+      transition_class.transaction(requires_new: true) do
+        if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(transition_class)
           # Set all transitions' most_recent to FALSE
           transition_class.where(parent_fk => models.map(&:id)).
             update_all(most_recent: false, updated_at: updated_at)
diff --git a/spec/generators/statesman/active_record_transition_generator_spec.rb b/spec/generators/statesman/active_record_transition_generator_spec.rb
index 68431391..ed71359e 100644
--- a/spec/generators/statesman/active_record_transition_generator_spec.rb
+++ b/spec/generators/statesman/active_record_transition_generator_spec.rb
@@ -5,6 +5,13 @@
 require "generators/statesman/active_record_transition_generator"
 
 describe Statesman::ActiveRecordTransitionGenerator, type: :generator do
+  before do
+    stub_const("Bacon", Class.new(ActiveRecord::Base))
+    stub_const("BaconTransition", Class.new(ActiveRecord::Base))
+    stub_const("Yummy::Bacon", Class.new(ActiveRecord::Base))
+    stub_const("Yummy::BaconTransition", Class.new(ActiveRecord::Base))
+  end
+
   it_behaves_like "a generator" do
     let(:migration_name) { "db/migrate/create_bacon_transitions.rb" }
   end
diff --git a/spec/generators/statesman/migration_generator_spec.rb b/spec/generators/statesman/migration_generator_spec.rb
index 03966bae..c62d1dc7 100644
--- a/spec/generators/statesman/migration_generator_spec.rb
+++ b/spec/generators/statesman/migration_generator_spec.rb
@@ -5,6 +5,11 @@
 require "generators/statesman/migration_generator"
 
 describe Statesman::MigrationGenerator, type: :generator do
+  before do
+    stub_const("Yummy::Bacon", Class.new(ActiveRecord::Base))
+    stub_const("Yummy::BaconTransition", Class.new(ActiveRecord::Base))
+  end
+
   it_behaves_like "a generator" do
     let(:migration_name) { "db/migrate/add_statesman_to_bacon_transitions.rb" }
   end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index f0775b9e..7723ecf3 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -5,13 +5,14 @@
 require "mysql2"
 require "pg"
 require "active_record"
+require "active_record/database_configurations"
 # We have to include all of Rails to make rspec-rails work
 require "rails"
 require "action_view"
 require "action_dispatch"
 require "action_controller"
 require "rspec/rails"
-require "support/active_record"
+require "support/exactly_query_databases"
 require "rspec/its"
 require "pry"
 
@@ -28,10 +29,31 @@ def connection_failure
   if config.exclusion_filter[:active_record]
     puts "Skipping ActiveRecord tests"
   else
-    # Connect to the database for activerecord tests
-    db_conn_spec = ENV["DATABASE_URL"]
-    db_conn_spec ||= { adapter: "sqlite3", database: ":memory:" }
-    ActiveRecord::Base.establish_connection(db_conn_spec)
+    current_env = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+
+    # We have to parse this to a hash since ActiveRecord::Base.configurations
+    # will only consider a single URL config.
+    url_config = if ENV["DATABASE_URL"]
+                   ActiveRecord::DatabaseConfigurations::ConnectionUrlResolver.
+                     new(ENV["DATABASE_URL"]).to_hash.merge({ sslmode: "disable" })
+                 end
+
+    db_config = {
+      current_env => {
+        primary: url_config || {
+          adapter: "sqlite3",
+          database: "/tmp/statesman.db",
+        },
+        secondary: url_config || {
+          adapter: "sqlite3",
+          database: "/tmp/statesman.db",
+        },
+      },
+    }
+
+    # Connect to the primary database for activerecord tests.
+    ActiveRecord::Base.configurations = db_config
+    ActiveRecord::Base.establish_connection(:primary)
 
     db_adapter = ActiveRecord::Base.connection.adapter_name
     puts "Running with database adapter '#{db_adapter}'"
@@ -40,6 +62,8 @@ def connection_failure
     ActiveRecord::Migration.verbose = false
   end
 
+  # Since our primary and secondary connections point to the same database, we don't
+  # need to worry about applying these actions to both.
   config.before(:each, :active_record) do
     tables = %w[
       my_active_record_models
@@ -53,6 +77,7 @@ def connection_failure
     ]
     tables.each do |table_name|
       sql = "DROP TABLE IF EXISTS #{table_name};"
+
       ActiveRecord::Base.connection.execute(sql)
     end
 
@@ -84,3 +109,6 @@ def prepare_sti_transitions_table
     end
   end
 end
+
+# We have to require this after the databases are configured.
+require "support/active_record"
diff --git a/spec/statesman/adapters/active_record_spec.rb b/spec/statesman/adapters/active_record_spec.rb
index a0950b4b..9e9d3697 100644
--- a/spec/statesman/adapters/active_record_spec.rb
+++ b/spec/statesman/adapters/active_record_spec.rb
@@ -10,8 +10,6 @@
     prepare_model_table
     prepare_transitions_table
 
-    # MyActiveRecordModelTransition.serialize(:metadata, JSON)
-
     prepare_sti_model_table
     prepare_sti_transitions_table
 
@@ -26,8 +24,10 @@
 
   after { Statesman.configure { storage_adapter(Statesman::Adapters::Memory) } }
 
+  let(:model_class) { MyActiveRecordModel }
+  let(:transition_class) { MyActiveRecordModelTransition }
   let(:observer) { double(Statesman::Machine, execute: nil) }
-  let(:model) { MyActiveRecordModel.create(current_state: :pending) }
+  let(:model) { model_class.create(current_state: :pending) }
 
   it_behaves_like "an adapter", described_class, MyActiveRecordModelTransition
 
@@ -36,8 +36,8 @@
       before do
         metadata_column = double
         allow(metadata_column).to receive_messages(sql_type: "")
-        allow(MyActiveRecordModelTransition).to receive_messages(columns_hash:
-                                           { "metadata" => metadata_column })
+        allow(MyActiveRecordModelTransition).
+          to receive_messages(columns_hash: { "metadata" => metadata_column })
         expect(MyActiveRecordModelTransition).
           to receive(:type_for_attribute).with("metadata").
           and_return(ActiveRecord::Type::Value.new)
@@ -105,9 +105,7 @@
   describe "#create" do
     subject(:transition) { create }
 
-    let!(:adapter) do
-      described_class.new(MyActiveRecordModelTransition, model, observer)
-    end
+    let!(:adapter) { described_class.new(transition_class, model, observer) }
     let(:from) { :x }
     let(:to) { :y }
     let(:create) { adapter.create(from, to) }
@@ -343,12 +341,35 @@
         end
       end
     end
+
+    context "when using the secondary database" do
+      let(:model_class) { SecondaryActiveRecordModel }
+      let(:transition_class) { SecondaryActiveRecordModelTransition }
+
+      it "retrieves the new transition from the database" do
+        expect { transition }.to exactly_query_databases({ secondary: [:writing] })
+
+        expect(adapter.last.to_state).to eq("y")
+      end
+
+      context "when there is a race" do
+        it "raises a TransitionConflictError and uses the correct database" do
+          adapter2 = adapter.dup
+          adapter2.create(:x, :y)
+          adapter.last
+          adapter2.create(:y, :z)
+
+          expect { adapter.create(:y, :z) }.
+            to raise_exception(Statesman::TransitionConflictError).
+            and exactly_query_databases({ secondary: [:writing] })
+        end
+      end
+    end
   end
 
   describe "#last" do
-    let(:adapter) do
-      described_class.new(MyActiveRecordModelTransition, model, observer)
-    end
+    let(:transition_class) { MyActiveRecordModelTransition }
+    let(:adapter) { described_class.new(transition_class, model, observer) }
 
     context "with a previously looked up transition" do
       before { adapter.create(:x, :y) }
@@ -367,6 +388,17 @@
         it "retrieves the new transition from the database" do
           expect(adapter.last.to_state).to eq("z")
         end
+
+        context "when using the secondary database" do
+          let(:model_class) { SecondaryActiveRecordModel }
+          let(:transition_class) { SecondaryActiveRecordModelTransition }
+
+          it "retrieves the new transition from the database" do
+            expect { adapter.last.to_state }.to exactly_query_databases({ secondary: [:writing] })
+
+            expect(adapter.last.to_state).to eq("z")
+          end
+        end
       end
 
       context "when a new transition has been created elsewhere" do
diff --git a/spec/support/active_record.rb b/spec/support/active_record.rb
index 82d739ed..a8c7dbc0 100644
--- a/spec/support/active_record.rb
+++ b/spec/support/active_record.rb
@@ -81,7 +81,7 @@ def change
         t.text :metadata, default: "{}"
       end
 
-      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
         t.boolean :most_recent, default: true, null: false
       else
         t.boolean :most_recent, default: true
@@ -98,7 +98,7 @@ def change
               %i[my_active_record_model_id sort_key],
               unique: true, name: "sort_key_index"
 
-    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
       add_index :my_active_record_model_transitions,
                 %i[my_active_record_model_id most_recent],
                 unique: true,
@@ -134,6 +134,48 @@ class OtherActiveRecordModelTransition < ActiveRecord::Base
   belongs_to :other_active_record_model
 end
 
+class SecondaryRecord < ActiveRecord::Base
+  self.abstract_class = true
+
+  connects_to database: { writing: :secondary, reading: :secondary }
+end
+
+class SecondaryActiveRecordModelTransition < SecondaryRecord
+  self.table_name = "my_active_record_model_transitions"
+
+  include Statesman::Adapters::ActiveRecordTransition
+
+  belongs_to :my_active_record_model,
+             class_name: "SecondaryActiveRecordModel",
+             foreign_key: "my_active_record_model_transition_id"
+end
+
+class SecondaryActiveRecordModel < SecondaryRecord
+  self.table_name = "my_active_record_models"
+
+  has_many :my_active_record_model_transitions,
+           class_name: "SecondaryActiveRecordModelTransition",
+           foreign_key: "my_active_record_model_id",
+           autosave: false
+
+  alias_method :transitions, :my_active_record_model_transitions
+
+  include Statesman::Adapters::ActiveRecordQueries[
+            transition_class: SecondaryActiveRecordModelTransition,
+            initial_state: :initial
+          ]
+
+  def state_machine
+    @state_machine ||= MyStateMachine.new(
+      self, transition_class: SecondaryActiveRecordModelTransition
+    )
+  end
+
+  def metadata
+    super || {}
+  end
+end
+
 class CreateOtherActiveRecordModelMigration < MIGRATION_CLASS
   def change
     create_table :other_active_record_models do |t|
@@ -158,7 +200,7 @@ def change
         t.text :metadata, default: "{}"
       end
 
-      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
         t.boolean :most_recent, default: true, null: false
       else
         t.boolean :most_recent, default: true
@@ -171,7 +213,7 @@ def change
               %i[other_active_record_model_id sort_key],
               unique: true, name: "other_sort_key_index"
 
-    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
       add_index :other_active_record_model_transitions,
                 %i[other_active_record_model_id most_recent],
                 unique: true,
@@ -253,7 +295,7 @@ def change
         t.text :metadata, default: "{}"
       end
 
-      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
         t.boolean :most_recent, default: true, null: false
       else
         t.boolean :most_recent, default: true
@@ -265,7 +307,7 @@ def change
     add_index :my_namespace_my_active_record_model_transitions, :sort_key,
               unique: true, name: "my_namespaced_key"
 
-    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
       add_index :my_namespace_my_active_record_model_transitions,
                 %i[my_active_record_model_id most_recent],
                 unique: true,
@@ -342,7 +384,7 @@ def change
         t.text :metadata, default: "{}"
       end
 
-      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+      if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
         t.boolean :most_recent, default: true, null: false
       else
         t.boolean :most_recent, default: true
@@ -355,7 +397,7 @@ def change
               %i[type sti_active_record_model_id sort_key],
               unique: true, name: "sti_sort_key_index"
 
-    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?
+    if Statesman::Adapters::ActiveRecord.database_supports_partial_indexes?(ActiveRecord::Base)
       add_index :sti_active_record_model_transitions,
                 %i[type sti_active_record_model_id most_recent],
                 unique: true,
diff --git a/spec/support/exactly_query_databases.rb b/spec/support/exactly_query_databases.rb
new file mode 100644
index 00000000..209b6231
--- /dev/null
+++ b/spec/support/exactly_query_databases.rb
@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+# `expected_dbs` should be a Hash of the form:
+# {
+#   primary: [:writing, :reading],
+#   replica: [:reading],
+# }
+RSpec::Matchers.define :exactly_query_databases do |expected_dbs|
+  match do |block|
+    @expected_dbs = expected_dbs.transform_values(&:to_set).with_indifferent_access
+    @actual_dbs = Hash.new { |h, k| h[k] = Set.new }.with_indifferent_access
+
+    ActiveSupport::Notifications.
+      subscribe("sql.active_record") do |_name, _start, _finish, _id, payload|
+      pool = payload.fetch(:connection).pool
+
+      next if pool.is_a?(ActiveRecord::ConnectionAdapters::NullPool)
+
+      name = pool.db_config.name
+      role = pool.role
+
+      @actual_dbs[name] << role
+    end
+
+    block.call
+
+    @actual_dbs == @expected_dbs
+  end
+
+  failure_message do |_block|
+    "expected to query exactly #{@expected_dbs}, but queried #{@actual_dbs}"
+  end
+
+  supports_block_expectations
+end